mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-11-12 21:44:22 +08:00
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
This commit is contained in:
@@ -19,6 +19,8 @@ sidebar: false
|
|||||||
|
|
||||||
### 🛠 修复
|
### 🛠 修复
|
||||||
|
|
||||||
|
- 修复 Agent 自动注册响应凭证字段不一致导致节点不断重复创建的问题;新增持久化注册重试凭证与退避,响应丢失、并发重试或配置写入失败后仍复用原节点,且不会覆盖管理员修改或恢复已更换的节点凭证。支持目录刷盘的平台会确认原子替换后的目录持久化,失败时阻止使用未确认的注册凭证;接口文档同步列出两种兼容凭证字段。
|
||||||
|
|
||||||
- SQLite 文件数据库启动时自动启用 WAL 模式,改善读写并发;数据库导出改用一致性快照,确保包含 WAL 中已提交的数据。
|
- SQLite 文件数据库启动时自动启用 WAL 模式,改善读写并发;数据库导出改用一致性快照,确保包含 WAL 中已提交的数据。
|
||||||
|
|
||||||
- 修复 IP 组按规则自动抓取时会收录内网、回环及链路本地等地址的问题;规则测试同步过滤这些地址,历史抓取结果会在下次同步时清理。
|
- 修复 IP 组按规则自动抓取时会收录内网、回环及链路本地等地址的问题;规则测试同步过滤这些地址,历史抓取结果会在下次同步时清理。
|
||||||
|
|||||||
@@ -37,10 +37,13 @@ Agent 主要由以下核心子模块组成,共同配合完成其完整的生
|
|||||||
Agent 在生命周期中主要通过 **基于 Token 的自动注册** 和 **心跳/WebSocket 双通道** 与控制面通信。
|
Agent 在生命周期中主要通过 **基于 Token 的自动注册** 和 **心跳/WebSocket 双通道** 与控制面通信。
|
||||||
|
|
||||||
### 1. 自动注册流程
|
### 1. 自动注册流程
|
||||||
若 Agent 启动时本地 `agent.json` 的 `access_token` 为空,但配置了 `discovery_token`,将触发自动注册流程:
|
若 Agent 启动时本地 `agent.json` 的 `agent_token` 为空,但配置了 `discovery_token`,将触发自动注册流程:
|
||||||
1. Agent 向控制面 `/api/v1/agent/nodes/register` 发送注册请求,携带本地硬件摘要、IP 及主机名。
|
1. Agent 在本地状态文件原子保存随机 256 位 `registration_key`,再向控制面 `/api/v1/agent/nodes/register` 发送注册请求,携带该重试凭证、本地硬件摘要、IP 及主机名。支持目录刷盘的平台会在文件刷盘及重命名后刷新所在目录;目录刷盘失败会返回错误,已有凭证也必须在重试时确认目录刷盘成功后才能发送。Windows 保留文件刷盘与替换,但目录持久化仅为尽力而为;上述行为针对所在目录已存在的文件更新,不保证所有文件系统或突然断电场景下的完整持久性。该凭证仅用于注册,不进入心跳、WebSocket 或观测数据,必须与节点 Token 一样保密。
|
||||||
2. Server 校验 `discovery_token` 有效后,在数据库生成唯一的 `NodeID` 与专属 `AccessToken`(即 `agent_token`)并返回。
|
2. Server 校验当前 `discovery_token` 有效后,从重试凭证使用不同的哈希域分别生成 `NodeID` 与专属 Token。相同凭证的重复或并发请求只保留一个节点,不按 IP、主机名或公开 `node_id` 合并,也不会覆盖管理员修改的名称、升级标志等数据。响应以 `agent_token` 为准,同时保留旧的 `access_token` 字段;新版 Agent 能识别两种字段,值冲突时拒绝响应。
|
||||||
3. Agent 将获取的专用 Token 写入本地配置文件,擦除一次性 `discovery_token`,后续所有的通信均基于专属 `AccessToken` 进行鉴权认证。
|
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. 双通道心跳与同步机制
|
### 2. 双通道心跳与同步机制
|
||||||
* **HTTP 轮询通道(兜底与探测)**:Agent 默认按设定的 `heartbeat_interval` 间隔发送 POST 心跳包。上报指标的同时获取当前激活版本的摘要信息(Version & Checksum)。
|
* **HTTP 轮询通道(兜底与探测)**:Agent 默认按设定的 `heartbeat_interval` 间隔发送 POST 心跳包。上报指标的同时获取当前激活版本的摘要信息(Version & Checksum)。
|
||||||
|
|||||||
+78
-4
@@ -4443,12 +4443,12 @@ const docTemplate = `{
|
|||||||
"summary": "注册或发现 Agent 节点",
|
"summary": "注册或发现 Agent 节点",
|
||||||
"parameters": [
|
"parameters": [
|
||||||
{
|
{
|
||||||
"description": "节点上报数据",
|
"description": "节点上报数据;registration_key 为首次注册重试凭证",
|
||||||
"name": "body",
|
"name": "body",
|
||||||
"in": "body",
|
"in": "body",
|
||||||
"required": true,
|
"required": true,
|
||||||
"schema": {
|
"schema": {
|
||||||
"$ref": "#/definitions/agent.NodePayload"
|
"$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
@@ -4464,7 +4464,7 @@ const docTemplate = `{
|
|||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
"data": {
|
"data": {
|
||||||
"$ref": "#/definitions/agent.RegistrationResponse"
|
"$ref": "#/definitions/agent.RegistrationResponseDocs"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -14989,12 +14989,15 @@ const docTemplate = `{
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"agent.RegistrationResponse": {
|
"agent.RegistrationResponseDocs": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
"access_token": {
|
"access_token": {
|
||||||
"type": "string"
|
"type": "string"
|
||||||
},
|
},
|
||||||
|
"agent_token": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
"name": {
|
"name": {
|
||||||
"type": "string"
|
"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": {
|
"github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
|
|||||||
+78
-4
@@ -4436,12 +4436,12 @@
|
|||||||
"summary": "注册或发现 Agent 节点",
|
"summary": "注册或发现 Agent 节点",
|
||||||
"parameters": [
|
"parameters": [
|
||||||
{
|
{
|
||||||
"description": "节点上报数据",
|
"description": "节点上报数据;registration_key 为首次注册重试凭证",
|
||||||
"name": "body",
|
"name": "body",
|
||||||
"in": "body",
|
"in": "body",
|
||||||
"required": true,
|
"required": true,
|
||||||
"schema": {
|
"schema": {
|
||||||
"$ref": "#/definitions/agent.NodePayload"
|
"$ref": "#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
],
|
],
|
||||||
@@ -4457,7 +4457,7 @@
|
|||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
"data": {
|
"data": {
|
||||||
"$ref": "#/definitions/agent.RegistrationResponse"
|
"$ref": "#/definitions/agent.RegistrationResponseDocs"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -14982,12 +14982,15 @@
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"agent.RegistrationResponse": {
|
"agent.RegistrationResponseDocs": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
"access_token": {
|
"access_token": {
|
||||||
"type": "string"
|
"type": "string"
|
||||||
},
|
},
|
||||||
|
"agent_token": {
|
||||||
|
"type": "string"
|
||||||
|
},
|
||||||
"name": {
|
"name": {
|
||||||
"type": "string"
|
"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": {
|
"github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup": {
|
||||||
"type": "object",
|
"type": "object",
|
||||||
"properties": {
|
"properties": {
|
||||||
|
|||||||
+53
-4
@@ -99,10 +99,12 @@ definitions:
|
|||||||
type: string
|
type: string
|
||||||
type: object
|
type: object
|
||||||
type: object
|
type: object
|
||||||
agent.RegistrationResponse:
|
agent.RegistrationResponseDocs:
|
||||||
properties:
|
properties:
|
||||||
access_token:
|
access_token:
|
||||||
type: string
|
type: string
|
||||||
|
agent_token:
|
||||||
|
type: string
|
||||||
name:
|
name:
|
||||||
type: string
|
type: string
|
||||||
node_id:
|
node_id:
|
||||||
@@ -955,6 +957,53 @@ definitions:
|
|||||||
total_size:
|
total_size:
|
||||||
type: integer
|
type: integer
|
||||||
type: object
|
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:
|
github_com_Rain-kl_Wavelet_pkg_protocol.WAFIPGroup:
|
||||||
properties:
|
properties:
|
||||||
checksum:
|
checksum:
|
||||||
@@ -7020,12 +7069,12 @@ paths:
|
|||||||
- application/json
|
- application/json
|
||||||
description: 使用节点 access token 重新注册,或使用全局 discovery token 发现新节点;请求头需携带 X-Agent-Token
|
description: 使用节点 access token 重新注册,或使用全局 discovery token 发现新节点;请求头需携带 X-Agent-Token
|
||||||
parameters:
|
parameters:
|
||||||
- description: 节点上报数据
|
- description: 节点上报数据;registration_key 为首次注册重试凭证
|
||||||
in: body
|
in: body
|
||||||
name: body
|
name: body
|
||||||
required: true
|
required: true
|
||||||
schema:
|
schema:
|
||||||
$ref: '#/definitions/agent.NodePayload'
|
$ref: '#/definitions/github_com_Rain-kl_Wavelet_pkg_protocol.RegisterNodeRequest'
|
||||||
produces:
|
produces:
|
||||||
- application/json
|
- application/json
|
||||||
responses:
|
responses:
|
||||||
@@ -7036,7 +7085,7 @@ paths:
|
|||||||
- $ref: '#/definitions/response.Any'
|
- $ref: '#/definitions/response.Any'
|
||||||
- properties:
|
- properties:
|
||||||
data:
|
data:
|
||||||
$ref: '#/definitions/agent.RegistrationResponse'
|
$ref: '#/definitions/agent.RegistrationResponseDocs'
|
||||||
type: object
|
type: object
|
||||||
"400":
|
"400":
|
||||||
description: 参数错误
|
description: 参数错误
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -23,7 +23,7 @@ import (
|
|||||||
|
|
||||||
// HeartbeatService handles node registration and periodic heartbeat reporting.
|
// HeartbeatService handles node registration and periodic heartbeat reporting.
|
||||||
type HeartbeatService interface {
|
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)
|
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
|
||||||
SetToken(token string)
|
SetToken(token string)
|
||||||
}
|
}
|
||||||
@@ -65,6 +65,7 @@ type Runner struct {
|
|||||||
configSyncRetry *configSyncRetry
|
configSyncRetry *configSyncRetry
|
||||||
restartOpenrestyNow bool
|
restartOpenrestyNow bool
|
||||||
websocketUpgradeEnabled bool
|
websocketUpgradeEnabled bool
|
||||||
|
discoveryRetry discoveryRetry
|
||||||
}
|
}
|
||||||
|
|
||||||
// Run starts the agent's main loop, performing heartbeats and upgrading to WebSocket when available.
|
// 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()
|
tryStartWebSocket()
|
||||||
|
|
||||||
for {
|
for {
|
||||||
|
if err := ctx.Err(); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
slog.Info("agent runner shutting down", "error", ctx.Err())
|
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()) {
|
func (r *Runner) handleHeartbeatTick(ctx context.Context, nodeID *string, heartbeatTicker *time.Ticker, tryStartWebSocket func()) {
|
||||||
if !r.hasAccessToken() {
|
if !r.hasAccessToken() {
|
||||||
|
if !r.discoveryRetry.ready(time.Now()) {
|
||||||
|
return
|
||||||
|
}
|
||||||
if err := r.tryRegister(ctx, nodeID); err != nil {
|
if err := r.tryRegister(ctx, nodeID); err != nil {
|
||||||
slog.Error("agent discovery register failed", "error", err)
|
slog.Error("agent discovery register failed", "error", err)
|
||||||
}
|
}
|
||||||
@@ -465,36 +472,46 @@ func (r *Runner) tryRestartOpenresty(ctx context.Context) {
|
|||||||
r.recordOpenrestyHealthy()
|
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) == "" {
|
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")
|
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 {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if response == nil || strings.TrimSpace(response.AccessToken) == "" || strings.TrimSpace(response.NodeID) == "" {
|
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 = r.StateStore.SetRegisteredNodeID(response.NodeID); err != nil {
|
||||||
if err != nil {
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
snapshot.NodeID = response.NodeID
|
updatedConfig := *r.Config
|
||||||
if err = r.StateStore.Save(snapshot); err != nil {
|
updatedConfig.AccessToken = response.AccessToken
|
||||||
return err
|
updatedConfig.DiscoveryToken = ""
|
||||||
}
|
if err = updatedConfig.Save(); err != nil {
|
||||||
r.Config.AccessToken = response.AccessToken
|
|
||||||
r.Config.DiscoveryToken = ""
|
|
||||||
if err = r.Config.Save(); err != nil {
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
*r.Config = updatedConfig
|
||||||
r.HeartbeatService.SetToken(response.AccessToken)
|
r.HeartbeatService.SetToken(response.AccessToken)
|
||||||
if r.WebSocketService != nil {
|
if r.WebSocketService != nil {
|
||||||
r.WebSocketService.SetToken(response.AccessToken)
|
r.WebSocketService.SetToken(response.AccessToken)
|
||||||
}
|
}
|
||||||
*nodeID = response.NodeID
|
*nodeID = response.NodeID
|
||||||
|
r.discoveryRetry = discoveryRetry{}
|
||||||
slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
|
slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
|
||||||
r.refreshOpenrestyHealth(ctx)
|
r.refreshOpenrestyHealth(ctx)
|
||||||
if _, err = r.HeartbeatCycle.Perform(ctx, *nodeID, true, r); err != nil {
|
if _, err = r.HeartbeatCycle.Perform(ctx, *nodeID, true, r); err != nil {
|
||||||
|
|||||||
@@ -42,13 +42,17 @@ type fakeHeartbeatService struct {
|
|||||||
heartbeatResults []*protocol.HeartbeatResult
|
heartbeatResults []*protocol.HeartbeatResult
|
||||||
heartbeatPayloads []protocol.NodePayload
|
heartbeatPayloads []protocol.NodePayload
|
||||||
onHeartbeat func(int)
|
onHeartbeat func(int)
|
||||||
|
onRegister func()
|
||||||
lastToken string
|
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()
|
f.mu.Lock()
|
||||||
defer f.mu.Unlock()
|
defer f.mu.Unlock()
|
||||||
f.registerCalls++
|
f.registerCalls++
|
||||||
|
if f.onRegister != nil {
|
||||||
|
f.onRegister()
|
||||||
|
}
|
||||||
return f.registerResp, f.registerErr
|
return f.registerResp, f.registerErr
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ import (
|
|||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/internal/apps/agent/localfile"
|
||||||
"github.com/Rain-kl/Wavelet/internal/apps/edge/nodeip"
|
"github.com/Rain-kl/Wavelet/internal/apps/edge/nodeip"
|
||||||
"github.com/Rain-kl/Wavelet/pkg/util"
|
"github.com/Rain-kl/Wavelet/pkg/util"
|
||||||
)
|
)
|
||||||
@@ -406,7 +407,7 @@ func (cfg *Config) Save() error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return os.WriteFile(cfg.configPath, data, configFilePerm)
|
return localfile.WriteAtomic(cfg.configPath, data, configFilePerm)
|
||||||
}
|
}
|
||||||
|
|
||||||
func detectHostname() string {
|
func detectHostname() string {
|
||||||
|
|||||||
@@ -11,14 +11,14 @@ import (
|
|||||||
|
|
||||||
// RemoteClient is the interface that abstracts the remote API calls performed by Service.
|
// RemoteClient is the interface that abstracts the remote API calls performed by Service.
|
||||||
type RemoteClient interface {
|
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)
|
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
|
||||||
SetToken(token string)
|
SetToken(token string)
|
||||||
}
|
}
|
||||||
|
|
||||||
// API abstracts registration and heartbeat operations used by Cycle.
|
// API abstracts registration and heartbeat operations used by Cycle.
|
||||||
type API interface {
|
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)
|
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
|
||||||
SetToken(token string)
|
SetToken(token string)
|
||||||
}
|
}
|
||||||
@@ -34,7 +34,7 @@ func New(client RemoteClient) *Service {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Register sends a node registration request to the server.
|
// 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)
|
return s.client.RegisterNode(ctx, payload)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -32,7 +32,7 @@ func New(baseURL string, token string, timeout time.Duration) *Client {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// RegisterNode registers the agent node with the control plane server.
|
// 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]{}
|
resp := protocol.APIResponse[protocol.RegisterNodeResponse]{}
|
||||||
if err := c.base.PostJSON(ctx, "/api/v1/agent/nodes/register", payload, &resp); err != nil {
|
if err := c.base.PostJSON(ctx, "/api/v1/agent/nodes/register", payload, &resp); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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()
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -51,6 +51,9 @@ type NodeHealthEvent = pkgprotocol.NodeHealthEvent
|
|||||||
// RegisterNodeResponse is an alias for pkgprotocol.RegisterNodeResponse.
|
// RegisterNodeResponse is an alias for pkgprotocol.RegisterNodeResponse.
|
||||||
type RegisterNodeResponse = 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.
|
// ApplyLogPayload is an alias for pkgprotocol.ApplyLogPayload.
|
||||||
type ApplyLogPayload = pkgprotocol.ApplyLogPayload
|
type ApplyLogPayload = pkgprotocol.ApplyLogPayload
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -10,12 +10,15 @@ import (
|
|||||||
"os"
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/internal/apps/agent/localfile"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
stateDirPerm = 0o750
|
stateDirPerm = 0o750
|
||||||
stateFilePerm = 0o600
|
stateFilePerm = 0o600
|
||||||
nodeIDRandomBytes = 8
|
nodeIDRandomBytes = 8
|
||||||
|
registrationKeyBytes = 32
|
||||||
)
|
)
|
||||||
|
|
||||||
// PagesDeployment records a Pages project tracked by the agent and the last
|
// 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.
|
// Snapshot represents the state of the agent at a given point in time.
|
||||||
type Snapshot struct {
|
type Snapshot struct {
|
||||||
NodeID string `json:"node_id"`
|
NodeID string `json:"node_id"`
|
||||||
|
RegistrationKey string `json:"registration_key,omitempty"`
|
||||||
CurrentVersion string `json:"current_version"`
|
CurrentVersion string `json:"current_version"`
|
||||||
CurrentChecksum string `json:"current_checksum"`
|
CurrentChecksum string `json:"current_checksum"`
|
||||||
PagesDeployments []PagesDeployment `json:"pages_deployments"`
|
PagesDeployments []PagesDeployment `json:"pages_deployments"`
|
||||||
@@ -47,15 +51,42 @@ type Snapshot struct {
|
|||||||
AccessLogOffset int64 `json:"access_log_offset"`
|
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.
|
// Store manages the storage and retrieval of the agent state snapshot.
|
||||||
type Store struct {
|
type Store struct {
|
||||||
path string
|
path string
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
|
syncDirectory func(string) error
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewStore creates a new Store instance at the given path.
|
// NewStore creates a new Store instance at the given path.
|
||||||
func NewStore(path string) *Store {
|
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.
|
// Load loads the snapshot from the store.
|
||||||
@@ -91,6 +122,29 @@ func (s *Store) EnsureNodeID() (string, error) {
|
|||||||
func (s *Store) Save(snapshot *Snapshot) error {
|
func (s *Store) Save(snapshot *Snapshot) error {
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
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)
|
return s.saveUnlocked(snapshot)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -120,7 +174,7 @@ func (s *Store) saveUnlocked(snapshot *Snapshot) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
return os.WriteFile(s.path, data, stateFilePerm)
|
return localfile.WriteAtomic(s.path, data, stateFilePerm)
|
||||||
}
|
}
|
||||||
|
|
||||||
func newNodeID() (string, error) {
|
func newNodeID() (string, error) {
|
||||||
|
|||||||
@@ -4,17 +4,19 @@
|
|||||||
package agent
|
package agent
|
||||||
|
|
||||||
const (
|
const (
|
||||||
errMissingAgentToken = "缺少 Agent Token" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
|
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
|
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
|
errInvalidDiscoveryToken = "无权进行此操作,注册 Token 无效" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
|
||||||
errNodeMissingFromContext = "Node object missing from context"
|
errNodeMissingFromContext = "Node object missing from context"
|
||||||
errNoActiveConfig = "当前没有激活版本"
|
errNoActiveConfig = "当前没有激活版本"
|
||||||
errNodeNotFound = "节点不存在"
|
errNodeNotFound = "节点不存在"
|
||||||
errNodeIDRequired = "node_id 不能为空"
|
errNodeIDRequired = "node_id 不能为空"
|
||||||
errVersionRequired = "version 不能为空"
|
errVersionRequired = "version 不能为空"
|
||||||
errInvalidApplyResult = "result 仅支持 success、warning 或 failed"
|
errInvalidApplyResult = "result 仅支持 success、warning 或 failed"
|
||||||
errIPRequired = "ip 不能为空"
|
errIPRequired = "ip 不能为空"
|
||||||
errIPInvalid = "ip 格式无效"
|
errIPInvalid = "ip 格式无效"
|
||||||
errAgentVersionRequired = "version 不能为空"
|
errAgentVersionRequired = "version 不能为空"
|
||||||
errNodeIDConflict = "节点标识生成冲突,请重试"
|
errNodeIDConflict = "节点标识生成冲突,请重试"
|
||||||
|
errRegistrationKeyInvalid = "invalid discovery registration key"
|
||||||
|
errRegistrationCredentialChanged = "discovery enrollment credential has changed; use the assigned agent token"
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ package agent
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"crypto/subtle"
|
||||||
"errors"
|
"errors"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -41,18 +42,14 @@ func RegisterWithAccessToken(ctx context.Context, authNode *model.OpenFlareNode,
|
|||||||
}
|
}
|
||||||
|
|
||||||
// RegisterWithDiscovery registers a new node using the global discovery token.
|
// 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)
|
_ = ofgeoip.EnsureRuntimeProvider(ctx)
|
||||||
payload = normalizeNodePayload(payload)
|
payload = normalizeNodePayload(payload)
|
||||||
if err := validateNodePayload(payload); err != nil {
|
if err := validateNodePayload(payload); err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
|
|
||||||
nodeID, err := newServerNodeID()
|
nodeID, accessToken, err := newDiscoveryIdentity(registrationKey)
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
accessToken, err := newRandomToken()
|
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
return nil, err
|
||||||
}
|
}
|
||||||
@@ -73,7 +70,15 @@ func RegisterWithDiscovery(ctx context.Context, payload NodePayload) (*Registrat
|
|||||||
}
|
}
|
||||||
applyNodeRuntime(ctx, record, payload, false)
|
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) {
|
if isUniqueConstraintError(err) {
|
||||||
return nil, errors.New(errNodeIDConflict)
|
return nil, errors.New(errNodeIDConflict)
|
||||||
}
|
}
|
||||||
@@ -87,6 +92,18 @@ func RegisterWithDiscovery(ctx context.Context, payload NodePayload) (*Registrat
|
|||||||
}, nil
|
}, 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.
|
// HeartbeatNode updates runtime state and returns agent settings.
|
||||||
func HeartbeatNode(ctx context.Context, authNode *model.OpenFlareNode, payload NodePayload) (*HeartbeatResponse, error) {
|
func HeartbeatNode(ctx context.Context, authNode *model.OpenFlareNode, payload NodePayload) (*HeartbeatResponse, error) {
|
||||||
_ = ofgeoip.EnsureRuntimeProvider(ctx)
|
_ = ofgeoip.EnsureRuntimeProvider(ctx)
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -22,16 +22,17 @@ import (
|
|||||||
// @Accept json
|
// @Accept json
|
||||||
// @Produce json
|
// @Produce json
|
||||||
// @Security AgentTokenAuth
|
// @Security AgentTokenAuth
|
||||||
// @Param body body agent.NodePayload true "节点上报数据"
|
// @Param body body protocol.RegisterNodeRequest true "节点上报数据;registration_key 为首次注册重试凭证"
|
||||||
// @Success 200 {object} response.Any{data=agent.RegistrationResponse} "注册成功"
|
// @Success 200 {object} response.Any{data=agent.RegistrationResponseDocs} "注册成功"
|
||||||
// @Failure 400 {object} response.Any "参数错误"
|
// @Failure 400 {object} response.Any "参数错误"
|
||||||
// @Failure 401 {object} response.Any "Token 无效"
|
// @Failure 401 {object} response.Any "Token 无效"
|
||||||
// @Router /api/v1/agent/nodes/register [post]
|
// @Router /api/v1/agent/nodes/register [post]
|
||||||
func RegisterHandler(c *gin.Context) {
|
func RegisterHandler(c *gin.Context) {
|
||||||
var payload NodePayload
|
var request protocol.RegisterNodeRequest
|
||||||
if !apiutil.BindJSON(c, &payload) {
|
if !apiutil.BindJSON(c, &request) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
payload := request.NodePayload
|
||||||
payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)
|
payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)
|
||||||
|
|
||||||
var (
|
var (
|
||||||
@@ -41,7 +42,7 @@ func RegisterHandler(c *gin.Context) {
|
|||||||
if authNode, ok := NodeFromContext(c); ok {
|
if authNode, ok := NodeFromContext(c); ok {
|
||||||
result, err = RegisterWithAccessToken(c.Request.Context(), authNode, payload)
|
result, err = RegisterWithAccessToken(c.Request.Context(), authNode, payload)
|
||||||
} else {
|
} else {
|
||||||
result, err = RegisterWithDiscovery(c.Request.Context(), payload)
|
result, err = RegisterWithDiscovery(c.Request.Context(), payload, request.RegistrationKey)
|
||||||
}
|
}
|
||||||
if apiutil.AbortBadRequestOnError(c, err) {
|
if apiutil.AbortBadRequestOnError(c, err) {
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/Rain-kl/Wavelet/internal/model"
|
"github.com/Rain-kl/Wavelet/internal/model"
|
||||||
|
"github.com/Rain-kl/Wavelet/pkg/protocol"
|
||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
@@ -17,9 +18,13 @@ const (
|
|||||||
)
|
)
|
||||||
|
|
||||||
// RegistrationResponse is returned after agent registration.
|
// RegistrationResponse is returned after agent registration.
|
||||||
// Server uses access_token; the agent client expects agent_token via RegisterNodeResponse.
|
type RegistrationResponse = protocol.RegisterNodeResponse
|
||||||
type RegistrationResponse struct {
|
|
||||||
|
// 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"`
|
NodeID string `json:"node_id"`
|
||||||
|
AgentToken string `json:"agent_token"`
|
||||||
AccessToken string `json:"access_token"`
|
AccessToken string `json:"access_token"`
|
||||||
Name string `json:"name"`
|
Name string `json:"name"`
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
+43
-1
@@ -4,7 +4,11 @@
|
|||||||
// Package protocol defines the communication protocol between OpenFlare server, agent, and relay components.
|
// Package protocol defines the communication protocol between OpenFlare server, agent, and relay components.
|
||||||
package protocol
|
package protocol
|
||||||
|
|
||||||
import "encoding/json"
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
// APIResponse is a generic API response wrapper.
|
// APIResponse is a generic API response wrapper.
|
||||||
type APIResponse[T any] struct {
|
type APIResponse[T any] struct {
|
||||||
@@ -166,6 +170,12 @@ type NodeHealthEvent struct {
|
|||||||
Metadata map[string]string `json:"metadata,omitempty"`
|
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.
|
// RegisterNodeResponse is the node registration response.
|
||||||
type RegisterNodeResponse struct {
|
type RegisterNodeResponse struct {
|
||||||
NodeID string `json:"node_id"`
|
NodeID string `json:"node_id"`
|
||||||
@@ -173,6 +183,38 @@ type RegisterNodeResponse struct {
|
|||||||
Name string `json:"name"`
|
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.
|
// ActiveConfigResponse is the active configuration response.
|
||||||
type ActiveConfigResponse struct {
|
type ActiveConfigResponse struct {
|
||||||
Version string `json:"version"`
|
Version string `json:"version"`
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user