From 72d67c80f2d48536d866b5845a93acc1aa33c79c Mon Sep 17 00:00:00 2001 From: ArthurKoba <61914203+ArthurKoba@users.noreply.github.com> Date: Sat, 7 Nov 2026 16:29:47 +0300 Subject: [PATCH] fix(agent): make discovery enrollment compatible and idempotent (#53) * fix(agent): make discovery enrollment compatible and idempotent * fix(agent): confirm enrollment persistence and document token aliases --- docs/changelog/index.md | 2 + docs/design/agent-design.md | 11 +- docs/docs.go | 82 ++++++- docs/swagger.json | 82 ++++++- docs/swagger.yaml | 57 ++++- .../agent/agent/discovery_integration_test.go | 223 ++++++++++++++++++ .../agent/agent/discovery_persistence_test.go | 85 +++++++ internal/apps/agent/agent/discovery_retry.go | 30 +++ .../apps/agent/agent/discovery_retry_test.go | 30 +++ internal/apps/agent/agent/runner.go | 45 ++-- internal/apps/agent/agent/runner_test.go | 6 +- internal/apps/agent/config/config.go | 3 +- internal/apps/agent/heartbeat/service.go | 6 +- internal/apps/agent/httpclient/client.go | 2 +- internal/apps/agent/localfile/atomic.go | 66 ++++++ internal/apps/agent/localfile/atomic_test.go | 88 +++++++ .../agent/localfile/sync_directory_unix.go | 12 + .../agent/localfile/sync_directory_windows.go | 12 + internal/apps/agent/protocol/alias.go | 3 + .../apps/agent/state/registration_test.go | 84 +++++++ internal/apps/agent/state/state.go | 68 +++++- internal/apps/openflare/agent/errs.go | 28 ++- internal/apps/openflare/agent/logics.go | 31 ++- internal/apps/openflare/agent/registration.go | 25 ++ .../apps/openflare/agent/registration_test.go | 168 +++++++++++++ internal/apps/openflare/agent/routers.go | 11 +- internal/apps/openflare/agent/types.go | 9 +- .../repository/openflare_node_registration.go | 28 +++ pkg/protocol/agent.go | 44 +++- pkg/protocol/agent_registration_test.go | 54 +++++ 30 files changed, 1324 insertions(+), 71 deletions(-) create mode 100644 internal/apps/agent/agent/discovery_integration_test.go create mode 100644 internal/apps/agent/agent/discovery_persistence_test.go create mode 100644 internal/apps/agent/agent/discovery_retry.go create mode 100644 internal/apps/agent/agent/discovery_retry_test.go create mode 100644 internal/apps/agent/localfile/atomic.go create mode 100644 internal/apps/agent/localfile/atomic_test.go create mode 100644 internal/apps/agent/localfile/sync_directory_unix.go create mode 100644 internal/apps/agent/localfile/sync_directory_windows.go create mode 100644 internal/apps/agent/state/registration_test.go create mode 100644 internal/apps/openflare/agent/registration.go create mode 100644 internal/apps/openflare/agent/registration_test.go create mode 100644 internal/repository/openflare_node_registration.go create mode 100644 pkg/protocol/agent_registration_test.go diff --git a/docs/changelog/index.md b/docs/changelog/index.md index 95a9ac29..9c6aa068 100644 --- a/docs/changelog/index.md +++ b/docs/changelog/index.md @@ -19,6 +19,8 @@ sidebar: false ### 🛠 修复 +- 修复 Agent 自动注册响应凭证字段不一致导致节点不断重复创建的问题；新增持久化注册重试凭证与退避，响应丢失、并发重试或配置写入失败后仍复用原节点，且不会覆盖管理员修改或恢复已更换的节点凭证。支持目录刷盘的平台会确认原子替换后的目录持久化，失败时阻止使用未确认的注册凭证；接口文档同步列出两种兼容凭证字段。 + - SQLite 文件数据库启动时自动启用 WAL 模式，改善读写并发；数据库导出改用一致性快照，确保包含 WAL 中已提交的数据。 - 修复 IP 组按规则自动抓取时会收录内网、回环及链路本地等地址的问题；规则测试同步过滤这些地址，历史抓取结果会在下次同步时清理。 diff --git a/docs/design/agent-design.md b/docs/design/agent-design.md index ce111616..310bcfe9 100644 --- a/docs/design/agent-design.md +++ b/docs/design/agent-design.md @@ -37,10 +37,13 @@ Agent 主要由以下核心子模块组成，共同配合完成其完整的生 Agent 在生命周期中主要通过 **基于 Token 的自动注册** 和 **心跳/WebSocket 双通道** 与控制面通信。 ### 1. 自动注册流程 -若 Agent 启动时本地 `agent.json` 的 `access_token` 为空，但配置了 `discovery_token`，将触发自动注册流程： -1. Agent 向控制面 `/api/v1/agent/nodes/register` 发送注册请求，携带本地硬件摘要、IP 及主机名。 -2. Server 校验 `discovery_token` 有效后，在数据库生成唯一的 `NodeID` 与专属 `AccessToken`（即 `agent_token`）并返回。 -3. Agent 将获取的专用 Token 写入本地配置文件，擦除一次性 `discovery_token`，后续所有的通信均基于专属 `AccessToken` 进行鉴权认证。 +若 Agent 启动时本地 `agent.json` 的 `agent_token` 为空，但配置了 `discovery_token`，将触发自动注册流程： +1. Agent 在本地状态文件原子保存随机 256 位 `registration_key`，再向控制面 `/api/v1/agent/nodes/register` 发送注册请求，携带该重试凭证、本地硬件摘要、IP 及主机名。支持目录刷盘的平台会在文件刷盘及重命名后刷新所在目录；目录刷盘失败会返回错误，已有凭证也必须在重试时确认目录刷盘成功后才能发送。Windows 保留文件刷盘与替换，但目录持久化仅为尽力而为；上述行为针对所在目录已存在的文件更新，不保证所有文件系统或突然断电场景下的完整持久性。该凭证仅用于注册，不进入心跳、WebSocket 或观测数据，必须与节点 Token 一样保密。 +2. Server 校验当前 `discovery_token` 有效后，从重试凭证使用不同的哈希域分别生成 `NodeID` 与专属 Token。相同凭证的重复或并发请求只保留一个节点，不按 IP、主机名或公开 `node_id` 合并，也不会覆盖管理员修改的名称、升级标志等数据。响应以 `agent_token` 为准，同时保留旧的 `access_token` 字段；新版 Agent 能识别两种字段，值冲突时拒绝响应。 +3. Agent 原子保存服务端分配的节点 ID 和配置 Token，全部写入成功后才切换运行时鉴权并清空配置中的 `discovery_token`，后续通信均使用节点专属 Token。响应丢失或配置写入失败时，重启仍沿用原重试凭证；状态文件中的重试凭证与节点 ID 不会被旧的后台任务快照覆盖。 +4. 注册失败按 3、6、12、24、48、60 秒退避，之后每次至少间隔 60 秒重试（实际请求在下一个心跳定时器周期触发），不再随每次心跳立即注册。成功后重置退避。 + +若节点 Token 已被清空或更换，旧重试凭证不能取回新 Token 或恢复旧 Token，必须使用管理员分配的节点凭证。物理删除节点不构成永久封禁：持有当前有效 discovery Token 的客户端仍能重新注册。旧版 Agent 不携带 `registration_key`，修复后的响应可使其正常完成注册，但其响应丢失后的重试仍不具备幂等保证。请保留每个实例独立的持久化状态与配置，不要把已初始化的状态目录复制给其他实例。 ### 2. 双通道心跳与同步机制 * **HTTP 轮询通道（兜底与探测）**：Agent 默认按设定的 `heartbeat_interval` 间隔发送 POST 心跳包。上报指标的同时获取当前激活版本的摘要信息（Version & Checksum）。 diff --git a/docs/docs.go b/docs/docs.go index 693d38fa..9a68de80 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -4443,12 +4443,12 @@ const docTemplate = `{ "summary": "注册或发现 Agent 节点", "parameters": [ { - "description": "节点上报数据", + "description": "节点上报数据；registration_key 为首次注册重试凭证", "name": "body", "in": "body", "required": true, "schema": { - "$ref": "#/definitions/agent.NodePayload" + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest" } } ], @@ -4464,7 +4464,7 @@ const docTemplate = `{ "type": "object", "properties": { "data": { - "$ref": "#/definitions/agent.RegistrationResponse" + "$ref": "#/definitions/agent.RegistrationResponseDocs" } } } @@ -14989,12 +14989,15 @@ const docTemplate = `{ } } }, - "agent.RegistrationResponse": { + "agent.RegistrationResponseDocs": { "type": "object", "properties": { "access_token": { "type": "string" }, + "agent_token": { + "type": "string" + }, "name": { "type": "string" }, @@ -16293,6 +16296,77 @@ const docTemplate = `{ } } }, + "github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest": { + "type": "object", + "properties": { + "access_logs": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeAccessLog" + } + }, + "buffered": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.BufferedObservabilityRecord" + } + }, + "current_version": { + "type": "string" + }, + "edge_health": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeEdgeHealth" + }, + "ext_version": { + "type": "string" + }, + "health_events": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeHealthEvent" + } + }, + "host_metrics": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeMetricSnapshot" + }, + "ip": { + "type": "string" + }, + "last_error": { + "type": "string" + }, + "name": { + "type": "string" + }, + "node_id": { + "type": "string" + }, + "openresty_message": { + "type": "string" + }, + "openresty_status": { + "type": "string" + }, + "profile": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeSystemProfile" + }, + "registration_key": { + "type": "string" + }, + "schema_version": { + "type": "integer" + }, + "version": { + "type": "string" + }, + "waf_ip_group_checksums": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + }, "github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup": { "type": "object", "properties": { diff --git a/docs/swagger.json b/docs/swagger.json index 36c583c0..7c3f0c5a 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -4436,12 +4436,12 @@ "summary": "注册或发现 Agent 节点", "parameters": [ { - "description": "节点上报数据", + "description": "节点上报数据；registration_key 为首次注册重试凭证", "name": "body", "in": "body", "required": true, "schema": { - "$ref": "#/definitions/agent.NodePayload" + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest" } } ], @@ -4457,7 +4457,7 @@ "type": "object", "properties": { "data": { - "$ref": "#/definitions/agent.RegistrationResponse" + "$ref": "#/definitions/agent.RegistrationResponseDocs" } } } @@ -14982,12 +14982,15 @@ } } }, - "agent.RegistrationResponse": { + "agent.RegistrationResponseDocs": { "type": "object", "properties": { "access_token": { "type": "string" }, + "agent_token": { + "type": "string" + }, "name": { "type": "string" }, @@ -16286,6 +16289,77 @@ } } }, + "github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest": { + "type": "object", + "properties": { + "access_logs": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeAccessLog" + } + }, + "buffered": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.BufferedObservabilityRecord" + } + }, + "current_version": { + "type": "string" + }, + "edge_health": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeEdgeHealth" + }, + "ext_version": { + "type": "string" + }, + "health_events": { + "type": "array", + "items": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeHealthEvent" + } + }, + "host_metrics": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeMetricSnapshot" + }, + "ip": { + "type": "string" + }, + "last_error": { + "type": "string" + }, + "name": { + "type": "string" + }, + "node_id": { + "type": "string" + }, + "openresty_message": { + "type": "string" + }, + "openresty_status": { + "type": "string" + }, + "profile": { + "$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeSystemProfile" + }, + "registration_key": { + "type": "string" + }, + "schema_version": { + "type": "integer" + }, + "version": { + "type": "string" + }, + "waf_ip_group_checksums": { + "type": "object", + "additionalProperties": { + "type": "string" + } + } + } + }, "github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup": { "type": "object", "properties": { diff --git a/docs/swagger.yaml b/docs/swagger.yaml index f32d2aeb..20b3adc8 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -99,10 +99,12 @@ definitions: type: string type: object type: object - agent.RegistrationResponse: + agent.RegistrationResponseDocs: properties: access_token: type: string + agent_token: + type: string name: type: string node_id: @@ -955,6 +957,53 @@ definitions: total_size: type: integer type: object + github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest: + properties: + access_logs: + items: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeAccessLog' + type: array + buffered: + items: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.BufferedObservabilityRecord' + type: array + current_version: + type: string + edge_health: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeEdgeHealth' + ext_version: + type: string + health_events: + items: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeHealthEvent' + type: array + host_metrics: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeMetricSnapshot' + ip: + type: string + last_error: + type: string + name: + type: string + node_id: + type: string + openresty_message: + type: string + openresty_status: + type: string + profile: + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.NodeSystemProfile' + registration_key: + type: string + schema_version: + type: integer + version: + type: string + waf_ip_group_checksums: + additionalProperties: + type: string + type: object + type: object github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup: properties: checksum: @@ -7020,12 +7069,12 @@ paths: - application/json description: 使用节点 access token 重新注册，或使用全局 discovery token 发现新节点；请求头需携带 X-Agent-Token parameters: - - description: 节点上报数据 + - description: 节点上报数据；registration_key 为首次注册重试凭证 in: body name: body required: true schema: - $ref: '#/definitions/agent.NodePayload' + $ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest' produces: - application/json responses: @@ -7036,7 +7085,7 @@ paths: - $ref: '#/definitions/response.Any' - properties: data: - $ref: '#/definitions/agent.RegistrationResponse' + $ref: '#/definitions/agent.RegistrationResponseDocs' type: object "400": description: 参数错误 diff --git a/internal/apps/agent/agent/discovery_integration_test.go b/internal/apps/agent/agent/discovery_integration_test.go new file mode 100644 index 00000000..612f90c3 --- /dev/null +++ b/internal/apps/agent/agent/discovery_integration_test.go @@ -0,0 +1,223 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Rain-kl/Wavelet/internal/apps/agent/config" + agentheartbeat "github.com/Rain-kl/Wavelet/internal/apps/agent/heartbeat" + "github.com/Rain-kl/Wavelet/internal/apps/agent/httpclient" + "github.com/Rain-kl/Wavelet/internal/apps/agent/protocol" + "github.com/Rain-kl/Wavelet/internal/apps/agent/state" + serveragent "github.com/Rain-kl/Wavelet/internal/apps/openflare/agent" + db "github.com/Rain-kl/Wavelet/internal/infra/persistence" + "github.com/Rain-kl/Wavelet/internal/model" + "github.com/Rain-kl/Wavelet/internal/repository" + "github.com/Rain-kl/Wavelet/internal/shared/response" + "github.com/Rain-kl/Wavelet/internal/testhelper" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/pressly/goose/v3" + "gorm.io/gorm" +) + +func newDiscoveryTestRunner(t *testing.T, configPath string) (*Runner, string) { + t.Helper() + cfg, err := config.Load(configPath) + if err != nil { + t.Fatal(err) + } + store := state.NewStore(cfg.StatePath) + nodeID, err := store.EnsureNodeID() + if err != nil { + t.Fatal(err) + } + client := httpclient.New(cfg.ServerURL, cfg.InitialAuthToken(), time.Second) + runner := withHeartbeatCycle(&Runner{Config: cfg, StateStore: store, + HeartbeatService: agentheartbeat.New(client), SyncService: &fakeSyncService{}}, nil) + return runner, nodeID +} + +func setupDiscoveryIntegrationDB(t *testing.T) *gorm.DB { + t.Helper() + ctx := context.Background() + conn, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "discovery.db")), &gorm.Config{DisableForeignKeyConstraintWhenMigrating: true}) + if err != nil { + t.Fatal(err) + } + sqlDB, err := conn.DB() + if err != nil { + t.Fatal(err) + } + previous := db.DB(ctx) + t.Cleanup(func() { + if invalidateErr := repository.InvalidateAllSystemConfigCaches(ctx); invalidateErr != nil { + t.Error(invalidateErr) + } + db.SetDB(previous) + if closeErr := sqlDB.Close(); closeErr != nil { + t.Error(closeErr) + } + }) + provider, err := goose.NewProvider(goose.DialectSQLite3, sqlDB, os.DirFS("../../../infra/persistence/migrator/goose/sqlite")) + if err != nil { + t.Fatal(err) + } + if _, err = provider.Up(ctx); err != nil { + t.Fatal(err) + } + db.SetDB(conn) + return conn +} + +func TestDiscoveryHTTPRecoveryAndRestart(t *testing.T) { + ctx := context.Background() + conn := setupDiscoveryIntegrationDB(t) + var err error + if err = repository.SaveOrUpdateSystemConfig(ctx, model.ConfigKeyAgentDiscoveryToken, "discovery-test"); err != nil { + t.Fatal(err) + } + if err = repository.SaveOrUpdateSystemConfig(ctx, model.ConfigKeyGeoIPProvider, "disabled"); err != nil { + t.Fatal(err) + } + var registrations, heartbeats atomic.Int32 + var dropResponse atomic.Bool + dropResponse.Store(true) + router := testhelper.NewTestGinEngine() + router.POST("/api/v1/agent/nodes/register", serveragent.RegisterAuth(), func(c *gin.Context) { + registrations.Add(1) + serveragent.RegisterHandler(c) + }) + router.POST("/api/v1/agent/nodes/heartbeat", serveragent.Auth(), func(c *gin.Context) { + var payload map[string]json.RawMessage + if decodeErr := c.ShouldBindJSON(&payload); decodeErr != nil { + t.Error(decodeErr) + } + if _, exists := payload["registration_key"]; exists { + t.Error("heartbeat exposed registration_key") + } + heartbeats.Add(1) + c.JSON(http.StatusOK, response.OKNil()) + }) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/agent/nodes/register" && dropResponse.Swap(false) { + recorder := httptest.NewRecorder() + router.ServeHTTP(recorder, r) + if recorder.Code != http.StatusOK { + t.Errorf("First registration status = %d, want 200", recorder.Code) + } + w.WriteHeader(http.StatusBadGateway) // Enrollment committed, but its response was lost. + return + } + router.ServeHTTP(w, r) + })) + t.Cleanup(server.Close) + configPath := filepath.Join(t.TempDir(), "agent.json") + data, err := json.Marshal(map[string]string{"server_url": server.URL, "discovery_token": "discovery-test", "node_ip": "10.0.0.8", "node_name": "edge"}) + if err != nil { + t.Fatal(err) + } + if err = os.WriteFile(configPath, data, 0o600); err != nil { + t.Fatal(err) + } + + first, nodeID := newDiscoveryTestRunner(t, configPath) + if err = first.tryRegister(ctx, &nodeID); err == nil { + t.Fatal("tryRegister(lost response) succeeded, want transport failure") + } + key, err := first.StateStore.EnsureRegistrationKey() + if err != nil { + t.Fatal(err) + } + if first.hasAccessToken() { + t.Fatal("tryRegister(lost response) activated token") + } + + // Restart with the same private key, then fail the second durable identity write. + second, secondID := newDiscoveryTestRunner(t, configPath) + backup := configPath + ".original" + if err = os.Rename(configPath, backup); err != nil { + t.Fatal(err) + } + if err = os.Mkdir(configPath, 0o700); err != nil { + t.Fatal(err) + } + if err = second.tryRegister(ctx, &secondID); err == nil { + t.Fatal("tryRegister(config save failure) succeeded, want write error") + } + if second.hasAccessToken() || second.Config.DiscoveryToken == "" { + t.Fatal("Config write failure activated partial in-memory credentials") + } + if err = os.Remove(configPath); err != nil { + t.Fatal(err) + } + if err = os.Rename(backup, configPath); err != nil { + t.Fatal(err) + } + + third, thirdID := newDiscoveryTestRunner(t, configPath) + if err = third.tryRegister(ctx, &thirdID); err != nil { + t.Fatalf("tryRegister(after restart) error = %v", err) + } + if !third.hasAccessToken() || third.Config.DiscoveryToken != "" { + t.Fatal("Successful enrollment did not activate node credentials") + } + persistedKey, err := third.StateStore.EnsureRegistrationKey() + if err != nil || persistedKey != key { + t.Errorf("Retry key changed across partial-write restarts: %v", err) + } + + restarted, restartedID := newDiscoveryTestRunner(t, configPath) + restarted.runStartupAuth(ctx, &restartedID) + if restartedID != thirdID { + t.Errorf("Restarted node ID = %q, want %q", restartedID, thirdID) + } + if registrations.Load() != 3 || heartbeats.Load() != 2 { + t.Errorf("Enrollment/restart calls = %d registrations, %d heartbeats; want 3, 2", registrations.Load(), heartbeats.Load()) + } + var count int64 + if err = conn.Model(&model.OpenFlareNode{}).Count(&count).Error; err != nil { + t.Fatal(err) + } + if count != 1 { + t.Errorf("Enrollment through lost response and partial write created %d nodes, want 1", count) + } + + // Knowing the private retry key alone is insufficient after the global discovery token rotates. + if err = repository.SaveOrUpdateSystemConfig(ctx, model.ConfigKeyAgentDiscoveryToken, "rotated-discovery"); err != nil { + t.Fatal(err) + } + request := protocol.RegisterNodeRequest{NodePayload: protocol.NodePayload{IP: "10.0.0.8", Version: "test"}, RegistrationKey: key} + for _, bearer := range []string{"", "invalid", "discovery-test"} { + client := httpclient.New(server.URL, bearer, time.Second) + result, authErr := client.RegisterNode(ctx, request) + if authErr == nil || result != nil { + t.Errorf("RegisterNode(unauthorized bearer) = %+v, %v; want rejection", result, authErr) + } + } + encoded, err := json.Marshal(request) + if err != nil { + t.Fatal(err) + } + unauthorized := httptest.NewRecorder() + router.ServeHTTP(unauthorized, httptest.NewRequest(http.MethodPost, "/api/v1/agent/nodes/register", strings.NewReader(string(encoded)))) + if unauthorized.Code != http.StatusUnauthorized || strings.Contains(unauthorized.Body.String(), third.Config.AccessToken) { + t.Errorf("Unauthorized retry response status = %d; want 401 without credential", unauthorized.Code) + } + client := httpclient.New(server.URL, "rotated-discovery", time.Second) + result, err := client.RegisterNode(ctx, request) + if err != nil || result.AccessToken != third.Config.AccessToken { + t.Errorf("RegisterNode(current discovery bearer) error = %v, want same identity", err) + } +} diff --git a/internal/apps/agent/agent/discovery_persistence_test.go b/internal/apps/agent/agent/discovery_persistence_test.go new file mode 100644 index 00000000..9d3d2354 --- /dev/null +++ b/internal/apps/agent/agent/discovery_persistence_test.go @@ -0,0 +1,85 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/Rain-kl/Wavelet/internal/apps/agent/config" + "github.com/Rain-kl/Wavelet/internal/apps/agent/protocol" + "github.com/Rain-kl/Wavelet/internal/apps/agent/state" +) + +func TestDiscoveryDoesNotRegisterBeforeRetrySecretPersists(t *testing.T) { + blockedParent := filepath.Join(t.TempDir(), "file") + if err := os.WriteFile(blockedParent, []byte("not a directory"), 0o600); err != nil { + t.Fatal(err) + } + service := &fakeHeartbeatService{} + runner := &Runner{Config: &config.Config{DiscoveryToken: "discovery"}, + StateStore: state.NewStore(filepath.Join(blockedParent, "state.json")), HeartbeatService: service} + nodeID := "local" + if err := runner.tryRegister(context.Background(), &nodeID); err == nil { + t.Fatal("tryRegister(unwritable state) succeeded, want persistence error") + } + if service.registerCalls != 0 { + t.Errorf("tryRegister(unwritable retry key) sent %d requests, want 0", service.registerCalls) + } +} + +func TestDiscoveryStateWriteFailureDoesNotActivateCredentials(t *testing.T) { + configPath := filepath.Join(t.TempDir(), "agent.json") + if err := os.WriteFile(configPath, []byte(`{"server_url":"http://127.0.0.1:3000","discovery_token":"discovery","node_ip":"10.0.0.8"}`), 0o600); err != nil { + t.Fatal(err) + } + runner, nodeID := newDiscoveryTestRunner(t, configPath) + statePath := runner.Config.StatePath + backup := statePath + ".original" + service := &fakeHeartbeatService{registerResp: &protocol.RegisterNodeResponse{NodeID: "server-assigned", AccessToken: "issued"}} + var originalKey string + service.onRegister = func() { + snapshot, err := runner.StateStore.Load() + if err != nil { + t.Fatal(err) + } + originalKey = snapshot.RegistrationKey + if err = os.Rename(statePath, backup); err != nil { + t.Fatal(err) + } + if err = os.Mkdir(statePath, 0o700); err != nil { + t.Fatal(err) + } + } + runner.HeartbeatService = service + runner.HeartbeatCycle.Heartbeat = service + if err := runner.tryRegister(context.Background(), &nodeID); err == nil { + t.Fatal("tryRegister(state write failure) succeeded, want error") + } + if runner.hasAccessToken() || service.lastToken != "" { + t.Fatal("Failed state persistence activated credentials") + } + if err := os.Remove(statePath); err != nil { + t.Fatal(err) + } + if err := os.Rename(backup, statePath); err != nil { + t.Fatal(err) + } + restarted, restartedID := newDiscoveryTestRunner(t, configPath) + service.onRegister = nil + restarted.HeartbeatService = service + restarted.HeartbeatCycle.Heartbeat = service + if err := restarted.tryRegister(context.Background(), &restartedID); err != nil { + t.Fatal(err) + } + key, err := restarted.StateStore.EnsureRegistrationKey() + if err != nil || key != originalKey { + t.Errorf("Key after state write failure/restart changed: %v", err) + } + if restartedID != "server-assigned" || service.lastToken != "issued" { + t.Errorf("Recovered identity/token = %q/%q, want server-assigned/issued", restartedID, service.lastToken) + } +} diff --git a/internal/apps/agent/agent/discovery_retry.go b/internal/apps/agent/agent/discovery_retry.go new file mode 100644 index 00000000..b1200486 --- /dev/null +++ b/internal/apps/agent/agent/discovery_retry.go @@ -0,0 +1,30 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import "time" + +const ( + discoveryRetryInitialDelay = 3 * time.Second + discoveryRetryMaxDelay = time.Minute + discoveryRetryMultiplier = 2 +) + +type discoveryRetry struct { + delay time.Duration + next time.Time +} + +func (r *discoveryRetry) ready(now time.Time) bool { + return !now.Before(r.next) +} + +func (r *discoveryRetry) failed(now time.Time) { + if r.delay == 0 { + r.delay = discoveryRetryInitialDelay + } else { + r.delay = min(r.delay*discoveryRetryMultiplier, discoveryRetryMaxDelay) + } + r.next = now.Add(r.delay) +} diff --git a/internal/apps/agent/agent/discovery_retry_test.go b/internal/apps/agent/agent/discovery_retry_test.go new file mode 100644 index 00000000..c134064e --- /dev/null +++ b/internal/apps/agent/agent/discovery_retry_test.go @@ -0,0 +1,30 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "testing" + "time" +) + +func TestDiscoveryRetryBackoffIsIndependentOfHeartbeat(t *testing.T) { + now := time.Unix(100, 0) + var retry discoveryRetry + for _, delay := range []time.Duration{3, 6, 12, 24, 48, 60, 60} { + if !retry.ready(now) { + t.Fatal("Discovery retry not ready at scheduled time") + } + retry.failed(now) + wantNext := now.Add(delay * time.Second) + if retry.next != wantNext { + t.Errorf("Discovery retry next = %v, want %v", retry.next, wantNext) + } + for simulatedHeartbeat := now; simulatedHeartbeat.Before(wantNext); simulatedHeartbeat = simulatedHeartbeat.Add(time.Millisecond) { + if retry.ready(simulatedHeartbeat) { + t.Fatal("Discovery retry ready before independent enrollment backoff expires") + } + } + now = wantNext + } +} diff --git a/internal/apps/agent/agent/runner.go b/internal/apps/agent/agent/runner.go index 7c5991eb..66cd2ea1 100644 --- a/internal/apps/agent/agent/runner.go +++ b/internal/apps/agent/agent/runner.go @@ -23,7 +23,7 @@ import ( // HeartbeatService handles node registration and periodic heartbeat reporting. type HeartbeatService interface { - Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) + Register(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) SetToken(token string) } @@ -65,6 +65,7 @@ type Runner struct { configSyncRetry *configSyncRetry restartOpenrestyNow bool websocketUpgradeEnabled bool + discoveryRetry discoveryRetry } // Run starts the agent's main loop, performing heartbeats and upgrading to WebSocket when available. @@ -120,6 +121,9 @@ func (r *Runner) Run(ctx context.Context) error { tryStartWebSocket() for { + if err := ctx.Err(); err != nil { + return err + } select { case <-ctx.Done(): slog.Info("agent runner shutting down", "error", ctx.Err()) @@ -162,6 +166,9 @@ func (r *Runner) handleWSDisconnect(ctx context.Context, nodeID string) { func (r *Runner) handleHeartbeatTick(ctx context.Context, nodeID *string, heartbeatTicker *time.Ticker, tryStartWebSocket func()) { if !r.hasAccessToken() { + if !r.discoveryRetry.ready(time.Now()) { + return + } if err := r.tryRegister(ctx, nodeID); err != nil { slog.Error("agent discovery register failed", "error", err) } @@ -465,36 +472,46 @@ func (r *Runner) tryRestartOpenresty(ctx context.Context) { r.recordOpenrestyHealthy() } -func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error { +func (r *Runner) tryRegister(ctx context.Context, nodeID *string) (resultErr error) { + defer func() { + if resultErr != nil { + r.discoveryRetry.failed(time.Now()) + } + }() if strings.TrimSpace(r.Config.DiscoveryToken) == "" { - return errors.New("agent_token 为空且未配置 discovery_token") + return errors.New("agent token is empty and no discovery token is configured") + } + registrationKey, err := r.StateStore.EnsureRegistrationKey() + if err != nil { + return err } slog.Info("agent discovery registration started") - response, err := r.HeartbeatService.Register(ctx, r.HeartbeatCycle.NodePayload(ctx, *nodeID)) + response, err := r.HeartbeatService.Register(ctx, protocol.RegisterNodeRequest{ + NodePayload: r.HeartbeatCycle.NodePayload(ctx, *nodeID), + RegistrationKey: registrationKey, + }) if err != nil { return err } if response == nil || strings.TrimSpace(response.AccessToken) == "" || strings.TrimSpace(response.NodeID) == "" { - return errors.New("discovery register response 缺少 node_id 或 agent_token") + return errors.New("discovery registration response is missing node_id or agent_token") } - snapshot, err := r.StateStore.Load() - if err != nil { + if err = r.StateStore.SetRegisteredNodeID(response.NodeID); err != nil { return err } - snapshot.NodeID = response.NodeID - if err = r.StateStore.Save(snapshot); err != nil { - return err - } - r.Config.AccessToken = response.AccessToken - r.Config.DiscoveryToken = "" - if err = r.Config.Save(); err != nil { + updatedConfig := *r.Config + updatedConfig.AccessToken = response.AccessToken + updatedConfig.DiscoveryToken = "" + if err = updatedConfig.Save(); err != nil { return err } + *r.Config = updatedConfig r.HeartbeatService.SetToken(response.AccessToken) if r.WebSocketService != nil { r.WebSocketService.SetToken(response.AccessToken) } *nodeID = response.NodeID + r.discoveryRetry = discoveryRetry{} slog.Info("agent discovery registration succeeded", "node_id", response.NodeID) r.refreshOpenrestyHealth(ctx) if _, err = r.HeartbeatCycle.Perform(ctx, *nodeID, true, r); err != nil { diff --git a/internal/apps/agent/agent/runner_test.go b/internal/apps/agent/agent/runner_test.go index 1ec211a1..43d97cab 100644 --- a/internal/apps/agent/agent/runner_test.go +++ b/internal/apps/agent/agent/runner_test.go @@ -42,13 +42,17 @@ type fakeHeartbeatService struct { heartbeatResults []*protocol.HeartbeatResult heartbeatPayloads []protocol.NodePayload onHeartbeat func(int) + onRegister func() lastToken string } -func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) { +func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) { f.mu.Lock() defer f.mu.Unlock() f.registerCalls++ + if f.onRegister != nil { + f.onRegister() + } return f.registerResp, f.registerErr } diff --git a/internal/apps/agent/config/config.go b/internal/apps/agent/config/config.go index 06a2de7a..46b88377 100644 --- a/internal/apps/agent/config/config.go +++ b/internal/apps/agent/config/config.go @@ -15,6 +15,7 @@ import ( "strings" "time" + "github.com/Rain-kl/Wavelet/internal/apps/agent/localfile" "github.com/Rain-kl/Wavelet/internal/apps/edge/nodeip" "github.com/Rain-kl/Wavelet/pkg/util" ) @@ -406,7 +407,7 @@ func (cfg *Config) Save() error { if err != nil { return err } - return os.WriteFile(cfg.configPath, data, configFilePerm) + return localfile.WriteAtomic(cfg.configPath, data, configFilePerm) } func detectHostname() string { diff --git a/internal/apps/agent/heartbeat/service.go b/internal/apps/agent/heartbeat/service.go index eae24939..56b51083 100644 --- a/internal/apps/agent/heartbeat/service.go +++ b/internal/apps/agent/heartbeat/service.go @@ -11,14 +11,14 @@ import ( // RemoteClient is the interface that abstracts the remote API calls performed by Service. type RemoteClient interface { - RegisterNode(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) + RegisterNode(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) SetToken(token string) } // API abstracts registration and heartbeat operations used by Cycle. type API interface { - Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) + Register(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) SetToken(token string) } @@ -34,7 +34,7 @@ func New(client RemoteClient) *Service { } // Register sends a node registration request to the server. -func (s *Service) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) { +func (s *Service) Register(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) { return s.client.RegisterNode(ctx, payload) } diff --git a/internal/apps/agent/httpclient/client.go b/internal/apps/agent/httpclient/client.go index 4e4e2efb..e1b5358c 100644 --- a/internal/apps/agent/httpclient/client.go +++ b/internal/apps/agent/httpclient/client.go @@ -32,7 +32,7 @@ func New(baseURL string, token string, timeout time.Duration) *Client { } // RegisterNode registers the agent node with the control plane server. -func (c *Client) RegisterNode(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) { +func (c *Client) RegisterNode(ctx context.Context, payload protocol.RegisterNodeRequest) (*protocol.RegisterNodeResponse, error) { resp := protocol.APIResponse[protocol.RegisterNodeResponse]{} if err := c.base.PostJSON(ctx, "/api/v1/agent/nodes/register", payload, &resp); err != nil { return nil, err diff --git a/internal/apps/agent/localfile/atomic.go b/internal/apps/agent/localfile/atomic.go new file mode 100644 index 00000000..07e16ddc --- /dev/null +++ b/internal/apps/agent/localfile/atomic.go @@ -0,0 +1,66 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +// Package localfile persists agent identity and configuration without truncating existing files. +package localfile + +import ( + "errors" + "os" + "path/filepath" +) + +// WriteAtomic syncs the complete contents, replaces path, then syncs its containing directory. +// An error from the final sync means the replacement is visible but its durability is unconfirmed. +func WriteAtomic(path string, data []byte, perm os.FileMode) error { + return writeAtomic(path, data, perm, SyncDirectory) +} + +func writeAtomic(path string, data []byte, perm os.FileMode, syncDir func(string) error) error { + file, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".tmp-*") + if err != nil { + return err + } + defer func() { + _ = file.Close() // Best effort after an earlier write/close error. + _ = os.Remove(file.Name()) // Best effort; rename normally removes the temporary path. + }() + if err = file.Chmod(perm); err != nil { + return err + } + if _, err = file.Write(data); err != nil { + return err + } + if err = file.Sync(); err != nil { + return err + } + if err = file.Close(); err != nil { + return err + } + if err = os.Rename(file.Name(), path); err != nil { + return err + } + return syncDir(filepath.Dir(path)) +} + +// SyncDirectory confirms a directory entry update on platforms with directory flushing. +// Windows provides best-effort replacement without a portable directory flush guarantee. +func SyncDirectory(path string) (resultErr error) { + directory, err := os.Open(path) //nolint:gosec // callers provide the directory of a managed local state/config file + if err != nil { + return err + } + defer func() { + if closeErr := directory.Close(); resultErr == nil { + resultErr = closeErr + } + }() + info, err := directory.Stat() + if err != nil { + return err + } + if !info.IsDir() { + return errors.New("cannot sync a non-directory path") + } + return syncDirectory(directory) +} diff --git a/internal/apps/agent/localfile/atomic_test.go b/internal/apps/agent/localfile/atomic_test.go new file mode 100644 index 00000000..b0c81569 --- /dev/null +++ b/internal/apps/agent/localfile/atomic_test.go @@ -0,0 +1,88 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package localfile + +import ( + "errors" + "os" + "path/filepath" + "testing" +) + +func TestWriteAtomicReplacesCompleteFileAndCleansFailure(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "identity.json") + if err := WriteAtomic(path, []byte("old identity"), 0o600); err != nil { + t.Fatal(err) + } + if err := WriteAtomic(path, []byte("new identity"), 0o600); err != nil { + t.Fatal(err) + } + data, err := os.ReadFile(path) + if err != nil || string(data) != "new identity" { + t.Errorf("WriteAtomic replacement = %q, %v; want complete new identity", data, err) + } + blocked := filepath.Join(dir, "directory") + if err = os.Mkdir(blocked, 0o700); err != nil { + t.Fatal(err) + } + if err = WriteAtomic(blocked, []byte("cannot replace directory"), 0o600); err == nil { + t.Fatal("WriteAtomic(directory) succeeded, want rename failure") + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Errorf("WriteAtomic(after failure) directory entries = %d, want 2 with no temporary file", len(entries)) + } +} + +func TestWriteAtomicReportsDirectorySyncFailureAfterReplacement(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "identity.json") + if err := os.WriteFile(path, []byte("old identity"), 0o600); err != nil { + t.Fatal(err) + } + syncFailure := errors.New("directory I/O failure") + syncCalls := 0 + err := writeAtomic(path, []byte("new identity"), 0o600, func(gotDir string) error { + syncCalls++ + if gotDir != dir { + t.Errorf("writeAtomic sync path = %q, want %q", gotDir, dir) + } + data, readErr := os.ReadFile(path) + if readErr != nil || string(data) != "new identity" { + t.Errorf("Directory sync ran before replacement: contents = %q, error %v", data, readErr) + } + return syncFailure + }) + if !errors.Is(err, syncFailure) || syncCalls != 1 { + t.Errorf("writeAtomic(directory sync failure) = %v, %d sync calls; want I/O failure, 1", err, syncCalls) + } + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Errorf("writeAtomic(failed directory sync) left %d entries, want only destination", len(entries)) + } +} + +func TestSyncDirectoryRejectsMissingAndNonDirectoryPaths(t *testing.T) { + dir := t.TempDir() + if err := SyncDirectory(dir); err != nil { + t.Fatalf("SyncDirectory(existing directory) = %v, want nil", err) + } + if err := SyncDirectory(filepath.Join(dir, "missing")); !errors.Is(err, os.ErrNotExist) { + t.Errorf("SyncDirectory(missing directory) = %v, want not-exist error", err) + } + path := filepath.Join(dir, "file") + if err := os.WriteFile(path, []byte("contents"), 0o600); err != nil { + t.Fatal(err) + } + if err := SyncDirectory(path); err == nil { + t.Error("SyncDirectory(regular file) = nil, want non-directory error") + } +} diff --git a/internal/apps/agent/localfile/sync_directory_unix.go b/internal/apps/agent/localfile/sync_directory_unix.go new file mode 100644 index 00000000..3175073c --- /dev/null +++ b/internal/apps/agent/localfile/sync_directory_unix.go @@ -0,0 +1,12 @@ +//go:build !windows + +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package localfile + +import "os" + +func syncDirectory(directory *os.File) error { + return directory.Sync() +} diff --git a/internal/apps/agent/localfile/sync_directory_windows.go b/internal/apps/agent/localfile/sync_directory_windows.go new file mode 100644 index 00000000..1332305d --- /dev/null +++ b/internal/apps/agent/localfile/sync_directory_windows.go @@ -0,0 +1,12 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package localfile + +import "os" + +// Windows does not support a portable directory fsync through os.File.Sync. +// The file is flushed before replacement; persisting the directory entry is best effort. +func syncDirectory(_ *os.File) error { + return nil +} diff --git a/internal/apps/agent/protocol/alias.go b/internal/apps/agent/protocol/alias.go index 653670b2..1a82c9c5 100644 --- a/internal/apps/agent/protocol/alias.go +++ b/internal/apps/agent/protocol/alias.go @@ -51,6 +51,9 @@ type NodeHealthEvent = pkgprotocol.NodeHealthEvent // RegisterNodeResponse is an alias for pkgprotocol.RegisterNodeResponse. type RegisterNodeResponse = pkgprotocol.RegisterNodeResponse +// RegisterNodeRequest is an enrollment-only request, including its retry credential. +type RegisterNodeRequest = pkgprotocol.RegisterNodeRequest + // ApplyLogPayload is an alias for pkgprotocol.ApplyLogPayload. type ApplyLogPayload = pkgprotocol.ApplyLogPayload diff --git a/internal/apps/agent/state/registration_test.go b/internal/apps/agent/state/registration_test.go new file mode 100644 index 00000000..dc05303c --- /dev/null +++ b/internal/apps/agent/state/registration_test.go @@ -0,0 +1,84 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package state + +import ( + "encoding/hex" + "errors" + "path/filepath" + "testing" + + "github.com/Rain-kl/Wavelet/internal/apps/agent/localfile" +) + +func TestRegistrationKeySurvivesStaleSnapshotAndRestart(t *testing.T) { + path := filepath.Join(t.TempDir(), "state.json") + store := NewStore(path) + stale, err := store.Load() + if err != nil { + t.Fatal(err) + } + key, err := store.EnsureRegistrationKey() + if err != nil { + t.Fatal(err) + } + if err = store.SetRegisteredNodeID("node-server-assigned"); err != nil { + t.Fatal(err) + } + decoded, err := hex.DecodeString(key) + if err != nil || len(decoded) != 32 { + t.Fatalf("EnsureRegistrationKey() length = %d, error %v; want 32, nil", len(decoded), err) + } + stale.LastError = "background sync failed" + if err = store.Save(stale); err != nil { + t.Fatal(err) + } + restartedKey, err := NewStore(path).EnsureRegistrationKey() + if err != nil { + t.Fatal(err) + } + if restartedKey != key { + t.Error("EnsureRegistrationKey(after stale save and restart) changed enrollment identity") + } + snapshot, err := NewStore(path).Load() + if err != nil { + t.Fatal(err) + } + if snapshot.NodeID != "node-server-assigned" { + t.Errorf("Load(after stale save).NodeID = %q, want server-assigned ID", snapshot.NodeID) + } +} + +func TestExistingRegistrationKeyRequiresDirectorySyncBeforeUse(t *testing.T) { + path := filepath.Join(t.TempDir(), "state.json") + store := NewStore(path) + key, err := store.EnsureRegistrationKey() + if err != nil { + t.Fatal(err) + } + // Model a prior successful rename whose final directory flush was not confirmed. + syncFailure := errors.New("directory I/O failure") + syncCalls := 0 + store.syncDirectory = func(dir string) error { + syncCalls++ + if dir != filepath.Dir(path) { + t.Errorf("EnsureRegistrationKey sync directory = %q, want state parent", dir) + } + return syncFailure + } + for range 2 { + got, retryErr := store.EnsureRegistrationKey() + if got != "" || !errors.Is(retryErr, syncFailure) { + t.Errorf("EnsureRegistrationKey(unconfirmed directory) returned usable key or lost error: %v", retryErr) + } + } + if syncCalls != 2 { + t.Errorf("EnsureRegistrationKey(repeated failures) sync calls = %d, want 2", syncCalls) + } + store.syncDirectory = localfile.SyncDirectory + confirmedKey, err := store.EnsureRegistrationKey() + if err != nil || confirmedKey != key { + t.Errorf("EnsureRegistrationKey(after successful sync) changed identity or failed: %v", err) + } +} diff --git a/internal/apps/agent/state/state.go b/internal/apps/agent/state/state.go index e57b45b6..2498ea0a 100644 --- a/internal/apps/agent/state/state.go +++ b/internal/apps/agent/state/state.go @@ -10,12 +10,15 @@ import ( "os" "path/filepath" "sync" + + "github.com/Rain-kl/Wavelet/internal/apps/agent/localfile" ) const ( - stateDirPerm = 0o750 - stateFilePerm = 0o600 - nodeIDRandomBytes = 8 + stateDirPerm = 0o750 + stateFilePerm = 0o600 + nodeIDRandomBytes = 8 + registrationKeyBytes = 32 ) // PagesDeployment records a Pages project tracked by the agent and the last @@ -31,6 +34,7 @@ type PagesDeployment struct { // Snapshot represents the state of the agent at a given point in time. type Snapshot struct { NodeID string `json:"node_id"` + RegistrationKey string `json:"registration_key,omitempty"` CurrentVersion string `json:"current_version"` CurrentChecksum string `json:"current_checksum"` PagesDeployments []PagesDeployment `json:"pages_deployments"` @@ -47,15 +51,42 @@ type Snapshot struct { AccessLogOffset int64 `json:"access_log_offset"` } +// EnsureRegistrationKey persists the private retry credential before enrollment can reach the server. +func (s *Store) EnsureRegistrationKey() (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + snapshot, err := s.loadUnlocked() + if err != nil { + return "", err + } + if snapshot.RegistrationKey != "" { + // A previous rename may have succeeded even though its directory sync failed. + if err = s.syncDirectory(filepath.Dir(s.path)); err != nil { + return "", err + } + return snapshot.RegistrationKey, nil + } + key := make([]byte, registrationKeyBytes) + if _, err = rand.Read(key); err != nil { + return "", err + } + snapshot.RegistrationKey = hex.EncodeToString(key) + if err = s.saveUnlocked(snapshot); err != nil { + return "", err + } + return snapshot.RegistrationKey, nil +} + // Store manages the storage and retrieval of the agent state snapshot. type Store struct { - path string - mu sync.Mutex + path string + mu sync.Mutex + syncDirectory func(string) error } // NewStore creates a new Store instance at the given path. func NewStore(path string) *Store { - return &Store{path: filepath.Clean(path)} + return &Store{path: filepath.Clean(path), syncDirectory: localfile.SyncDirectory} } // Load loads the snapshot from the store. @@ -91,6 +122,29 @@ func (s *Store) EnsureNodeID() (string, error) { func (s *Store) Save(snapshot *Snapshot) error { s.mu.Lock() defer s.mu.Unlock() + current, err := s.loadUnlocked() + if err != nil { + return err + } + updated := *snapshot + + // Other workers can hold a snapshot loaded before enrollment first persisted its key. + if current.RegistrationKey != "" { + updated.RegistrationKey = current.RegistrationKey + updated.NodeID = current.NodeID + } + return s.saveUnlocked(&updated) +} + +// SetRegisteredNodeID commits the server-assigned identity without overwriting concurrent runtime state. +func (s *Store) SetRegisteredNodeID(nodeID string) error { + s.mu.Lock() + defer s.mu.Unlock() + snapshot, err := s.loadUnlocked() + if err != nil { + return err + } + snapshot.NodeID = nodeID return s.saveUnlocked(snapshot) } @@ -120,7 +174,7 @@ func (s *Store) saveUnlocked(snapshot *Snapshot) error { if err != nil { return err } - return os.WriteFile(s.path, data, stateFilePerm) + return localfile.WriteAtomic(s.path, data, stateFilePerm) } func newNodeID() (string, error) { diff --git a/internal/apps/openflare/agent/errs.go b/internal/apps/openflare/agent/errs.go index 21bc37f7..9244622e 100644 --- a/internal/apps/openflare/agent/errs.go +++ b/internal/apps/openflare/agent/errs.go @@ -4,17 +4,19 @@ package agent const ( - errMissingAgentToken = "缺少 Agent Token" //nolint:gosec // false positive: this is an error message, not hardcoded credentials - errInvalidAgentToken = "无权进行此操作，Agent Token 无效" //nolint:gosec // false positive: this is an error message, not hardcoded credentials - errInvalidDiscoveryToken = "无权进行此操作，注册 Token 无效" //nolint:gosec // false positive: this is an error message, not hardcoded credentials - errNodeMissingFromContext = "Node object missing from context" - errNoActiveConfig = "当前没有激活版本" - errNodeNotFound = "节点不存在" - errNodeIDRequired = "node_id 不能为空" - errVersionRequired = "version 不能为空" - errInvalidApplyResult = "result 仅支持 success、warning 或 failed" - errIPRequired = "ip 不能为空" - errIPInvalid = "ip 格式无效" - errAgentVersionRequired = "version 不能为空" - errNodeIDConflict = "节点标识生成冲突，请重试" + errMissingAgentToken = "缺少 Agent Token" //nolint:gosec // false positive: this is an error message, not hardcoded credentials + errInvalidAgentToken = "无权进行此操作，Agent Token 无效" //nolint:gosec // false positive: this is an error message, not hardcoded credentials + errInvalidDiscoveryToken = "无权进行此操作，注册 Token 无效" //nolint:gosec // false positive: this is an error message, not hardcoded credentials + errNodeMissingFromContext = "Node object missing from context" + errNoActiveConfig = "当前没有激活版本" + errNodeNotFound = "节点不存在" + errNodeIDRequired = "node_id 不能为空" + errVersionRequired = "version 不能为空" + errInvalidApplyResult = "result 仅支持 success、warning 或 failed" + errIPRequired = "ip 不能为空" + errIPInvalid = "ip 格式无效" + errAgentVersionRequired = "version 不能为空" + errNodeIDConflict = "节点标识生成冲突，请重试" + errRegistrationKeyInvalid = "invalid discovery registration key" + errRegistrationCredentialChanged = "discovery enrollment credential has changed; use the assigned agent token" ) diff --git a/internal/apps/openflare/agent/logics.go b/internal/apps/openflare/agent/logics.go index 7967963d..a922363b 100644 --- a/internal/apps/openflare/agent/logics.go +++ b/internal/apps/openflare/agent/logics.go @@ -5,6 +5,7 @@ package agent import ( "context" + "crypto/subtle" "errors" "strings" "time" @@ -41,18 +42,14 @@ func RegisterWithAccessToken(ctx context.Context, authNode *model.OpenFlareNode, } // RegisterWithDiscovery registers a new node using the global discovery token. -func RegisterWithDiscovery(ctx context.Context, payload NodePayload) (*RegistrationResponse, error) { +func RegisterWithDiscovery(ctx context.Context, payload NodePayload, registrationKey string) (*RegistrationResponse, error) { _ = ofgeoip.EnsureRuntimeProvider(ctx) payload = normalizeNodePayload(payload) if err := validateNodePayload(payload); err != nil { return nil, err } - nodeID, err := newServerNodeID() - if err != nil { - return nil, err - } - accessToken, err := newRandomToken() + nodeID, accessToken, err := newDiscoveryIdentity(registrationKey) if err != nil { return nil, err } @@ -73,7 +70,15 @@ func RegisterWithDiscovery(ctx context.Context, payload NodePayload) (*Registrat } applyNodeRuntime(ctx, record, payload, false) - if err = repository.CreateOpenFlareNode(ctx, record); err != nil { + if registrationKey != "" { + record, err = repository.CreateOrGetOpenFlareDiscoveryNode(ctx, record) + if err == nil && (record.NodeType != "edge_node" || subtle.ConstantTimeCompare([]byte(record.AccessToken), []byte(accessToken)) != 1) { + return nil, errors.New(errRegistrationCredentialChanged) + } + } else { + err = repository.CreateOpenFlareNode(ctx, record) + } + if err != nil { if isUniqueConstraintError(err) { return nil, errors.New(errNodeIDConflict) } @@ -87,6 +92,18 @@ func RegisterWithDiscovery(ctx context.Context, payload NodePayload) (*Registrat }, nil } +func newDiscoveryIdentity(registrationKey string) (string, string, error) { + if registrationKey != "" { + return discoveryIdentity(registrationKey) + } + nodeID, err := newServerNodeID() + if err != nil { + return "", "", err + } + token, err := newRandomToken() + return nodeID, token, err +} + // HeartbeatNode updates runtime state and returns agent settings. func HeartbeatNode(ctx context.Context, authNode *model.OpenFlareNode, payload NodePayload) (*HeartbeatResponse, error) { _ = ofgeoip.EnsureRuntimeProvider(ctx) diff --git a/internal/apps/openflare/agent/registration.go b/internal/apps/openflare/agent/registration.go new file mode 100644 index 00000000..11534ce7 --- /dev/null +++ b/internal/apps/openflare/agent/registration.go @@ -0,0 +1,25 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "errors" +) + +const registrationKeyBytes = 32 + +// discoveryIdentity derives separate public identity and private credential from a random retry secret. +// Public node IDs cannot be used as registration keys to recover credentials. +func discoveryIdentity(key string) (string, string, error) { + decoded, err := hex.DecodeString(key) + if err != nil || len(decoded) != registrationKeyBytes { + return "", "", errors.New(errRegistrationKeyInvalid) + } + nodeHash := sha256.Sum256(append([]byte("openflare:discovery:node:v1:"), decoded...)) + tokenHash := sha256.Sum256(append([]byte("openflare:discovery:token:v1:"), decoded...)) + return "node-" + base64.RawURLEncoding.EncodeToString(nodeHash[:]), hex.EncodeToString(tokenHash[:]), nil +} diff --git a/internal/apps/openflare/agent/registration_test.go b/internal/apps/openflare/agent/registration_test.go new file mode 100644 index 00000000..d0a24e21 --- /dev/null +++ b/internal/apps/openflare/agent/registration_test.go @@ -0,0 +1,168 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package agent + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + agentclient "github.com/Rain-kl/Wavelet/internal/apps/agent/httpclient" + db "github.com/Rain-kl/Wavelet/internal/infra/persistence" + "github.com/Rain-kl/Wavelet/internal/model" + "github.com/Rain-kl/Wavelet/internal/repository" + "github.com/Rain-kl/Wavelet/internal/shared/response" + "github.com/Rain-kl/Wavelet/pkg/protocol" +) + +func TestRegistrationResponseServerClientContract(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(response.OK(RegistrationResponse{NodeID: "node-contract", AccessToken: "issued", Name: "edge"})); err != nil { + t.Error(err) + } + })) + t.Cleanup(server.Close) + client := agentclient.New(server.URL, "discovery", time.Second) + got, err := client.RegisterNode(context.Background(), protocol.RegisterNodeRequest{}) + if err != nil { + t.Fatalf("RegisterNode(server response) error = %v", err) + } + if got.NodeID != "node-contract" || got.AccessToken != "issued" || got.Name != "edge" { + t.Errorf("RegisterNode(server response) = %+v, want node-contract/issued/edge", got) + } + data, err := json.Marshal(RegistrationResponse{AccessToken: "issued"}) + if err != nil { + t.Fatal(err) + } + var legacy struct { + AccessToken string `json:"access_token"` + } + if err = json.Unmarshal(data, &legacy); err != nil || legacy.AccessToken != "issued" { + t.Errorf("Decode legacy server response = %q, %v; want issued, nil", legacy.AccessToken, err) + } + var documented RegistrationResponseDocs + if err = json.Unmarshal(data, &documented); err != nil { + t.Fatal(err) + } + if documented.AgentToken != "issued" || documented.AccessToken != documented.AgentToken { + t.Error("RegistrationResponseDocs does not describe both emitted token aliases") + } +} + +func TestDiscoveryRetriesKeepOneIdentityAndPreserveChanges(t *testing.T) { + t.Cleanup(setupAgentAuthTestDB(t)) + ctx := context.Background() + key := strings.Repeat("a1", 32) + payload := NodePayload{IP: "10.0.0.8", Version: "test", Name: "edge"} + first, err := RegisterWithDiscovery(ctx, payload, key) + if err != nil { + t.Fatal(err) + } + record, err := repository.GetOpenFlareNodeByNodeID(ctx, first.NodeID) + if err != nil { + t.Fatal(err) + } + record.Name = "administrator name" + record.UpdateRequested = true + if err = repository.SaveOpenFlareNode(ctx, record); err != nil { + t.Fatal(err) + } + second, err := RegisterWithDiscovery(ctx, payload, key) + if err != nil { + t.Fatal(err) + } + if first.NodeID != second.NodeID || first.AccessToken != second.AccessToken { + t.Errorf("RegisterWithDiscovery(retry) identity = %+v, want original identity", second) + } + record, err = repository.GetOpenFlareNodeByNodeID(ctx, first.NodeID) + if err != nil { + t.Fatal(err) + } + if record.Name != "administrator name" || !record.UpdateRequested { + t.Errorf("RegisterWithDiscovery(retry) overwrote administrator metadata: %+v", record) + } + other, err := RegisterWithDiscovery(ctx, payload, strings.Repeat("b2", 32)) + if err != nil { + t.Fatal(err) + } + if other.NodeID == first.NodeID || other.AccessToken == first.AccessToken { + t.Error("RegisterWithDiscovery(distinct key, same host/IP) reused identity") + } + var count int64 + if err = db.DB(ctx).Model(&model.OpenFlareNode{}).Count(&count).Error; err != nil { + t.Fatal(err) + } + if count != 2 { + t.Errorf("RegisterWithDiscovery(two keys, retry) node count = %d, want 2", count) + } +} + +func TestDiscoveryConcurrentRetriesAndRevokedCredential(t *testing.T) { + t.Cleanup(setupAgentAuthTestDB(t)) + ctx := context.Background() + sqlDB, err := db.DB(ctx).DB() + if err != nil { + t.Fatal(err) + } + sqlDB.SetMaxOpenConns(1) + key := strings.Repeat("c3", 32) + payload := NodePayload{IP: "10.0.0.8", Version: "test"} + const attempts = 12 + results := make(chan *RegistrationResponse, attempts) + errors := make(chan error, attempts) + var wg sync.WaitGroup + for range attempts { + wg.Go(func() { + result, registrationErr := RegisterWithDiscovery(ctx, payload, key) + results <- result + errors <- registrationErr + }) + } + wg.Wait() + close(results) + close(errors) + for registrationErr := range errors { + if registrationErr != nil { + t.Fatalf("Concurrent RegisterWithDiscovery error = %v", registrationErr) + } + } + var first *RegistrationResponse + for result := range results { + if first == nil { + first = result + } + if result.NodeID != first.NodeID || result.AccessToken != first.AccessToken { + t.Error("Concurrent enrollment returned different identities") + } + } + var count int64 + if err = db.DB(ctx).Model(&model.OpenFlareNode{}).Count(&count).Error; err != nil { + t.Fatal(err) + } + if count != 1 { + t.Errorf("Concurrent enrollment node count = %d, want 1", count) + } + record, err := repository.GetOpenFlareNodeByNodeID(ctx, first.NodeID) + if err != nil { + t.Fatal(err) + } + record.AccessToken = "rotated-by-administrator" + if err = repository.SaveOpenFlareNode(ctx, record); err != nil { + t.Fatal(err) + } + if result, retryErr := RegisterWithDiscovery(ctx, payload, key); retryErr == nil || result != nil { + t.Errorf("RegisterWithDiscovery(rotated credential) = %+v, %v; want rejection", result, retryErr) + } + for _, invalid := range []string{"short", strings.Repeat("g", 64), first.NodeID} { + if result, retryErr := RegisterWithDiscovery(ctx, payload, invalid); retryErr == nil || result != nil { + t.Errorf("RegisterWithDiscovery(invalid/public key) = %+v, %v; want rejection", result, retryErr) + } + } +} diff --git a/internal/apps/openflare/agent/routers.go b/internal/apps/openflare/agent/routers.go index 7ea38574..15cfc9a5 100644 --- a/internal/apps/openflare/agent/routers.go +++ b/internal/apps/openflare/agent/routers.go @@ -22,16 +22,17 @@ import ( // @Accept json // @Produce json // @Security AgentTokenAuth -// @Param body body agent.NodePayload true "节点上报数据" -// @Success 200 {object} response.Any{data=agent.RegistrationResponse} "注册成功" +// @Param body body protocol.RegisterNodeRequest true "节点上报数据；registration_key 为首次注册重试凭证" +// @Success 200 {object} response.Any{data=agent.RegistrationResponseDocs} "注册成功" // @Failure 400 {object} response.Any "参数错误" // @Failure 401 {object} response.Any "Token 无效" // @Router /api/v1/agent/nodes/register [post] func RegisterHandler(c *gin.Context) { - var payload NodePayload - if !apiutil.BindJSON(c, &payload) { + var request protocol.RegisterNodeRequest + if !apiutil.BindJSON(c, &request) { return } + payload := request.NodePayload payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr) var ( @@ -41,7 +42,7 @@ func RegisterHandler(c *gin.Context) { if authNode, ok := NodeFromContext(c); ok { result, err = RegisterWithAccessToken(c.Request.Context(), authNode, payload) } else { - result, err = RegisterWithDiscovery(c.Request.Context(), payload) + result, err = RegisterWithDiscovery(c.Request.Context(), payload, request.RegistrationKey) } if apiutil.AbortBadRequestOnError(c, err) { return diff --git a/internal/apps/openflare/agent/types.go b/internal/apps/openflare/agent/types.go index 52c76203..a3eece04 100644 --- a/internal/apps/openflare/agent/types.go +++ b/internal/apps/openflare/agent/types.go @@ -7,6 +7,7 @@ import ( "time" "github.com/Rain-kl/Wavelet/internal/model" + "github.com/Rain-kl/Wavelet/pkg/protocol" ) const ( @@ -17,9 +18,13 @@ const ( ) // RegistrationResponse is returned after agent registration. -// Server uses access_token; the agent client expects agent_token via RegisterNodeResponse. -type RegistrationResponse struct { +type RegistrationResponse = protocol.RegisterNodeResponse + +// RegistrationResponseDocs describes both token aliases emitted by the custom JSON marshaler. +// It is used only by Swagger; the runtime contract remains protocol.RegisterNodeResponse. +type RegistrationResponseDocs struct { NodeID string `json:"node_id"` + AgentToken string `json:"agent_token"` AccessToken string `json:"access_token"` Name string `json:"name"` } diff --git a/internal/repository/openflare_node_registration.go b/internal/repository/openflare_node_registration.go new file mode 100644 index 00000000..4c97dfc7 --- /dev/null +++ b/internal/repository/openflare_node_registration.go @@ -0,0 +1,28 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package repository + +import ( + "context" + "errors" + + db "github.com/Rain-kl/Wavelet/internal/infra/persistence" + "github.com/Rain-kl/Wavelet/internal/model" + "gorm.io/gorm/clause" +) + +// CreateOrGetOpenFlareDiscoveryNode inserts an enrollment once, without changing an existing identity. +func CreateOrGetOpenFlareDiscoveryNode(ctx context.Context, node *model.OpenFlareNode) (*model.OpenFlareNode, error) { + conn := db.DB(ctx) + if conn == nil { + return nil, errors.New(errDatabaseNotInitialized) + } + if err := conn.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "node_id"}}, + DoNothing: true, + }).Create(node).Error; err != nil { + return nil, err + } + return GetOpenFlareNodeByNodeID(ctx, node.NodeID) +} diff --git a/pkg/protocol/agent.go b/pkg/protocol/agent.go index cd0ba6ee..19183d5b 100644 --- a/pkg/protocol/agent.go +++ b/pkg/protocol/agent.go @@ -4,7 +4,11 @@ // Package protocol defines the communication protocol between OpenFlare server, agent, and relay components. package protocol -import "encoding/json" +import ( + "encoding/json" + "errors" + "strings" +) // APIResponse is a generic API response wrapper. type APIResponse[T any] struct { @@ -166,6 +170,12 @@ type NodeHealthEvent struct { Metadata map[string]string `json:"metadata,omitempty"` } +// RegisterNodeRequest carries enrollment data separately from heartbeat telemetry. +type RegisterNodeRequest struct { + NodePayload + RegistrationKey string `json:"registration_key,omitempty"` +} + // RegisterNodeResponse is the node registration response. type RegisterNodeResponse struct { NodeID string `json:"node_id"` @@ -173,6 +183,38 @@ type RegisterNodeResponse struct { Name string `json:"name"` } +// MarshalJSON retains access_token for consumers of older server responses. +func (r RegisterNodeResponse) MarshalJSON() ([]byte, error) { + type response RegisterNodeResponse + return json.Marshal(struct { + response + LegacyAccessToken string `json:"access_token"` + }{response(r), r.AccessToken}) +} + +// UnmarshalJSON accepts either spelling used by released servers. +func (r *RegisterNodeResponse) UnmarshalJSON(data []byte) error { + type response RegisterNodeResponse + var decoded struct { + response + LegacyAccessToken string `json:"access_token"` + } + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + canonical := strings.TrimSpace(decoded.AccessToken) + legacy := strings.TrimSpace(decoded.LegacyAccessToken) + if canonical != "" && legacy != "" && canonical != legacy { + return errors.New("registration response contains conflicting agent tokens") + } + *r = RegisterNodeResponse(decoded.response) + r.AccessToken = canonical + if canonical == "" { + r.AccessToken = legacy + } + return nil +} + // ActiveConfigResponse is the active configuration response. type ActiveConfigResponse struct { Version string `json:"version"` diff --git a/pkg/protocol/agent_registration_test.go b/pkg/protocol/agent_registration_test.go new file mode 100644 index 00000000..61b2d04d --- /dev/null +++ b/pkg/protocol/agent_registration_test.go @@ -0,0 +1,54 @@ +// Copyright 2026 Arctel.net +// SPDX-License-Identifier: Apache-2.0 + +package protocol + +import ( + "encoding/json" + "strings" + "testing" +) + +func TestRegisterNodeResponseTokenAliases(t *testing.T) { + for _, tt := range []struct { + name string + body string + want string + wantErr bool + }{ + {name: "canonical", body: `{"node_id":"node-1","agent_token":"issued"}`, want: "issued"}, + {name: "legacy", body: `{"node_id":"node-1","access_token":"issued"}`, want: "issued"}, + {name: "matching aliases", body: `{"agent_token":"issued","access_token":"issued"}`, want: "issued"}, + {name: "conflicting aliases", body: `{"agent_token":"a","access_token":"b"}`, wantErr: true}, + {name: "invalid JSON", body: `{`, wantErr: true}, + } { + t.Run(tt.name, func(t *testing.T) { + var got RegisterNodeResponse + err := json.Unmarshal([]byte(tt.body), &got) + if (err != nil) != tt.wantErr { + t.Fatalf("Unmarshal(%s) error = %v, want error %t", tt.name, err, tt.wantErr) + } + if got.AccessToken != tt.want { + t.Errorf("Unmarshal(%s).AccessToken = %q, want %q", tt.name, got.AccessToken, tt.want) + } + }) + } +} + +func TestRegistrationRequestKeepsRetryCredentialOutOfTelemetry(t *testing.T) { + request := RegisterNodeRequest{NodePayload: NodePayload{NodeID: "node-1"}, RegistrationKey: "private"} + registration, err := json.Marshal(request) + if err != nil { + t.Fatal(err) + } + if !strings.Contains(string(registration), `"registration_key":"private"`) { + t.Errorf("Marshal(registration) = %s, want retry credential", registration) + } + telemetry, err := json.Marshal(request.NodePayload) + if err != nil { + t.Fatal(err) + } + if strings.Contains(string(telemetry), "registration_key") || strings.Contains(string(telemetry), "private") { + t.Errorf("Marshal(heartbeat) = %s, want no enrollment credential", telemetry) + } +}