refactor(repo): consolidate openflare-server to root and move subprojects to internal/apps

- Merge all files inside openflare-server to the repository root directory.
- Relocate agent, relay, and flared subprojects from internal/ to internal/apps/.
- Combine docker-compose files and update build context paths to root.
- Update GitHub workflows and Dockerfiles to refer to new directories and package names.
- Rewrite Go package imports across all files.
- Resolve database renew test race condition and clean up docs.
This commit is contained in:
ryan
2026-06-19 14:23:29 +08:00
parent 19d476ed7f
commit 63cd906cfc
1064 changed files with 366 additions and 1397 deletions
+245
View File
@@ -0,0 +1,245 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package auth_source 提供认证源管理功能
package auth_source
import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/Rain-kl/Wavelet/internal/apps/admin"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// AuthSourceRequest 创建或更新认证源的请求参数
type AuthSourceRequest struct {
Name string `json:"name"`
Type string `json:"type"`
DisplayName string `json:"display_name"`
IsActive bool `json:"is_active"`
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret"`
OpenIDDiscoveryURL string `json:"openid_discovery_url"`
Scopes string `json:"scopes"`
IconURL string `json:"icon_url"`
}
// ToggleAuthSourceRequest 切换认证源启用状态的请求参数
type ToggleAuthSourceRequest struct {
IsActive bool `json:"is_active"`
}
// ListAuthSources 获取认证源列表
// @Summary 获取认证源列表
// @Description 返回所有已配置的 OAuth/OIDC 认证源列表,包括已启用和未启用的,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]model.AuthSource} "认证源列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/auth-sources [get]
func ListAuthSources(c *gin.Context) {
sources, err := model.GetAuthSources(c.Request.Context())
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(sources))
}
// CreateAuthSource 创建认证源
// @Summary 创建认证源
// @Description 创建一个新的 OAuth/OIDC 认证源配置,认证源名称必须唯一且符合命名规范,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body auth_source.AuthSourceRequest true "创建认证源参数"
// @Success 200 {object} response.Any{data=model.AuthSource} "创建成功,返回认证源信息"
// @Failure 400 {object} response.Any "参数错误或验证失败"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/auth-sources [post]
func CreateAuthSource(c *gin.Context) {
var req AuthSourceRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
source := model.AuthSource{
Name: req.Name,
Type: req.Type,
DisplayName: req.DisplayName,
IsActive: req.IsActive,
ClientID: req.ClientID,
ClientSecret: req.ClientSecret,
OpenIDDiscoveryURL: req.OpenIDDiscoveryURL,
Scopes: req.Scopes,
IconURL: req.IconURL,
}
if err := model.CreateAuthSource(c.Request.Context(), &source); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
source.Sanitize()
c.JSON(http.StatusOK, response.OK(source))
}
// UpdateAuthSource 更新认证源
// @Summary 更新认证源
// @Description 更新指定 ID 的认证源配置。若 client_secret 字段为空,则保留原有密钥不变,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path uint64 true "认证源 ID 或名称"
// @Param request body auth_source.AuthSourceRequest true "更新认证源参数"
// @Success 200 {object} response.Any{data=model.AuthSource} "更新成功,返回更新后的认证源信息"
// @Failure 400 {object} response.Any "参数错误或验证失败"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/auth-sources/{id} [put]
func UpdateAuthSource(c *gin.Context) {
id, err := parseSourceID(c)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
var req AuthSourceRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
// 记录更新前的 Discovery URL,以便更新成功后清除旧缓存条目。
existing, _ := model.GetAuthSourceByID(c.Request.Context(), id)
source := model.AuthSource{
ID: id,
Name: req.Name,
Type: req.Type,
DisplayName: req.DisplayName,
IsActive: req.IsActive,
ClientID: req.ClientID,
ClientSecret: req.ClientSecret,
OpenIDDiscoveryURL: req.OpenIDDiscoveryURL,
Scopes: req.Scopes,
IconURL: req.IconURL,
}
keepSecret := source.ClientSecret == ""
if err := model.UpdateAuthSource(c.Request.Context(), &source, keepSecret); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
// Discovery URL 可能已变更,清除旧、新 issuer 的 provider 缓存,
// 确保下次登录时重新拉取最新 OIDC 元数据。
if existing != nil {
oauth.InvalidateOIDCProviderCache(normalizeIssuer(existing.OpenIDDiscoveryURL))
}
oauth.InvalidateOIDCProviderCache(normalizeIssuer(req.OpenIDDiscoveryURL))
updated, err := model.GetAuthSourceByID(c.Request.Context(), id)
if err != nil {
response.AbortInternal(c, err.Error())
return
}
updated.Sanitize()
c.JSON(http.StatusOK, response.OK(updated))
}
// ToggleAuthSource 切换认证源启用状态
// @Summary 切换认证源启用状态
// @Description 启用或禁用指定认证源。尝试启用时将验证 Client ID 和 Client Secret 是否已配置,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path uint64 true "认证源 ID 或名称"
// @Param request body auth_source.ToggleAuthSourceRequest true "启用状态"
// @Success 200 {object} response.Any{data=string} "切换成功"
// @Failure 400 {object} response.Any "验证失败或认证源不存在"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/auth-sources/{id}/toggle [put]
func ToggleAuthSource(c *gin.Context) {
id, err := parseSourceID(c)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
var req ToggleAuthSourceRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if err := model.ToggleAuthSource(c.Request.Context(), id, req.IsActive); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// DeleteAuthSource 删除认证源
// @Summary 删除认证源
// @Description 删除指定认证源及其关联的所有外部帐号绑定记录,警告:删除后相关用户将无法通过该源登录,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path uint64 true "认证源 ID 或名称"
// @Success 200 {object} response.Any{data=string} "删除成功"
// @Failure 400 {object} response.Any "ID 无效或删除失败"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/auth-sources/{id} [delete]
func DeleteAuthSource(c *gin.Context) {
id, err := parseSourceID(c)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if err := model.DeleteAuthSource(c.Request.Context(), id); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
func parseSourceID(c *gin.Context) (uint64, error) {
raw := c.Param("id")
if raw == "" {
return 0, errors.New(admin.InvalidAuthSourceID)
}
source, err := model.GetAuthSourceByName(c.Request.Context(), raw)
if err == nil {
return source.ID, nil
}
var id uint64
if _, scanErr := fmt.Sscanf(raw, "%d", &id); scanErr != nil || id == 0 {
return 0, errors.New(admin.InvalidAuthSourceID)
}
return id, nil
}
// normalizeIssuer 将 Discovery URL 规范化为 issuer 基础 URL,
// 与 oauth.buildOAuthConfig 中的规范化逻辑保持一致。
func normalizeIssuer(discoveryURL string) string {
issuer := strings.TrimSuffix(strings.TrimSpace(discoveryURL), "/")
issuer = strings.TrimSuffix(issuer, "/.well-known/openid-configuration")
issuer = strings.TrimSuffix(issuer, "/.well-known/oauth-authorization-server")
return issuer
}
@@ -0,0 +1,333 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package auth_source
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
// Mock authentication middleware
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, oauth.UserObjKey, authUser)
}
c.Next()
})
adminGroup.GET("/auth-sources", ListAuthSources)
adminGroup.POST("/auth-sources", CreateAuthSource)
adminGroup.PUT("/auth-sources/:id", UpdateAuthSource)
adminGroup.PUT("/auth-sources/:id/toggle", ToggleAuthSource)
adminGroup.DELETE("/auth-sources/:id", DeleteAuthSource)
return r
}
func TestListAuthSources(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
// Seed source
source := model.AuthSource{
ID: 1,
Name: "google",
Type: "oidc",
DisplayName: "Google Auth",
IsActive: true,
ClientID: "client_id_123",
ClientSecret: "client_secret_456",
OpenIDDiscoveryURL: "https://accounts.google.com",
}
dbConn.Create(&source)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
req, _ := http.NewRequest("GET", "/api/v1/admin/auth-sources", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d", w.Code)
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var sources []model.AuthSource
_ = json.Unmarshal(dataBytes, &sources)
if len(sources) != 1 {
t.Errorf("expected 1 auth source, got %d", len(sources))
}
if sources[0].Name != "google" {
t.Errorf("expected name 'google', got '%s'", sources[0].Name)
}
// Verify sanitize removed the secret
if sources[0].ClientSecret != "" {
t.Error("client secret should be sanitized")
}
if !sources[0].ClientSecretConfigured {
t.Error("client secret configured flag should be true")
}
}
func TestCreateAuthSource(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("create successfully", func(t *testing.T) {
reqPayload := AuthSourceRequest{
Name: "github",
Type: "oidc",
DisplayName: "GitHub OIDC",
IsActive: true,
ClientID: "client_id_gh",
ClientSecret: "client_secret_gh",
OpenIDDiscoveryURL: "https://github.com",
}
body, _ := json.Marshal(reqPayload)
req, _ := http.NewRequest("POST", "/api/v1/admin/auth-sources", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
// Verify database
var src model.AuthSource
dbConn.Where("name = ?", "github").First(&src)
if src.ClientID != "client_id_gh" {
t.Errorf("expected client_id_gh, got '%s'", src.ClientID)
}
})
t.Run("create invalid validation failure", func(t *testing.T) {
reqPayload := AuthSourceRequest{
Name: "invalid name!",
Type: "oidc",
DisplayName: "Invalid",
IsActive: true,
ClientID: "client_id_val",
ClientSecret: "client_secret_val",
OpenIDDiscoveryURL: "https://discovery.url",
}
body, _ := json.Marshal(reqPayload)
req, _ := http.NewRequest("POST", "/api/v1/admin/auth-sources", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d", w.Code)
}
})
}
func TestUpdateAuthSource(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
// Seed source
source := model.AuthSource{
ID: 1,
Name: "microsoft",
Type: "oidc",
DisplayName: "Microsoft",
IsActive: true,
ClientID: "old_client_id",
ClientSecret: "old_secret",
OpenIDDiscoveryURL: "https://login.microsoftonline.com",
}
dbConn.Create(&source)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("update keep client secret", func(t *testing.T) {
reqPayload := AuthSourceRequest{
Name: "microsoft",
Type: "oidc",
DisplayName: "Microsoft Updated",
IsActive: true,
ClientID: "new_client_id",
ClientSecret: "", // empty implies keeping existing secret
OpenIDDiscoveryURL: "https://login.microsoftonline.com",
}
body, _ := json.Marshal(reqPayload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/auth-sources/1", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var src model.AuthSource
dbConn.First(&src, 1)
if src.DisplayName != "Microsoft Updated" {
t.Errorf("expected display name update, got '%s'", src.DisplayName)
}
if src.ClientSecret != "old_secret" {
t.Errorf("expected old secret to be preserved, got '%s'", src.ClientSecret)
}
})
t.Run("update new client secret", func(t *testing.T) {
reqPayload := AuthSourceRequest{
Name: "microsoft",
Type: "oidc",
DisplayName: "Microsoft Updated Again",
IsActive: true,
ClientID: "new_client_id",
ClientSecret: "brand_new_secret",
OpenIDDiscoveryURL: "https://login.microsoftonline.com",
}
body, _ := json.Marshal(reqPayload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/auth-sources/microsoft", bytes.NewBuffer(body)) // Using Name instead of ID
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var src model.AuthSource
dbConn.First(&src, 1)
if src.ClientSecret != "brand_new_secret" {
t.Errorf("expected secret update, got '%s'", src.ClientSecret)
}
})
}
func TestToggleAuthSource(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
source := model.AuthSource{
ID: 1,
Name: "test_source",
Type: "oidc",
DisplayName: "Test Source",
IsActive: false,
ClientID: "",
ClientSecret: "",
OpenIDDiscoveryURL: "https://test.discovery.url",
}
dbConn.Create(&source)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("cannot activate without credentials", func(t *testing.T) {
payload := ToggleAuthSourceRequest{IsActive: true}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/auth-sources/1/toggle", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request when activating without client_id/secret, got %d", w.Code)
}
})
t.Run("toggle success after setting credentials", func(t *testing.T) {
// Set credentials first
dbConn.Model(&model.AuthSource{}).Where("id = ?", 1).Updates(map[string]interface{}{
"client_id": "id",
"client_secret": "secret",
})
payload := ToggleAuthSourceRequest{IsActive: true}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/auth-sources/1/toggle", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var src model.AuthSource
dbConn.First(&src, 1)
if !src.IsActive {
t.Error("auth source should be activated")
}
})
}
func TestDeleteAuthSource(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
source := model.AuthSource{
ID: 1,
Name: "delete_me",
Type: "oidc",
DisplayName: "Delete Me",
IsActive: true,
ClientID: "id",
ClientSecret: "secret",
OpenIDDiscoveryURL: "https://delete.me",
}
dbConn.Create(&source)
externalAccount := model.ExternalAccount{
ID: 10,
AuthSourceID: 1,
UserID: 50,
ExternalID: "ext_50",
}
dbConn.Create(&externalAccount)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
req, _ := http.NewRequest("DELETE", "/api/v1/admin/auth-sources/1", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d", w.Code)
}
// Verify AuthSource is deleted
var srcCount int64
dbConn.Model(&model.AuthSource{}).Where("id = ?", 1).Count(&srcCount)
if srcCount != 0 {
t.Error("AuthSource should be deleted from the database")
}
// Verify ExternalAccount bindings are also deleted
var extCount int64
dbConn.Model(&model.ExternalAccount{}).Where("auth_source_id = ?", 1).Count(&extCount)
if extCount != 0 {
t.Error("related ExternalAccount bindings should be deleted")
}
}
+14
View File
@@ -0,0 +1,14 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package cache
import (
"context"
"github.com/Rain-kl/Wavelet/internal/repository"
)
func saveOrUpdateConfig(ctx context.Context, key, value string) error {
return repository.SaveOrUpdateSystemConfig(ctx, key, value)
}
+100
View File
@@ -0,0 +1,100 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package cache provides HTTP handlers for managing disk cache.
package cache
import (
"net/http"
"strconv"
"github.com/Rain-kl/Wavelet/internal/diskcache"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
type updateCacheConfigRequest struct {
MaxSizeMB int64 `json:"max_size_mb" binding:"required,min=1"`
TTLMinutes int64 `json:"ttl_minutes" binding:"required,min=0"`
LRUEnabled bool `json:"lru_enabled"`
}
// GetCacheStatus 获取磁盘缓存状态与当前统计数据
// @Summary 获取缓存状态
// @Description 获取当前系统磁盘缓存的使用情况(已占用字节、Key 数量等)与策略配置
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=diskcache.Status} "获取成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/cache/status [get]
func GetCacheStatus(c *gin.Context) {
status := diskcache.GetGlobalCache().Status()
c.JSON(http.StatusOK, response.OK(status))
}
// UpdateCacheConfig 更新磁盘缓存策略配置
// @Summary 更新缓存配置
// @Description 更改磁盘缓存最大容量限制、文件生存时间(TTL)以及是否启用 LRU 淘汰淘汰算法,并进行热更新
// @Tags admin
// @Accept json
// @Produce json
// @Param request body cache.updateCacheConfigRequest true "缓存配置请求体"
// @Security SessionCookie
// @Success 200 {object} response.Any "更新成功"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "服务内部错误"
// @Router /api/v1/admin/cache/config [post]
func UpdateCacheConfig(c *gin.Context) {
var req updateCacheConfigRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
ctx := c.Request.Context()
if err := saveOrUpdateConfig(ctx, model.ConfigKeyDiskCacheMaxSizeMB, strconv.FormatInt(req.MaxSizeMB, 10)); err != nil {
response.AbortInternal(c, err.Error())
return
}
if err := saveOrUpdateConfig(ctx, model.ConfigKeyDiskCacheTTLMinutes, strconv.FormatInt(req.TTLMinutes, 10)); err != nil {
response.AbortInternal(c, err.Error())
return
}
if err := saveOrUpdateConfig(ctx, model.ConfigKeyDiskCacheLRUEnabled, strconv.FormatBool(req.LRUEnabled)); err != nil {
response.AbortInternal(c, err.Error())
return
}
diskcache.GetGlobalCache().ReloadConfig(ctx)
c.JSON(http.StatusOK, response.OKNil())
}
// ClearCache 一键清空所有磁盘缓存数据
// @Summary 清空缓存
// @Description 清除系统磁盘缓存目录中的所有临时文件,并重置缓存容量和 Key 追踪数据
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any "清理成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "服务内部错误"
// @Router /api/v1/admin/cache/clear [post]
func ClearCache(c *gin.Context) {
if err := diskcache.GetGlobalCache().Clear(); err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
+498
View File
@@ -0,0 +1,498 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package db_manage provides router handlers for managing database tables,
// overview information, and executing custom SQL queries.
package db_manage
import (
"database/sql"
"fmt"
"math"
"net/http"
"os"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/config"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
const (
binaryKB = 0
binaryMB = 1
binaryGB = 2
valueThreshold = 10
maxStringLength = 200
)
// DBOverviewResponse 数据库运行概览响应结构体
type DBOverviewResponse struct {
Type string `json:"type"`
Version string `json:"version"`
Name string `json:"name"`
Size string `json:"size"`
TableCount int64 `json:"table_count"`
Connections int64 `json:"connections"`
}
// GetTableDataRequest 分页拉取表数据请求结构体
type GetTableDataRequest struct {
Table string `form:"table" binding:"required"`
Page int `form:"page,default=1"`
PageSize int `form:"pageSize,default=10"`
}
// TableDataResponse 动态数据表响应结构体
type TableDataResponse struct {
Columns []string `json:"columns"`
Total int64 `json:"total"`
Results []map[string]interface{} `json:"results"`
}
// ExecuteSQLRequest 执行自定义 SQL 请求结构体
type ExecuteSQLRequest struct {
SQL string `json:"sql" binding:"required"`
}
// ExecuteSQLResponse 执行自定义 SQL 响应结构体
type ExecuteSQLResponse struct {
Type string `json:"type"` // "select" 或 "exec"
Columns []string `json:"columns,omitempty"`
Results []map[string]interface{} `json:"results,omitempty"`
AffectedRows int64 `json:"affected_rows"`
ExecutionTimeMs int64 `json:"execution_time_ms"`
}
// formatBytes 格式化字节大小为可读字符串
func formatBytes(bytes uint64) string {
const unit = 1024
if bytes < unit {
return fmt.Sprintf("%d B", bytes)
}
div, exp := int64(unit), 0
for n := bytes / unit; n >= unit; n /= unit {
div *= unit
exp++
}
value := float64(bytes) / float64(div)
var suffix string
switch exp {
case binaryKB:
suffix = "KiB"
case binaryMB:
suffix = "MiB"
case binaryGB:
suffix = "GiB"
default:
suffix = "TiB"
}
if value == math.Trunc(value) {
if value >= valueThreshold {
return fmt.Sprintf("%.0f %s", value, suffix)
}
return fmt.Sprintf("%.1f %s", value, suffix)
}
return fmt.Sprintf("%.1f %s", value, suffix)
}
// getSQLiteOverview 获取 SQLite 数据库概览信息
func getSQLiteOverview(gormDB *gorm.DB) (DBOverviewResponse, error) {
name := config.Config.Database.SQLitePath
if name == "" {
name = "./data/openflare.db"
}
var version string
var ver string
if err := gormDB.Raw("SELECT sqlite_version()").Scan(&ver).Error; err == nil {
version = "SQLite " + ver
} else {
version = "SQLite"
}
var sizeStr string
if fi, err := os.Stat(name); err == nil {
size := fi.Size()
if size < 0 {
size = 0
}
sizeStr = formatBytes(uint64(size))
} else {
sizeStr = "0 B"
}
var tableCount int64
if err := gormDB.Raw("SELECT count(*) FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'").Scan(&tableCount).Error; err != nil {
tableCount = 0
}
var connCount int64
if sqlDB, err := gormDB.DB(); err == nil {
connCount = int64(sqlDB.Stats().OpenConnections)
} else {
connCount = 1
}
return DBOverviewResponse{
Type: "sqlite",
Version: version,
Name: name,
Size: sizeStr,
TableCount: tableCount,
Connections: connCount,
}, nil
}
// getPostgresOverview 获取 PostgreSQL 数据库概览信息
func getPostgresOverview(gormDB *gorm.DB) (DBOverviewResponse, error) {
name := config.Config.Database.Database
var version string
var ver string
if err := gormDB.Raw("SELECT version()").Scan(&ver).Error; err == nil {
version = ver
} else {
version = "PostgreSQL"
}
var sizeStr string
var sizeBytes sql.NullInt64
if err := gormDB.Raw("SELECT pg_database_size(current_database())").Scan(&sizeBytes).Error; err == nil && sizeBytes.Valid {
size := sizeBytes.Int64
if size < 0 {
size = 0
}
sizeStr = formatBytes(uint64(size))
} else {
sizeStr = "0 B"
}
var tableCount int64
if err := gormDB.Raw("SELECT count(*) FROM information_schema.tables WHERE table_schema = current_schema()").Scan(&tableCount).Error; err != nil {
tableCount = 0
}
var connCount int64
var pgc sql.NullInt64
if err := gormDB.Raw("SELECT count(*) FROM pg_stat_activity WHERE datname = current_database()").Scan(&pgc).Error; err == nil && pgc.Valid {
connCount = pgc.Int64
} else {
if sqlDB, err := gormDB.DB(); err == nil {
connCount = int64(sqlDB.Stats().OpenConnections)
} else {
connCount = 1
}
}
return DBOverviewResponse{
Type: "postgres",
Version: version,
Name: name,
Size: sizeStr,
TableCount: tableCount,
Connections: connCount,
}, nil
}
// GetDBOverview 获取数据库运行概览
// @Summary 获取数据库运行概览
// @Description 获取数据库类型、版本、名称、文件大小、表数量及当前连接数,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=db_manage.DBOverviewResponse} "获取成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/db-manage/overview [get]
func GetDBOverview(c *gin.Context) {
gormDB := db.DB(c.Request.Context())
if gormDB == nil {
response.AbortInternal(c, "数据库未初始化")
return
}
var overview DBOverviewResponse
var err error
if !config.Config.Database.Enabled {
overview, err = getSQLiteOverview(gormDB)
} else {
overview, err = getPostgresOverview(gormDB)
}
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(overview))
}
// ListDBTables 获取数据库所有表名
// @Summary 获取数据库所有表名
// @Description 返回当前数据库的所有用户自定义表名称列表,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]string} "获取成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/db-manage/tables [get]
func ListDBTables(c *gin.Context) {
gormDB := db.DB(c.Request.Context())
if gormDB == nil {
response.AbortInternal(c, "数据库未初始化")
return
}
var tables []string
var err error
if !config.Config.Database.Enabled {
err = gormDB.Raw("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name").Scan(&tables).Error
} else {
err = gormDB.Raw("SELECT table_name FROM information_schema.tables WHERE table_schema = current_schema() ORDER BY table_name").Scan(&tables).Error
}
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(tables))
}
// GetDBTableData 获取数据表 data
func GetDBTableData(c *gin.Context) {
var req GetTableDataRequest
if err := c.ShouldBindQuery(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
gormDB := db.DB(c.Request.Context())
if gormDB == nil {
response.AbortInternal(c, "数据库未初始化")
return
}
// 安全转义表名并拼接
quotedTable := `"` + strings.ReplaceAll(req.Table, `"`, `""`) + `"`
var total int64
if err := gormDB.Raw("SELECT count(*) FROM " + quotedTable).Scan(&total).Error; err != nil {
response.AbortBadRequest(c, err.Error())
return
}
offset := (req.Page - 1) * req.PageSize
if offset < 0 {
offset = 0
}
limit := req.PageSize
if limit <= 0 {
limit = 10
}
rows, err := gormDB.Raw("SELECT * FROM "+quotedTable+" LIMIT ? OFFSET ?", limit, offset).Rows()
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
defer func() {
_ = rows.Close()
}()
cols, err := rows.Columns()
if err != nil {
response.AbortInternal(c, err.Error())
return
}
results, err := scanTableRows(rows, cols)
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(TableDataResponse{
Columns: cols,
Total: total,
Results: results,
}))
}
// scanTableRows 扫描并提取数据表行数据,做截断处理
func scanTableRows(rows *sql.Rows, cols []string) ([]map[string]interface{}, error) {
results := make([]map[string]interface{}, 0)
for rows.Next() {
columns := make([]interface{}, len(cols))
columnPointers := make([]interface{}, len(cols))
for i := range columns {
columnPointers[i] = &columns[i]
}
if err := rows.Scan(columnPointers...); err != nil {
return nil, err
}
rowMap := make(map[string]interface{})
for i, colName := range cols {
val := columns[i]
if b, ok := val.([]byte); ok {
strVal := string(b)
runes := []rune(strVal)
if len(runes) > maxStringLength {
strVal = string(runes[:maxStringLength]) + "..."
}
rowMap[colName] = strVal
} else if str, ok := val.(string); ok {
runes := []rune(str)
if len(runes) > maxStringLength {
str = string(runes[:maxStringLength]) + "..."
}
rowMap[colName] = str
} else {
rowMap[colName] = val
}
}
results = append(results, rowMap)
}
return results, nil
}
// executeSQLQuery 执行并解析查询类 SQL 语句
func executeSQLQuery(gormDB *gorm.DB, sqlStr string, startTime time.Time) (ExecuteSQLResponse, error) {
rows, err := gormDB.Raw(sqlStr).Rows()
if err != nil {
return ExecuteSQLResponse{}, err
}
defer func() {
_ = rows.Close()
}()
cols, err := rows.Columns()
if err != nil {
return ExecuteSQLResponse{}, err
}
results := make([]map[string]interface{}, 0)
for rows.Next() {
columns := make([]interface{}, len(cols))
columnPointers := make([]interface{}, len(cols))
for i := range columns {
columnPointers[i] = &columns[i]
}
if err := rows.Scan(columnPointers...); err != nil {
return ExecuteSQLResponse{}, err
}
rowMap := make(map[string]interface{})
for i, colName := range cols {
val := columns[i]
if b, ok := val.([]byte); ok {
rowMap[colName] = string(b)
} else {
rowMap[colName] = val
}
}
results = append(results, rowMap)
}
executionTime := time.Since(startTime).Milliseconds()
return ExecuteSQLResponse{
Type: "select",
Columns: cols,
Results: results,
AffectedRows: int64(len(results)),
ExecutionTimeMs: executionTime,
}, nil
}
// executeSQLMutation 执行修改/更新类 SQL 语句
func executeSQLMutation(gormDB *gorm.DB, sqlStr string, startTime time.Time) (ExecuteSQLResponse, error) {
tx := gormDB.Exec(sqlStr)
if tx.Error != nil {
return ExecuteSQLResponse{}, tx.Error
}
executionTime := time.Since(startTime).Milliseconds()
return ExecuteSQLResponse{
Type: "exec",
AffectedRows: tx.RowsAffected,
ExecutionTimeMs: executionTime,
}, nil
}
// ExecuteSQL 执行 SQL 查询
// @Summary 执行 SQL 查询
// @Description 在当前数据库中执行任意自定义 SQL,如果是查询语句将返回格式化后的列与数据集,否则返回受影响行数,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body db_manage.ExecuteSQLRequest true "SQL 请求参数"
// @Success 200 {object} response.Any{data=db_manage.ExecuteSQLResponse} "执行完毕"
// @Failure 400 {object} response.Any "SQL 语句错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/db-manage/query [post]
func ExecuteSQL(c *gin.Context) {
var req ExecuteSQLRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
gormDB := db.DB(c.Request.Context())
if gormDB == nil {
response.AbortInternal(c, "数据库未初始化")
return
}
trimmedSQL := strings.TrimSpace(req.SQL)
if trimmedSQL == "" {
response.AbortBadRequest(c, "SQL 语句不能为空")
return
}
startTime := time.Now()
// 识别是否是查询语句(SELECT, SHOW, EXPLAIN 等)
isQuery := false
lowerSQL := strings.ToLower(trimmedSQL)
queryKeywords := []string{"select", "show", "explain", "describe", "pragma"}
for _, kw := range queryKeywords {
if strings.HasPrefix(lowerSQL, kw) {
isQuery = true
break
}
}
var resp ExecuteSQLResponse
var err error
if isQuery {
resp, err = executeSQLQuery(gormDB, trimmedSQL, startTime)
} else {
resp, err = executeSQLMutation(gormDB, trimmedSQL, startTime)
}
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(resp))
}
+15
View File
@@ -0,0 +1,15 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package admin 提供管理后台功能
package admin
// 管理后台错误消息常量
const (
AdminRequired = "未经授权访问"
TokenAdminRequired = "该访问令牌没有管理员权限,无法访问管理端点" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
InvalidAuthSourceID = "认证源 ID 无效"
InvalidCursorParam = "无效的 cursor 参数"
InvalidTaskExecutionID = "无效的任务执行记录 ID"
)
+430
View File
@@ -0,0 +1,430 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package logs 提供日志查询与分析功能
package logs
import (
"context"
"encoding/json"
"fmt"
"net/http"
"strconv"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/admin"
"github.com/Rain-kl/Wavelet/internal/config"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/model"
analyticsrepo "github.com/Rain-kl/Wavelet/internal/repository/analytics"
"github.com/Rain-kl/Wavelet/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
const (
defaultLimit = 200
maxLimit = 500
maxPageSize = 100
hoursInDay = 24
analyticsDays = 7
topActiveLimit = 10
)
// logsResponse 历史日志查询响应
type logsResponse struct {
Lines []logger.LogEntry `json:"lines"`
HasMore bool `json:"has_more"`
NextCursor int `json:"next_cursor"` // 用于加载更早日志的 cursor
}
// GetLogs 获取历史日志
// @Summary 获取系统日志
// @Description 分页获取系统历史日志,cursor=0 获取最新日志,cursor>0 获取更早日志
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param cursor query int false "日志游标,0=获取最新" default(0)
// @Param limit query int false "每页条数" default(200)
// @Success 200 {object} response.Any{data=logs.logsResponse} "日志列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/logs [get]
func GetLogs(c *gin.Context) {
cursorStr := c.DefaultQuery("cursor", "0")
limitStr := c.DefaultQuery("limit", "200")
var cursor, limit int
if _, err := parsePositiveInt(cursorStr, &cursor); err != nil {
response.AbortWithError(c, http.StatusBadRequest, admin.InvalidCursorParam)
return
}
if _, err := parsePositiveInt(limitStr, &limit); err != nil || limit <= 0 {
limit = defaultLimit
}
if limit > maxLimit {
limit = maxLimit
}
entries, hasMore := logger.GlobalRingBuffer.Query(cursor, limit)
resp := logsResponse{
Lines: entries,
HasMore: hasMore,
}
if len(entries) > 0 {
resp.NextCursor = entries[0].Index
}
c.JSON(http.StatusOK, response.OK(resp))
}
// wsMessage WebSocket 消息格式
type wsMessage struct {
Type string `json:"type"` // "log" | "error"
Data json.RawMessage `json:"data"`
}
// HandleLogWebSocket WebSocket 端点,实时推送系统日志
// @Summary 系统日志实时推送
// @Description 通过 WebSocket 实时推送系统日志,需要管理员权限
// @Tags admin
// @Router /api/v1/admin/logs/ws [get]
func HandleLogWebSocket(c *gin.Context) {
upgrader := getUpgrader()
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
return
}
defer func() { _ = conn.Close() }()
// 订阅 ring buffer
ch := logger.GlobalRingBuffer.Subscribe()
defer logger.GlobalRingBuffer.Unsubscribe(ch)
// 在独立 goroutine 中读取客户端消息(保持连接活跃 + 检测断开)
done := make(chan struct{})
go func() {
defer close(done)
for {
_, _, err := conn.ReadMessage()
if err != nil {
return
}
}
}()
// 主循环:推送日志
for {
select {
case <-done:
return
case entry, ok := <-ch:
if !ok {
return
}
data, _ := json.Marshal(entry)
msg := wsMessage{Type: "log", Data: data}
payload, _ := json.Marshal(msg)
if err := conn.WriteMessage(1, payload); err != nil {
return
}
}
}
}
// accessLogItem 访问日志单条数据
type accessLogItem struct {
ID uint64 `json:"id,string"`
UserID uint64 `json:"user_id,string"`
Username string `json:"username"`
Nickname string `json:"nickname"`
Path string `json:"path"`
Method string `json:"method"`
IP string `json:"ip"`
UserAgent string `json:"user_agent"`
Headers string `json:"headers"`
Status int32 `json:"status"`
Latency int64 `json:"latency"`
CreatedAt string `json:"created_at"`
}
// accessLogsResponse 访问日志查询响应
type accessLogsResponse struct {
Total uint64 `json:"total"`
List []accessLogItem `json:"list"`
}
func buildAccessLogFilter(ctx context.Context, c *gin.Context) (analyticsrepo.AccessLogFilter, error) {
filter := analyticsrepo.AccessLogFilter{}
username := c.Query("username")
if username != "" {
var userIDs []uint64
err := db.DB(ctx).Model(&model.User{}).
Where("username LIKE ?", "%"+username+"%").
Pluck("id", &userIDs).Error
if err != nil {
return filter, fmt.Errorf("查询用户信息失败: %w", err)
}
filter.UserIDs = userIDs
}
if path := c.Query("path"); path != "" {
filter.Path = path
}
if startTime := c.Query("start_time"); startTime != "" {
if t, err := parseAccessLogTime(startTime); err == nil {
filter.StartTime = &t
}
}
if endTime := c.Query("end_time"); endTime != "" {
if t, err := parseAccessLogTime(endTime); err == nil {
filter.EndTime = &t
}
}
return filter, nil
}
func parseAccessLogTime(value string) (time.Time, error) {
if t, err := time.Parse(time.RFC3339, value); err == nil {
return t, nil
}
return time.Parse("2006-01-02 15:04:05", value)
}
func enrichAccessLogsWithUsers(ctx context.Context, list []accessLogItem) {
if len(list) == 0 {
return
}
userIDs := make([]uint64, 0, len(list))
seen := make(map[uint64]struct{}, len(list))
for _, item := range list {
if _, ok := seen[item.UserID]; ok {
continue
}
seen[item.UserID] = struct{}{}
userIDs = append(userIDs, item.UserID)
}
userMap := make(map[uint64]struct{ Username, Nickname string })
var users []model.User
if err := db.DB(ctx).Where("id IN ?", userIDs).Find(&users).Error; err == nil {
for _, u := range users {
userMap[u.ID] = struct{ Username, Nickname string }{Username: u.Username, Nickname: u.Nickname}
}
}
for i := range list {
if info, ok := userMap[list[i].UserID]; ok {
list[i].Username = info.Username
list[i].Nickname = info.Nickname
}
}
}
// GetAccessLogs 获取 ClickHouse 异步采集的访问日志
// @Summary 获取用户访问日志
// @Description 分页并按照用户、接口路径、时间范围等维度检索 ClickHouse 用户访问日志列表(需要管理员权限,ClickHouse 未启用时报错)
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param page query int false "页码" default(1)
// @Param page_size query int false "每页条数" default(20)
// @Param username query string false "用户名模糊搜索"
// @Param path query string false "接口路径模糊搜索"
// @Param start_time query string false "起始时间(RFC3339 或 YYYY-MM-DD HH:MM:SS)"
// @Param end_time query string false "结束时间(RFC3339 或 YYYY-MM-DD HH:MM:SS)"
// @Success 200 {object} response.Any{data=logs.accessLogsResponse} "访问日志列表"
// @Failure 400 {object} response.Any "ClickHouse 未启用或参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/logs/access [get]
func GetAccessLogs(c *gin.Context) {
ctx := c.Request.Context()
if !config.Config.ClickHouse.Enabled || db.ChDB(ctx) == nil {
response.AbortWithError(c, http.StatusBadRequest, "ClickHouse 存储服务未启用,无法检索访问日志")
return
}
page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
if page < 1 {
page = 1
}
pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
if pageSize < 1 {
pageSize = 20
}
if pageSize > maxPageSize {
pageSize = maxPageSize
}
filter, err := buildAccessLogFilter(ctx, c)
if err != nil {
response.AbortWithError(c, http.StatusInternalServerError, err.Error())
return
}
if filter.UserIDs != nil && len(filter.UserIDs) == 0 {
c.JSON(http.StatusOK, response.OK(accessLogsResponse{Total: 0, List: []accessLogItem{}}))
return
}
logs, total, err := analyticsrepo.ListAccessLogs(ctx, filter, page, pageSize)
if err != nil {
response.AbortWithError(c, http.StatusInternalServerError, err.Error())
return
}
if total == 0 {
c.JSON(http.StatusOK, response.OK(accessLogsResponse{Total: 0, List: []accessLogItem{}}))
return
}
list := make([]accessLogItem, len(logs))
for i, logItem := range logs {
list[i] = accessLogItem{
ID: logItem.ID,
UserID: logItem.UserID,
Path: logItem.Path,
Method: logItem.Method,
IP: logItem.IP,
UserAgent: logItem.UserAgent,
Headers: logItem.Headers,
Status: logItem.Status,
Latency: logItem.Latency,
CreatedAt: logItem.CreatedAt.Format(time.RFC3339),
}
}
enrichAccessLogsWithUsers(ctx, list)
c.JSON(http.StatusOK, response.OK(accessLogsResponse{
Total: total,
List: list,
}))
}
// trendItem 趋势图数据点
type trendItem struct {
Date string `json:"date"`
Count uint64 `json:"count"`
}
// browserItem 浏览器占比排行
type browserItem struct {
Browser string `json:"browser"`
Count uint64 `json:"count"`
}
// topUserItem 活跃用户数据
type topUserItem struct {
UserID uint64 `json:"user_id,string"`
Username string `json:"username"`
Nickname string `json:"nickname"`
Count uint64 `json:"count"`
}
// logsAnalyticsResponse 访问日志数据分析结果
type logsAnalyticsResponse struct {
Trend []trendItem `json:"trend"`
Browsers []browserItem `json:"browsers"`
TopUsers []topUserItem `json:"top_users"`
}
// GetLogsAnalytics 获取 ClickHouse 访问日志图表聚合指标
// @Summary 获取访问日志分析数据
// @Description 聚合统计最近 7 天的每日访问趋势、浏览器分布以及前 10 名最活跃用户排行(需要管理员权限,ClickHouse 未启用时报错)
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=logs.logsAnalyticsResponse} "分析统计数据"
// @Failure 400 {object} response.Any "ClickHouse 未启用"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/logs/analytics [get]
func GetLogsAnalytics(c *gin.Context) {
ctx := c.Request.Context()
if !config.Config.ClickHouse.Enabled || db.ChDB(ctx) == nil {
response.AbortWithError(c, http.StatusBadRequest, "ClickHouse 存储服务未启用,无法获取分析数据")
return
}
startTime := time.Now().AddDate(0, 0, -(analyticsDays - 1)).Truncate(hoursInDay * time.Hour)
trendPoints, err := analyticsrepo.GetDailyTrend(ctx, analyticsDays)
if err != nil {
response.AbortWithError(c, http.StatusInternalServerError, "查询访问趋势失败: "+err.Error())
return
}
trendList := make([]trendItem, len(trendPoints))
for i, point := range trendPoints {
trendList[i] = trendItem{
Date: point.Date,
Count: point.Count,
}
}
browserPoints, err := analyticsrepo.GetBrowserDistribution(ctx, startTime)
if err != nil {
response.AbortWithError(c, http.StatusInternalServerError, "查询浏览器分布失败: "+err.Error())
return
}
browserList := make([]browserItem, len(browserPoints))
for i, point := range browserPoints {
browserList[i] = browserItem{
Browser: point.Browser,
Count: point.Count,
}
}
topUserPoints, err := analyticsrepo.GetTopActiveUsers(ctx, startTime, topActiveLimit)
if err != nil {
response.AbortWithError(c, http.StatusInternalServerError, "查询活跃用户失败: "+err.Error())
return
}
topUsers := make([]topUserItem, len(topUserPoints))
userIDs := make([]uint64, len(topUserPoints))
for i, point := range topUserPoints {
topUsers[i] = topUserItem{
UserID: point.UserID,
Count: point.Count,
}
userIDs[i] = point.UserID
}
if len(userIDs) > 0 {
userProfileMap := make(map[uint64]struct {
Username string
Nickname string
})
var users []model.User
if errProfile := db.DB(ctx).Where("id IN ?", userIDs).Find(&users).Error; errProfile == nil {
for _, u := range users {
userProfileMap[u.ID] = struct {
Username string
Nickname string
}{
Username: u.Username,
Nickname: u.Nickname,
}
}
}
for i := range topUsers {
if profile, ok := userProfileMap[topUsers[i].UserID]; ok {
topUsers[i].Username = profile.Username
topUsers[i].Nickname = profile.Nickname
}
}
}
c.JSON(http.StatusOK, response.OK(logsAnalyticsResponse{
Trend: trendList,
Browsers: browserList,
TopUsers: topUsers,
}))
}
+63
View File
@@ -0,0 +1,63 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package logs
import (
"net/http"
"net/url"
"strconv"
"strings"
"github.com/gorilla/websocket"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
)
// getUpgrader 返回 WebSocket 升级器并执行 Origin 安全检查以防止 CSWSH 攻击
func getUpgrader() *websocket.Upgrader {
return &websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
origin := r.Header.Get("Origin")
if origin == "" {
return true
}
// 1. 同源检查 (Same-origin check)
u, err := url.Parse(origin)
if err == nil && strings.EqualFold(u.Host, r.Host) {
return true
}
// 2. 检查配置的允许跨域 Origin (Check allowed origins in system config)
ctx := r.Context()
if sc, err := repository.GetSystemConfigByKey(ctx, model.ConfigKeyServerAddress); err == nil && sc.Value != "" {
originToCheck := strings.TrimRight(strings.TrimSpace(origin), "/")
allowedOrigins := strings.Split(sc.Value, ",")
for _, allowed := range allowedOrigins {
allowed = strings.TrimRight(strings.TrimSpace(allowed), "/")
if allowed != "" && strings.EqualFold(allowed, originToCheck) {
return true
}
}
}
return false
},
}
}
// parsePositiveInt 解析非负整数字符串
func parsePositiveInt(s string, result *int) (bool, error) {
if s == "" {
*result = 0
return true, nil
}
n, err := strconv.Atoi(s)
if err != nil || n < 0 {
return false, err
}
*result = n
return true, nil
}
+118
View File
@@ -0,0 +1,118 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package logs
import (
"context"
"net/http"
"testing"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/model"
)
func setupTestDB(t *testing.T) *gorm.DB {
dbConn, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
if err != nil {
t.Fatalf("failed to open sqlite in memory: %v", err)
}
err = dbConn.AutoMigrate(&model.SystemConfig{})
if err != nil {
t.Fatalf("failed to migrate schema: %v", err)
}
db.SetDB(dbConn)
return dbConn
}
func TestWebSocketCheckOrigin(t *testing.T) {
dbConn := setupTestDB(t)
// Clean up global DB after test
defer db.SetDB(nil)
// Seed ConfigKeyServerAddress with allowed frontend origin
allowedOrigin := "http://localhost:3000"
if err := dbConn.Create(&model.SystemConfig{
Key: model.ConfigKeyServerAddress,
Value: allowedOrigin,
}).Error; err != nil {
t.Fatalf("failed to seed server address config: %v", err)
}
upgrader := getUpgrader()
if upgrader.CheckOrigin == nil {
t.Fatal("expected CheckOrigin to be defined")
}
tests := []struct {
name string
origin string
host string
wantOK bool
}{
{
name: "empty origin (non-browser clients)",
origin: "",
host: "localhost:8000",
wantOK: true,
},
{
name: "same-origin request",
origin: "http://localhost:8000",
host: "localhost:8000",
wantOK: true,
},
{
name: "same-origin request case insensitive",
origin: "HTTP://LOCALHOST:8000",
host: "localhost:8000",
wantOK: true,
},
{
name: "configured allowed origin request",
origin: "http://localhost:3000",
host: "localhost:8000",
wantOK: true,
},
{
name: "configured allowed origin request with trailing slash",
origin: "http://localhost:3000/",
host: "localhost:8000",
wantOK: true,
},
{
name: "unauthorized third-party origin",
origin: "http://evil.com",
host: "localhost:8000",
wantOK: false,
},
{
name: "invalid origin format",
origin: "::not-a-valid-url",
host: "localhost:8000",
wantOK: false,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
req, err := http.NewRequestWithContext(context.Background(), "GET", "/api/v1/admin/logs/ws", nil)
if err != nil {
t.Fatalf("failed to create request: %v", err)
}
req.Host = tt.host
if tt.origin != "" {
req.Header.Set("Origin", tt.origin)
}
got := upgrader.CheckOrigin(req)
if got != tt.wantOK {
t.Errorf("CheckOrigin() = %v, want %v", got, tt.wantOK)
}
})
}
}
+46
View File
@@ -0,0 +1,46 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package admin
import (
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/pkg/logger"
otel_trace "github.com/Rain-kl/Wavelet/pkg/trace"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/gin-gonic/gin"
)
// LoginAdminRequired 返回管理员权限校验中间件
func LoginAdminRequired() gin.HandlerFunc {
return func(c *gin.Context) {
// init trace
ctx, span := otel_trace.Start(c.Request.Context(), "LoginAdminRequired")
defer span.End()
user, _ := oauth.GetFromContext[*model.User](c, oauth.UserObjKey)
// 如果是通过 Access Token 鉴权,需要检查令牌本身是否具有管理员权限
if tokenAuth, _ := oauth.GetFromContext[bool](c, oauth.TokenAuthKey); tokenAuth {
tokenAdmin, _ := oauth.GetFromContext[bool](c, oauth.TokenAdminKey)
if !tokenAdmin {
response.AbortNotFound(c, TokenAdminRequired)
return
}
}
if !user.IsAdmin {
response.AbortNotFound(c, AdminRequired)
return
}
// log
logger.InfoF(ctx, "[LoginAdminRequired] %d %s", user.ID, user.Username)
// next
c.Next()
}
}
+278
View File
@@ -0,0 +1,278 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
import (
"encoding/json"
"errors"
"net/http"
"strconv"
"strings"
pkgpush "github.com/Rain-kl/Wavelet/pkg/push"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/model"
)
// ListChannelDefinitions 获取各种消息通道的表单配置定义列表
// @Summary 获取所有消息通道配置字段定义
// @Description 返回系统支持的所有消息通道类型(如飞书、邮件、自定义、Telegram)的动态表单定义,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]Definition} "通道配置定义列表"
// @Router /api/v1/admin/push/channels/definitions [get]
func ListChannelDefinitions(c *gin.Context) {
c.JSON(http.StatusOK, response.OK(ListDefinitions()))
}
// ListChannels 获取消息通道列表
// @Summary 获取所有消息通道
// @Description 返回系统配置的所有消息通道列表,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]model.PushChannel} "消息通道列表"
// @Router /api/v1/admin/push/channels [get]
func ListChannels(c *gin.Context) {
channels, err := listPushChannels(c.Request.Context())
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(channels))
}
// CreateChannelRequest 创建通道参数
type CreateChannelRequest struct {
Name string `json:"name" binding:"required"`
Description string `json:"description"`
Type string `json:"type" binding:"required"`
Token string `json:"token"`
URL string `json:"url"`
Other string `json:"other"`
Enabled bool `json:"enabled"`
}
// CreateChannel 创建消息通道
// @Summary 创建消息通道
// @Description 新建一个消息通道配置,需要管理员权限
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body CreateChannelRequest true "创建参数"
// @Success 200 {object} response.Any{data=model.PushChannel} "创建成功"
// @Router /api/v1/admin/push/channels [post]
func CreateChannel(c *gin.Context) {
var req CreateChannelRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
channel, err := createPushChannel(c.Request.Context(), req)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(channel))
}
// UpdateChannelRequest 修改通道参数
type UpdateChannelRequest struct {
Description string `json:"description"`
Type string `json:"type" binding:"required"`
Token string `json:"token"`
URL string `json:"url"`
Other string `json:"other"`
Enabled bool `json:"enabled"`
}
// UpdateChannel 更新消息通道
// @Summary 更新消息通道
// @Description 修改消息通道配置,需要管理员权限
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path uint64 true "通道ID"
// @Param request body UpdateChannelRequest true "更新参数"
// @Success 200 {object} response.Any{data=model.PushChannel} "更新成功"
// @Router /api/v1/admin/push/channels/{id} [put]
func UpdateChannel(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "invalid channel id")
return
}
var req UpdateChannelRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
channel, err := updatePushChannel(c.Request.Context(), id, req)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, "channel not found")
return
}
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(channel))
}
// DeleteChannel 删除消息通道
// @Summary 删除消息通道
// @Description 根据ID删除消息通道,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Param id path uint64 true "通道ID"
// @Success 200 {object} response.Any "删除成功"
// @Router /api/v1/admin/push/channels/{id} [delete]
func DeleteChannel(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "invalid channel id")
return
}
if err := deletePushChannel(c.Request.Context(), id); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, "channel not found")
return
}
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// TestChannelRequest 测试通道连通性参数
type TestChannelRequest struct {
Name string `json:"name"`
Type string `json:"type"`
Token string `json:"token"`
URL string `json:"url"`
Other string `json:"other"`
Target string `json:"target"`
}
// TestChannel 测试通道连通性
// @Summary 测试通道连通性
// @Description 触发一次临时的或现有的通道连通性推送测试,需要管理员权限
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body TestChannelRequest true "测试参数"
// @Success 200 {object} response.Any "测试触发成功"
// @Router /api/v1/admin/push/channels/test [post]
func TestChannel(c *gin.Context) {
var req TestChannelRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
ctx := c.Request.Context()
url, token, other, channelType, err := loadChannelForTest(ctx, req)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if channelType == channelEmail {
url, token, other = resolveSMTPConfig(ctx, url, token, other)
}
tempChannel := model.PushChannel{
Name: "test_temp",
URL: url,
Token: token,
Other: other,
Type: channelType,
Enabled: true,
}
if err := tempChannel.Validate(); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
url = tempChannel.URL
var config pkgpush.Config
var renderedJSON string
switch channelType {
case channelLark:
config = pkgpush.Config{Channel: channelLark, URL: url, Secret: token}
renderedJSON = other
case channelEmail:
config = pkgpush.Config{Channel: channelEmail, URL: url, Key: token, Secret: other}
case channelTelegram:
config = pkgpush.Config{Channel: channelTelegram, URL: url, Secret: token, Key: other}
default:
config = pkgpush.Config{Channel: channelCustom, URL: url}
customPushReq := CustomPushRequest{
Title: "通道测试通知",
Content: "这是一条来自系统的消息通道连通性测试消息。",
Description: "系统通道测试",
URL: "https://example.com",
To: req.Target,
}
renderedJSON = renderCustomPayload(other, customPushReq)
}
payload := SendPayload{
EventKey: "test_channel",
Config: config,
Target: req.Target,
Body: NotificationMessage{
Title: "通道测试通知",
Content: "这是一条来自系统的消息通道连通性测试消息。",
Level: defaultLevelInfo,
},
Template: renderedJSON,
}
if err := enqueuePushTask(ctx, payload); err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// CustomPushRequest 外部公开推送请求参数
type CustomPushRequest struct {
Title string `json:"title" form:"title"`
Description string `json:"description" form:"description"`
Content string `json:"content" form:"content"`
URL string `json:"url" form:"url"`
To string `json:"to" form:"to"`
Token string `json:"token" form:"token"`
}
func escapeJSONString(s string) string {
b, _ := json.Marshal(s)
const minJSONLen = 2
if len(b) >= minJSONLen {
return string(b[1 : len(b)-1])
}
return s
}
func renderCustomPayload(template string, req CustomPushRequest) string {
result := template
result = strings.ReplaceAll(result, "$title", escapeJSONString(req.Title))
result = strings.ReplaceAll(result, "$description", escapeJSONString(req.Description))
result = strings.ReplaceAll(result, "$content", escapeJSONString(req.Content))
result = strings.ReplaceAll(result, "$url", escapeJSONString(req.URL))
result = strings.ReplaceAll(result, "$to", escapeJSONString(req.To))
return result
}
@@ -0,0 +1,183 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
import "sync"
const (
// KeyURL represents the URL field key
KeyURL = "url"
// KeyToken represents the Token field key
KeyToken = "token"
// KeyOther represents the Other field key
KeyOther = "other"
// TypeText represents standard text input type
TypeText = "text"
// TypePassword represents password input type
TypePassword = "password"
// TypeTextarea represents textarea input type
TypeTextarea = "textarea"
)
// Field represents a form field configuration for a channel.
type Field struct {
Key string `json:"key"` // unique key for the field (e.g. url, token, other)
Label string `json:"label"` // human readable label (e.g. "Webhook 地址")
Type string `json:"type"` // input type: "text" | "password" | "textarea"
Required bool `json:"required"` // whether this field is required
Placeholder string `json:"placeholder"` // input placeholder
Description string `json:"description"` // field explanation/help text
}
// Definition represents the metadata and form schema for a notification channel.
type Definition struct {
Type string `json:"type"` // channel type (e.g., custom, lark, email)
Name string `json:"name"` // display name
Description string `json:"description"` // short description
Fields []Field `json:"fields"` // form fields
}
var (
defMu sync.RWMutex
definitions = make(map[string]Definition)
)
// RegisterChannelDefinition registers a channel definition.
func RegisterChannelDefinition(def Definition) {
defMu.Lock()
defer defMu.Unlock()
definitions[def.Type] = def
}
// ListDefinitions returns all registered channel definitions.
func ListDefinitions() []Definition {
defMu.RLock()
defer defMu.RUnlock()
// We want a stable order: custom, lark, telegram, email
order := []string{channelCustom, channelLark, channelTelegram, channelEmail}
res := make([]Definition, 0, len(definitions))
for _, t := range order {
if d, ok := definitions[t]; ok {
res = append(res, d)
}
}
// Add any others
for t, d := range definitions {
found := false
for _, o := range order {
if o == t {
found = true
break
}
}
if !found {
res = append(res, d)
}
}
return res
}
func init() {
// Register custom webhook channel
RegisterChannelDefinition(Definition{
Type: channelCustom,
Name: "自定义消息通道",
Description: "使用自定义 HTTP POST 请求向外部 Webhook 发送数据。",
Fields: []Field{
{
Key: KeyURL,
Label: "请求地址",
Type: TypeText,
Required: true,
Placeholder: "在此填写完整的请求地址,必须使用 HTTPS 协议",
Description: "接口请求的完整 HTTPS URL,例如 https://api.example.com/webhook",
},
{
Key: KeyOther,
Label: "请求体 (JSON)",
Type: TypeTextarea,
Required: true,
Placeholder: "在此输入请求体,支持模板变量,必须为合法的 JSON 格式",
Description: "可使用的变量:$title, $description, $content, $url, $to。例如 {\"text\": \"$content\"}",
},
},
})
// Register Lark robot channel
RegisterChannelDefinition(Definition{
Type: channelLark,
Name: "飞书群机器人",
Description: "配置飞书群自定义机器人的 Webhook 接口投递。",
Fields: []Field{
{
Key: KeyURL,
Label: "Webhook 地址",
Type: TypeText,
Required: true,
Placeholder: "https://open.feishu.cn/open-apis/bot/v2/hook/YOUR_TOKEN",
Description: "从飞书群机器人设置中复制 of Webhook URL",
// Note: using 'of' was in feishu.go, let's keep original wording or fix it
},
{
Key: KeyToken,
Label: "签名校验密钥 (Secret) (可选)",
Type: TypeText,
Required: false,
Placeholder: "可选,若机器人启用了安全设置中的签名校验,请在此输入",
Description: "飞书群机器人安全设置中的签名校验 Key",
},
{
Key: KeyOther,
Label: "自定义卡片 JSON 模版 (可选)",
Type: TypeTextarea,
Required: false,
Placeholder: "可选,留空则默认使用系统内置的精美互动卡片",
Description: "若填写,必须是合法的飞书卡片 JSON 格式",
},
},
})
// Register Telegram channel
RegisterChannelDefinition(Definition{
Type: channelTelegram,
Name: "Telegram 机器人",
Description: "配置 Telegram 机器人推送消息。",
Fields: []Field{
{
Key: KeyURL,
Label: "API 基础地址 (可选)",
Type: TypeText,
Required: false,
Placeholder: "https://api.telegram.org",
Description: "接口请求的 HTTPS 基础地址,留空默认为 https://api.telegram.org",
},
{
Key: KeyToken,
Label: "机器人 Token (Bot Token)",
Type: TypePassword,
Required: true,
Placeholder: "在此输入 Telegram 机器人的 Bot Token",
Description: "通过 BotFather 申请到的机器人 Access Token",
},
{
Key: KeyOther,
Label: "默认会话 ID (Chat ID) (可选)",
Type: TypeText,
Required: false,
Placeholder: "例如 -100123456789 或 @channel_name",
Description: "默认的消息接收 Chat ID。如果通知事件中未配置 targets,将推送到此 ID",
},
},
})
// Register Email channel
RegisterChannelDefinition(Definition{
Type: channelEmail,
Name: "邮件推送通道",
Description: "邮件推送通道直接使用系统全局 SMTP 设置进行发送,无需在此填写服务器配置。",
Fields: []Field{},
})
}
+15
View File
@@ -0,0 +1,15 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
const (
channelCustom = "custom"
channelEmail = "email"
channelLark = "lark"
channelTelegram = "telegram"
defaultLevelInfo = "INFO"
keyTitle = "title"
keyContent = "content"
keyLevel = "level"
)
@@ -0,0 +1,38 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package custom_events defines custom push notification events.
package custom_events
import (
"context"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/admin/push"
"github.com/Rain-kl/Wavelet/internal/listener"
)
// AdminLogin is the metadata definition for the admin login event.
var AdminLogin = push.EventMetadata{
Key: "admin_login",
Name: "管理员登录",
DefaultTemplate: push.NotificationMessage{
Title: "管理员登录提醒",
Content: "管理员 {{user.username}} 于 {{time}} 从 IP {{ip}} 登录系统。",
Level: "INFO",
},
Description: "当管理员成功登录系统时触发此通知",
}
func handleAdminLogin(ctx context.Context, event listener.AdminLoggedIn) {
if event.User == nil {
return
}
body := map[string]any{
"user": event.User,
"ip": event.IP,
"time": time.Now().Format("2006-01-02 15:04:05"),
}
push.DefaultTrigger.Trigger(ctx, AdminLogin, body)
}
@@ -0,0 +1,177 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package custom_events
import (
"context"
"encoding/json"
"sync"
"testing"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/admin/push"
"github.com/Rain-kl/Wavelet/internal/listener"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/hibiken/asynq"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
var registerOnce sync.Once
func ensureRegistered() {
registerOnce.Do(Register)
}
func setupAdminLoginIntegrationTest(t *testing.T) (*gorm.DB, func()) {
t.Helper()
dbConn, mr, cleanup := testhelper.SetupTestEnvironment(t)
err := dbConn.AutoMigrate(
&model.PushEvent{},
&model.PushHistory{},
&model.PushChannel{},
)
require.NoError(t, err)
sysUser := &model.User{
ID: 999,
Username: "system",
Nickname: "系统",
Password: "*",
IsActive: true,
}
require.NoError(t, dbConn.Create(sysUser).Error)
task.AsynqClient = asynq.NewClient(asynq.RedisClientOpt{Addr: mr.Addr()})
task.RegisterHandler(push.SendNotificationTask, &push.PushHandler{})
task.RegisterTaskMeta(push.SendNotificationMeta)
ensureRegistered()
require.NoError(t, push.SyncEvents(context.Background()))
return dbConn, func() {
cleanup()
if task.AsynqClient != nil {
task.AsynqClient.Close()
task.AsynqClient = nil
}
}
}
func seedMockPushChannel(t *testing.T, dbConn *gorm.DB) *model.PushChannel {
t.Helper()
channel := &model.PushChannel{
Name: "mock_channel",
Type: "custom",
URL: "https://webhook.site/admin-login",
Other: `{"text": "$content"}`,
Enabled: true,
}
require.NoError(t, dbConn.Create(channel).Error)
return channel
}
func enableAdminLoginEvent(t *testing.T, dbConn *gorm.DB, channelName string, targets []string) {
t.Helper()
var event model.PushEvent
require.NoError(t, dbConn.Where("event_key = ?", AdminLogin.Key).First(&event).Error)
event.Enabled = true
event.Channels = []string{channelName}
event.Targets = targets
require.NoError(t, dbConn.Save(&event).Error)
}
func waitForAsyncTrigger(t *testing.T) {
t.Helper()
time.Sleep(100 * time.Millisecond)
}
func countPushTasks(t *testing.T, dbConn *gorm.DB) int64 {
t.Helper()
var count int64
require.NoError(t, dbConn.Model(&model.TaskExecution{}).
Where("task_type = ?", push.SendNotificationTask).
Count(&count).Error)
return count
}
func TestAdminLoginPushIntegration(t *testing.T) {
dbConn, cleanup := setupAdminLoginIntegrationTest(t)
defer cleanup()
channel := seedMockPushChannel(t, dbConn)
defer dbConn.Delete(channel)
enableAdminLoginEvent(t, dbConn, channel.Name, []string{"ops_team"})
adminUser := &model.User{
ID: 1001,
Username: "super_admin",
IsAdmin: true,
IsActive: true,
}
require.NoError(t, dbConn.Create(adminUser).Error)
t.Run("admin login emits push task with user and ip", func(t *testing.T) {
dbConn.Where("task_type = ?", push.SendNotificationTask).Delete(&model.TaskExecution{})
listener.EmitAdminLoggedIn(context.Background(), adminUser, "203.0.113.42")
waitForAsyncTrigger(t)
var execution model.TaskExecution
require.NoError(t, dbConn.Where("task_type = ?", push.SendNotificationTask).First(&execution).Error)
var payload push.SendPayload
require.NoError(t, json.Unmarshal([]byte(execution.Payload), &payload))
assert.Equal(t, AdminLogin.Key, payload.EventKey)
assert.Equal(t, "ops_team", payload.Target)
assert.Equal(t, "管理员登录提醒", payload.Body.Title)
assert.Contains(t, payload.Body.Content, "super_admin")
assert.Contains(t, payload.Body.Content, "203.0.113.42")
})
t.Run("non-admin login does not trigger push", func(t *testing.T) {
dbConn.Where("task_type = ?", push.SendNotificationTask).Delete(&model.TaskExecution{})
nonAdmin := &model.User{
ID: 2002,
Username: "regular_user",
IsAdmin: false,
IsActive: true,
}
require.NoError(t, dbConn.Create(nonAdmin).Error)
listener.EmitAdminLoggedIn(context.Background(), nonAdmin, "198.51.100.1")
waitForAsyncTrigger(t)
assert.Equal(t, int64(0), countPushTasks(t, dbConn))
})
t.Run("disabled admin login event does not enqueue push", func(t *testing.T) {
dbConn.Where("task_type = ?", push.SendNotificationTask).Delete(&model.TaskExecution{})
var event model.PushEvent
require.NoError(t, dbConn.Where("event_key = ?", AdminLogin.Key).First(&event).Error)
event.Enabled = false
require.NoError(t, dbConn.Save(&event).Error)
repository.DeleteActivePushEventCache(context.Background(), AdminLogin.Key)
listener.EmitAdminLoggedIn(context.Background(), adminUser, "10.0.0.1")
waitForAsyncTrigger(t)
assert.Equal(t, int64(0), countPushTasks(t, dbConn))
})
}
@@ -0,0 +1,17 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package custom_events
import (
"github.com/Rain-kl/Wavelet/internal/apps/admin/push"
"github.com/Rain-kl/Wavelet/internal/listener"
)
// Register wires push notification handlers for domain events and registers
// built-in event metadata. Must be called once during application bootstrap
// before push.SyncEvents.
func Register() {
push.RegisterBuiltInEvent(AdminLogin)
listener.OnAdminLoggedIn(handleAdminLogin)
}
+343
View File
@@ -0,0 +1,343 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package push defines push notification HTTP routes, background tasks, and events.
package push
import (
"context"
"encoding/json"
"errors"
"fmt"
"strings"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/pkg/logger"
pkgpush "github.com/Rain-kl/Wavelet/pkg/push"
"gorm.io/gorm"
)
// NotificationMessage represents the structured notification message payload.
type NotificationMessage struct {
Title string `json:"title"`
Content string `json:"content"`
Level string `json:"level"`
Ext map[string]any `json:"ext,omitempty"`
}
// Flatten converts the structured NotificationMessage back to a flat map (original json structure).
func (m NotificationMessage) Flatten() map[string]any {
res := map[string]any{
keyTitle: m.Title,
keyContent: m.Content,
keyLevel: m.Level,
}
for k, v := range m.Ext {
res[k] = v
}
return res
}
// EventMetadata represents the metadata of a push notification event.
type EventMetadata struct {
Key string `json:"key"`
Name string `json:"name"`
DefaultTemplate NotificationMessage `json:"default_template"`
Description string `json:"description"`
}
// SendPayload 异步投递推送载荷 (供 task/Worker 使用)
type SendPayload struct {
EventKey string `json:"event_key"`
Config pkgpush.Config `json:"config"`
Target string `json:"target"`
Body NotificationMessage `json:"body"`
Template string `json:"template"`
}
// BuiltInEvents lists all built-in events defined in custom_events.
var BuiltInEvents []EventMetadata
// RegisterBuiltInEvent registers a built-in event definition.
func RegisterBuiltInEvent(meta EventMetadata) {
BuiltInEvents = append(BuiltInEvents, meta)
}
// EventTrigger represents the unified event trigger class.
type EventTrigger struct{}
// DefaultTrigger is the singleton instance of EventTrigger.
var DefaultTrigger = &EventTrigger{}
// Trigger receives event metadata and processes the event notification dispatch asynchronously.
//
//nolint:contextcheck
func (t *EventTrigger) Trigger(ctx context.Context, meta EventMetadata, body map[string]any) {
asyncCtx := context.WithoutCancel(ctx)
go func() {
if body == nil {
body = make(map[string]any)
}
if _, hasUser := body["user"]; !hasUser || body["user"] == nil {
body["user"] = getSystemUser(asyncCtx)
}
eventPtr, err := repository.GetActivePushEventByKey(asyncCtx, meta.Key)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return
}
logger.ErrorF(asyncCtx, "push_event_trigger: failed to get active event %s: %v", meta.Key, err)
return
}
event := *eventPtr
if len(event.Channels) == 0 {
return
}
flatBody := getFlatBody(body)
msg, _ := t.buildMessage(&event, meta, flatBody, body)
t.enqueuePushTasks(asyncCtx, meta, &event, msg, flatBody)
}()
}
func (t *EventTrigger) buildMessage(event *model.PushEvent, meta EventMetadata, flatBody map[string]any, body map[string]any) (NotificationMessage, string) {
var msg NotificationMessage
renderedTemplate := ""
templateSource := event.Template
if templateSource != "" {
var err error
msg, renderedTemplate, err = t.parseCustomTemplate(event, templateSource, flatBody)
if err != nil {
msg.Title = event.Name
msg.Content = renderedTemplate
msg.Level = defaultLevelInfo
}
} else {
msg = t.parseDefaultTemplate(meta, flatBody)
}
if msg.Ext == nil {
msg.Ext = make(map[string]any)
}
for k, v := range body {
if k == keyTitle || k == keyContent || k == keyLevel {
continue
}
if _, exists := msg.Ext[k]; !exists {
msg.Ext[k] = v
}
}
return msg, renderedTemplate
}
func (t *EventTrigger) parseCustomTemplate(event *model.PushEvent, templateSource string, flatBody map[string]any) (NotificationMessage, string, error) {
var msg NotificationMessage
renderedTemplate := pkgpush.ParseTemplate(templateSource, flatBody)
var tMap map[string]any
if err := json.Unmarshal([]byte(renderedTemplate), &tMap); err != nil {
return msg, renderedTemplate, err
}
if title, ok := tMap[keyTitle].(string); ok && title != "" {
msg.Title = title
} else {
msg.Title = event.Name
}
delete(tMap, keyTitle)
if content, ok := tMap[keyContent].(string); ok && content != "" {
msg.Content = content
} else {
msg.Content = renderedTemplate
}
delete(tMap, keyContent)
if level, ok := tMap[keyLevel].(string); ok && level != "" {
msg.Level = level
} else {
msg.Level = defaultLevelInfo
}
delete(tMap, keyLevel)
msg.Ext = tMap
return msg, renderedTemplate, nil
}
func (t *EventTrigger) parseDefaultTemplate(meta EventMetadata, flatBody map[string]any) NotificationMessage {
var msg NotificationMessage
msg.Title = pkgpush.ParseTemplate(meta.DefaultTemplate.Title, flatBody)
msg.Content = pkgpush.ParseTemplate(meta.DefaultTemplate.Content, flatBody)
msg.Level = pkgpush.ParseTemplate(meta.DefaultTemplate.Level, flatBody)
if meta.DefaultTemplate.Ext != nil {
msg.Ext = make(map[string]any)
for k, v := range meta.DefaultTemplate.Ext {
if strVal, ok := v.(string); ok {
msg.Ext[k] = pkgpush.ParseTemplate(strVal, flatBody)
} else {
msg.Ext[k] = v
}
}
}
return msg
}
func (t *EventTrigger) enqueuePushTasks(ctx context.Context, meta EventMetadata, event *model.PushEvent, msg NotificationMessage, flatBody map[string]any) {
for _, channelName := range event.Channels {
customChannel, err := repository.GetActivePushChannelByName(ctx, channelName)
if err == nil {
t.enqueueCustomPushChannelTasks(ctx, meta, event, customChannel, msg, flatBody)
continue
}
logger.WarnF(ctx, "push_event_trigger: channel %q not found in DB or disabled: %v", channelName, err)
}
}
func (t *EventTrigger) enqueueCustomPushChannelTasks(ctx context.Context, meta EventMetadata, event *model.PushEvent, channel *model.PushChannel, msg NotificationMessage, flatBody map[string]any) {
if len(event.Targets) == 0 {
t.enqueueSingleCustomPushChannelTask(ctx, meta, channel, "", msg)
return
}
for _, target := range event.Targets {
resolvedTarget := resolveTarget(ctx, target, flatBody, channel.Name)
t.enqueueSingleCustomPushChannelTask(ctx, meta, channel, resolvedTarget, msg)
}
}
func (t *EventTrigger) enqueueSingleCustomPushChannelTask(ctx context.Context, meta EventMetadata, channel *model.PushChannel, target string, msg NotificationMessage) {
var config pkgpush.Config
var renderedTemplate string
switch channel.Type {
case channelLark:
config = pkgpush.Config{Channel: channelLark, URL: channel.URL, Secret: channel.Token}
renderedTemplate = channel.Other
case channelEmail:
url, token, other := resolveSMTPConfig(ctx, channel.URL, channel.Token, channel.Other)
config = pkgpush.Config{Channel: channelEmail, URL: url, Key: token, Secret: other}
case channelTelegram:
config = pkgpush.Config{Channel: channelTelegram, URL: channel.URL, Secret: channel.Token, Key: channel.Other}
default:
config = pkgpush.Config{Channel: channelCustom, URL: channel.URL}
customPushReq := CustomPushRequest{
Title: msg.Title,
Content: msg.Content,
Description: meta.Description,
To: target,
}
if urlVal, ok := msg.Ext["url"].(string); ok {
customPushReq.URL = urlVal
}
renderedTemplate = renderCustomPayload(channel.Other, customPushReq)
}
payload := SendPayload{
EventKey: meta.Key,
Config: config,
Target: target,
Body: msg,
Template: renderedTemplate,
}
if err := enqueuePushTask(ctx, payload); err != nil {
logger.ErrorF(ctx, "push_event_trigger: enqueuePushTask failed for %s channel %s -> %s: %v", channel.Type, channel.Name, target, err)
}
}
func enqueuePushTask(ctx context.Context, payload SendPayload) error {
payloadBytes, err := json.Marshal(payload)
if err != nil {
return err
}
_, err = task.DispatchTask(ctx, "send_notification", payloadBytes, "system")
return err
}
func getFlatBody(body map[string]any) map[string]any {
jsonBytes, err := json.Marshal(body)
if err != nil {
return body
}
var jsonMap map[string]any
if err := json.Unmarshal(jsonBytes, &jsonMap); err != nil {
return body
}
flatResult := make(map[string]any)
flattenMap("", jsonMap, flatResult)
return flatResult
}
func flattenMap(prefix string, m map[string]any, result map[string]any) {
for k, v := range m {
key := k
if prefix != "" {
key = prefix + "." + k
}
if nestedMap, ok := v.(map[string]any); ok {
flattenMap(key, nestedMap, result)
} else {
result[key] = v
}
}
}
func resolveTarget(ctx context.Context, target string, flatBody map[string]any, channel string) string {
target = strings.TrimSpace(target)
if target == "" {
return ""
}
resolved := resolveDynamicKeyword(target, flatBody)
if strings.Contains(resolved, "@") {
return resolved
}
if val, matched := resolveSystemTarget(ctx, resolved, channel); matched {
return val
}
user, found := resolveTargetUser(ctx, resolved, channel)
if !found {
return resolved
}
if channel == channelEmail && user.Email != "" {
return user.Email
}
if channel != channelEmail && user.Username != "" {
return user.Username
}
return resolved
}
func resolveDynamicKeyword(target string, flatBody map[string]any) string {
switch target {
case "user.id", "id":
if val, ok := flatBody["user.id"]; ok {
return fmt.Sprintf("%v", val)
}
if val, ok := flatBody["id"]; ok {
return fmt.Sprintf("%v", val)
}
case "user.username", "username":
if val, ok := flatBody["user.username"]; ok {
return fmt.Sprintf("%v", val)
}
if val, ok := flatBody["username"]; ok {
return fmt.Sprintf("%v", val)
}
case "user.email", channelEmail:
if val, ok := flatBody["user.email"]; ok {
return fmt.Sprintf("%v", val)
}
if val, ok := flatBody["email"]; ok {
return fmt.Sprintf("%v", val)
}
}
return target
}
+413
View File
@@ -0,0 +1,413 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
import (
"context"
"encoding/json"
"errors"
"strconv"
"strings"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/task"
pkgpush "github.com/Rain-kl/Wavelet/pkg/push"
"gorm.io/gorm"
)
type smtpConfig struct {
Host string
Port string
Username string
Password string
}
func loadSMTPConfig(ctx context.Context) smtpConfig {
host, _ := repository.GetSystemConfigByKey(ctx, model.ConfigKeySMTPHost)
port, _ := repository.GetSystemConfigByKey(ctx, model.ConfigKeySMTPPort)
user, _ := repository.GetSystemConfigByKey(ctx, model.ConfigKeySMTPUsername)
pass, _ := repository.GetSystemConfigByKey(ctx, model.ConfigKeySMTPPassword)
return smtpConfig{
Host: host.Value,
Port: port.Value,
Username: user.Value,
Password: pass.Value,
}
}
func syncBuiltInEvents(ctx context.Context) error {
for _, meta := range BuiltInEvents {
_, err := repository.GetPushEventByKey(ctx, meta.Key)
if errors.Is(err, gorm.ErrRecordNotFound) {
var defaultTemplateStr string
if defaultTemplateBytes, err := json.Marshal(meta.DefaultTemplate); err == nil {
defaultTemplateStr = string(defaultTemplateBytes)
}
event := model.PushEvent{
EventKey: meta.Key,
Name: meta.Name,
Channels: []string{},
Targets: []string{},
Template: defaultTemplateStr,
Enabled: false,
}
if err := repository.CreatePushEvent(ctx, &event); err != nil {
return err
}
} else if err != nil {
return err
}
}
return nil
}
func listPushEvents(ctx context.Context) ([]model.PushEvent, error) {
return repository.ListPushEvents(ctx)
}
func createPushEvent(ctx context.Context, req CreateEventRequest) (model.PushEvent, error) {
eventKey, eventName, defaultTemplateBytes, err := getEventInfo(req)
if err != nil {
return model.PushEvent{}, err
}
count, err := repository.CountPushEventsByKey(ctx, eventKey)
if err != nil {
return model.PushEvent{}, err
}
if count > 0 {
return model.PushEvent{}, errors.New("this notification event is already configured")
}
templateStr := strings.TrimSpace(req.Template)
if templateStr == "" {
templateStr = string(defaultTemplateBytes)
} else {
var tempMap map[string]any
if err := json.Unmarshal([]byte(templateStr), &tempMap); err != nil {
return model.PushEvent{}, errors.New("custom template is not a valid JSON format")
}
}
channels := req.Channels
if channels == nil {
channels = []string{}
}
targets := req.Targets
if targets == nil {
targets = []string{}
}
event := model.PushEvent{
EventKey: eventKey,
Name: eventName,
TaskType: req.TaskType,
Channels: channels,
Targets: targets,
Template: templateStr,
Enabled: req.Enabled,
}
if err := event.Validate(); err != nil {
return model.PushEvent{}, err
}
if err := repository.CreatePushEvent(ctx, &event); err != nil {
return model.PushEvent{}, err
}
return event, nil
}
func deletePushEvent(ctx context.Context, id uint64) error {
event, err := repository.GetPushEventByID(ctx, id)
if err != nil {
return err
}
return repository.DeletePushEvent(ctx, &event)
}
func updatePushEvent(ctx context.Context, id uint64, req UpdateEventRequest) error {
event, err := repository.GetPushEventByID(ctx, id)
if err != nil {
return err
}
event.Channels = req.Channels
event.Targets = req.Targets
event.Template = req.Template
event.Enabled = req.Enabled
if err := event.Validate(); err != nil {
return err
}
return repository.SavePushEvent(ctx, &event)
}
func togglePushEvent(ctx context.Context, id uint64) (bool, error) {
event, err := repository.GetPushEventByID(ctx, id)
if err != nil {
return false, err
}
enabled := !event.Enabled
if enabled && len(event.Channels) == 0 {
return false, errors.New("cannot enable event without any push channels configured")
}
if err := repository.UpdatePushEventEnabled(ctx, &event, enabled); err != nil {
return false, err
}
return enabled, nil
}
func listPushHistories(ctx context.Context, filter repository.PushHistoryListFilter) (int64, []model.PushHistory, error) {
return repository.ListPushHistories(ctx, filter)
}
func applySMTPFallbackToPushConfig(ctx context.Context, cfg *pkgpush.Config) {
if cfg.Channel != channelEmail || (cfg.URL != "" && cfg.Key != "") {
return
}
smtp := loadSMTPConfig(ctx)
if smtp.Host == "" || smtp.Username == "" {
return
}
port := smtp.Port
if port == "" {
port = "587"
}
cfg.URL = smtp.Host + ":" + port
cfg.Key = smtp.Username
cfg.Secret = smtp.Password
}
func listPushChannels(ctx context.Context) ([]model.PushChannel, error) {
return repository.ListPushChannels(ctx)
}
func createPushChannel(ctx context.Context, req CreateChannelRequest) (model.PushChannel, error) {
count, err := repository.CountPushChannelsByName(ctx, req.Name)
if err != nil {
return model.PushChannel{}, err
}
if count > 0 {
return model.PushChannel{}, errors.New("channel name already exists")
}
channel := model.PushChannel{
Name: req.Name,
Description: req.Description,
Type: req.Type,
Token: req.Token,
URL: req.URL,
Other: req.Other,
Enabled: req.Enabled,
}
if err := channel.Validate(); err != nil {
return model.PushChannel{}, err
}
if err := repository.CreatePushChannel(ctx, &channel); err != nil {
return model.PushChannel{}, err
}
return channel, nil
}
func updatePushChannel(ctx context.Context, id uint64, req UpdateChannelRequest) (model.PushChannel, error) {
channel, err := repository.GetPushChannelByID(ctx, id)
if err != nil {
return model.PushChannel{}, err
}
channel.Description = req.Description
channel.Type = req.Type
channel.Token = req.Token
channel.URL = req.URL
channel.Other = req.Other
channel.Enabled = req.Enabled
if err := channel.Validate(); err != nil {
return model.PushChannel{}, err
}
if err := repository.SavePushChannel(ctx, &channel); err != nil {
return model.PushChannel{}, err
}
return channel, nil
}
func deletePushChannel(ctx context.Context, id uint64) error {
channel, err := repository.GetPushChannelByID(ctx, id)
if err != nil {
return err
}
return repository.DeletePushChannel(ctx, &channel)
}
func loadChannelForTest(ctx context.Context, req TestChannelRequest) (string, string, string, string, error) {
if req.Name != "" {
channel, err := repository.GetPushChannelByName(ctx, req.Name)
if err != nil {
return "", "", "", "", errors.New("channel not found")
}
return channel.URL, channel.Token, channel.Other, channel.Type, nil
}
return req.URL, req.Token, req.Other, req.Type, nil
}
func listActivePushEventsByTaskType(ctx context.Context, taskType string) ([]model.PushEvent, error) {
return repository.ListActivePushEventsByTaskType(ctx, taskType)
}
func loadUserFromPayload(ctx context.Context, data map[string]any) any {
if u, exists := data["user"]; exists && u != nil {
return u
}
if userID, ok := extractUserID(data); ok && userID > 0 {
if user, err := repository.GetUserByID(ctx, userID); err == nil {
return &user
}
}
if username := extractUsername(data); username != "" {
if user, err := repository.GetUserByUsername(ctx, username); err == nil {
return &user
}
}
return nil
}
func recordPushHistory(ctx context.Context, req SendPayload, status, errMsg string) error {
title := req.Body.Title
content := req.Body.Content
level := req.Body.Level
if title == "" {
title = "系统通知"
}
if level == "" {
level = defaultLevelInfo
}
target := req.Target
if target == "" {
if req.Config.URL != "" {
target = req.Config.URL
const maxTargetLen = 50
const truncatedLen = 47
if len(target) > maxTargetLen {
target = target[:truncatedLen] + "..."
}
} else {
target = "default"
}
}
history := model.PushHistory{
EventKey: req.EventKey,
Channel: req.Config.Channel,
Target: target,
Title: title,
Content: content,
Level: level,
Status: status,
ErrorMsg: errMsg,
}
return repository.CreatePushHistory(ctx, &history)
}
func resolveTargetUser(ctx context.Context, resolved string, _ string) (model.User, bool) {
found := false
var user model.User
if id, err := strconv.ParseUint(resolved, 10, 64); err == nil {
if u, err := repository.GetUserByID(ctx, id); err == nil {
user = u
found = true
}
}
if !found {
if u, err := repository.GetUserByUsername(ctx, resolved); err == nil {
user = u
found = true
}
}
return user, found
}
func resolveSystemTarget(ctx context.Context, resolved string, channel string) (string, bool) {
if resolved != "系统" && resolved != "system" && resolved != "0" {
return "", false
}
adminUser, err := repository.GetFirstAdminUser(ctx)
if err != nil {
return resolved, true
}
if channel == channelEmail && adminUser.Email != "" {
return adminUser.Email, true
}
if channel != channelEmail && adminUser.Username != "" {
return adminUser.Username, true
}
return resolved, true
}
func resolveSMTPConfig(ctx context.Context, url, token, other string) (string, string, string) {
if url != "" && token != "" {
return url, token, other
}
smtp := loadSMTPConfig(ctx)
if smtp.Host == "" || smtp.Username == "" {
return url, token, other
}
port := smtp.Port
if port == "" {
port = "587"
}
if url == "" {
url = smtp.Host + ":" + port
}
if token == "" {
token = smtp.Username
}
if other == "" {
other = smtp.Password
}
return url, token, other
}
func getSystemUser(ctx context.Context) *model.User {
user := repository.GetSystemUser(ctx)
return &user
}
func getEventInfo(req CreateEventRequest) (string, string, []byte, error) {
if req.TaskType != "" {
meta := task.GetTaskMetaByAsynqTask(req.TaskType)
if meta == nil {
return "", "", nil, errors.New("unsupported task type")
}
eventKey := "task_completed:" + req.TaskType
eventName := "任务完成: " + meta.Name
defaultTemplate := NotificationMessage{
Title: "任务完成: " + meta.Name,
Content: "异步任务 {{task_name}} (ID: {{task_id}}) 已完成。状态: {{task_status}},耗时: {{task_duration}} ms。",
Level: defaultLevelInfo,
}
defaultTemplateBytes, err := json.Marshal(defaultTemplate)
if err != nil {
return "", "", nil, err
}
return eventKey, eventName, defaultTemplateBytes, nil
}
if req.EventKey == "" {
return "", "", nil, errors.New("either event_key or task_type must be provided")
}
meta, found := findBuiltInEvent(req.EventKey)
if !found {
return "", "", nil, errors.New("unsupported built-in event key")
}
defaultTemplateBytes, err := json.Marshal(meta.DefaultTemplate)
if err != nil {
return "", "", nil, err
}
return req.EventKey, meta.Name, defaultTemplateBytes, nil
}
+764
View File
@@ -0,0 +1,764 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"strconv"
"sync"
"testing"
"time"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/internal/testhelper"
pkgpush "github.com/Rain-kl/Wavelet/pkg/push"
"github.com/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
"github.com/hibiken/asynq"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/repository"
)
var adminLoginEvent = EventMetadata{
Key: "admin_login",
Name: "管理员登录",
DefaultTemplate: NotificationMessage{
Title: "管理员登录提醒",
Content: "管理员 {{user.username}} 于 {{time}} 从 IP {{ip}} 登录系统。",
Level: "INFO",
},
Description: "当管理员成功登录系统时触发此通知",
}
func init() {
RegisterBuiltInEvent(adminLoginEvent)
}
// mockPusher mock implementation of pkgpush.Pusher
type mockPusher struct {
mu sync.Mutex
sentBody map[string]any
sentTgt string
}
func (m *mockPusher) Send(ctx context.Context, cfg pkgpush.Config, target string, body map[string]any, template string, ext map[string]any) error {
m.mu.Lock()
defer m.mu.Unlock()
m.sentBody = body
m.sentTgt = target
return nil
}
func (m *mockPusher) ValidateConfig(cfg pkgpush.Config) error {
return nil
}
func setupPushTest(t *testing.T) (*gorm.DB, *miniredis.Miniredis, func()) {
dbConn, mr, cleanup := testhelper.SetupTestEnvironment(t)
// AutoMigrate push tables in SQLite test environment
err := dbConn.AutoMigrate(&model.PushEvent{}, &model.PushHistory{}, &model.User{}, &model.PushChannel{}, &model.SystemConfig{})
require.NoError(t, err)
// 写入数据库系统默认用户 Seed 记录
sysUser := &model.User{
ID: 999,
Username: "system",
Nickname: "系统",
Password: "*",
IsActive: true,
}
err = dbConn.Create(sysUser).Error
require.NoError(t, err)
// Initialize AsynqClient pointing to miniredis
task.AsynqClient = asynq.NewClient(asynq.RedisClientOpt{
Addr: mr.Addr(),
})
// Register the task handler and metadata
task.RegisterHandler(SendNotificationTask, &PushHandler{})
task.RegisterTaskMeta(SendNotificationMeta)
return dbConn, mr, func() {
cleanup()
if task.AsynqClient != nil {
task.AsynqClient.Close()
task.AsynqClient = nil
}
}
}
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin/push")
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, "user_obj", authUser)
}
c.Next()
})
adminGroup.GET("/events", ListEvents)
adminGroup.GET("/events/builtin", ListBuiltInEvents)
adminGroup.POST("/events", CreateEvent)
adminGroup.PUT("/events/:id", UpdateEvent)
adminGroup.DELETE("/events/:id", DeleteEvent)
adminGroup.POST("/events/:id/toggle", ToggleEvent)
adminGroup.GET("/histories", ListHistories)
adminGroup.POST("/test", TestPush)
return r
}
func TestSyncEvents(t *testing.T) {
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
// 1. SyncEvents first time
err := SyncEvents(context.Background())
require.NoError(t, err)
// Verify event exists in DB
var event model.PushEvent
err = dbConn.Where("event_key = ?", "admin_login").First(&event).Error
require.NoError(t, err)
assert.Equal(t, "管理员登录", event.Name)
assert.False(t, event.Enabled)
// Verify DefaultTemplate matches GORM template field
var defaultMsg NotificationMessage
err = json.Unmarshal([]byte(event.Template), &defaultMsg)
require.NoError(t, err)
assert.Equal(t, adminLoginEvent.DefaultTemplate.Title, defaultMsg.Title)
assert.Equal(t, adminLoginEvent.DefaultTemplate.Content, defaultMsg.Content)
}
func TestEventTrigger(t *testing.T) {
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
// SyncEvents
err := SyncEvents(context.Background())
require.NoError(t, err)
t.Run("trigger disabled event silently ignored", func(t *testing.T) {
body := map[string]any{
"user": map[string]any{"username": "test_admin"},
"ip": "127.0.0.1",
}
DefaultTrigger.Trigger(context.Background(), adminLoginEvent, body)
// Sleep briefly since Trigger runs in goroutine
time.Sleep(50 * time.Millisecond)
// Verify no tasks enqueued in TaskExecution GORM table
var count int64
dbConn.Model(&model.TaskExecution{}).Count(&count)
assert.Equal(t, int64(0), count)
})
t.Run("trigger enabled event enqueues task", func(t *testing.T) {
// Create an enabled custom channel in GORM
customChan := &model.PushChannel{
Name: "mock_channel",
Type: "custom",
URL: "https://webhook.site/trigger",
Other: `{"text": "$content"}`,
Enabled: true,
}
err = dbConn.Create(customChan).Error
require.NoError(t, err)
defer dbConn.Delete(customChan)
// Enable the push event in DB using struct to trigger JSON serializer
var event model.PushEvent
err = dbConn.Where("event_key = ?", "admin_login").First(&event).Error
require.NoError(t, err)
event.Enabled = true
event.Channels = []string{"mock_channel"}
event.Targets = []string{"admin_user"}
err = dbConn.Save(&event).Error
require.NoError(t, err)
repository.DeleteActivePushEventCache(context.Background(), "admin_login")
// Trigger
body := map[string]any{
"user": map[string]any{
"username": "super_admin",
},
"ip": "1.1.1.1",
"time": "2026-06-14 18:00:00",
}
DefaultTrigger.Trigger(context.Background(), adminLoginEvent, body)
// Wait for goroutine execution
time.Sleep(50 * time.Millisecond)
// Verify TaskExecution enqueued record
var execution model.TaskExecution
err = dbConn.Where("task_type = ?", SendNotificationTask).First(&execution).Error
require.NoError(t, err)
// Verify enqueued payload structure
var payload SendPayload
err = json.Unmarshal([]byte(execution.Payload), &payload)
require.NoError(t, err)
assert.Equal(t, "admin_login", payload.EventKey)
assert.Equal(t, "custom", payload.Config.Channel)
assert.Equal(t, "https://webhook.site/trigger", payload.Config.URL)
assert.Equal(t, "admin_user", payload.Target)
assert.Equal(t, "管理员登录提醒", payload.Body.Title)
assert.Contains(t, payload.Body.Content, "super_admin")
assert.Contains(t, payload.Body.Content, "1.1.1.1")
})
t.Run("trigger without user injects virtual system user", func(t *testing.T) {
// Create an enabled custom channel in GORM
customChan := &model.PushChannel{
Name: "mock_channel",
Type: "custom",
URL: "https://webhook.site/trigger",
Other: `{"text": "$content"}`,
Enabled: true,
}
err = dbConn.Create(customChan).Error
require.NoError(t, err)
defer dbConn.Delete(customChan)
// Enable the push event in DB
var event model.PushEvent
err = dbConn.Where("event_key = ?", "admin_login").First(&event).Error
require.NoError(t, err)
// 清理旧任务执行记录
dbConn.Where("task_type = ?", SendNotificationTask).Delete(&model.TaskExecution{})
event.Enabled = true
event.Channels = []string{"mock_channel"}
event.Targets = []string{"user.username"} // 动态目标
err = dbConn.Save(&event).Error
require.NoError(t, err)
repository.DeleteActivePushEventCache(context.Background(), "admin_login")
// Trigger with empty body (simulates cron scheduler triggering)
DefaultTrigger.Trigger(context.Background(), adminLoginEvent, nil)
// Wait for goroutine execution
time.Sleep(50 * time.Millisecond)
// Verify TaskExecution enqueued record
var execution model.TaskExecution
err = dbConn.Where("task_type = ?", SendNotificationTask).First(&execution).Error
require.NoError(t, err)
var payload SendPayload
err = json.Unmarshal([]byte(execution.Payload), &payload)
require.NoError(t, err)
// 检查 payload 是否将 target (user.username) 成功替换为 "system"
assert.Equal(t, "system", payload.Target)
// 检查 payload 中的 Content,应当被替换为 "system" 变量
assert.Contains(t, payload.Body.Content, "system")
})
}
func TestPushHandler(t *testing.T) {
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
mPusher := &mockPusher{}
pkgpush.Register("mock_channel", mPusher)
handler := &PushHandler{}
payload := SendPayload{
EventKey: "admin_login",
Config: pkgpush.Config{
Channel: "mock_channel",
URL: "http://mock-url",
},
Target: "admin_user",
Body: NotificationMessage{
Title: "Structured Alert",
Content: "Hello World",
Level: "WARNING",
Ext: map[string]any{"extra_val": 42},
},
}
payloadBytes, err := json.Marshal(payload)
require.NoError(t, err)
t.Run("validate payload", func(t *testing.T) {
validated, valErr := handler.ValidatePayload(payloadBytes)
require.NoError(t, valErr)
assert.NotEmpty(t, validated)
})
t.Run("execute task successfully", func(t *testing.T) {
res, execErr := handler.Execute(context.Background(), payloadBytes)
require.NoError(t, execErr)
assert.Contains(t, res.Message, "推送成功")
// Verify mock pusher received flattened variables
mPusher.mu.Lock()
assert.Equal(t, "admin_user", mPusher.sentTgt)
assert.Equal(t, "Structured Alert", mPusher.sentBody["title"])
assert.Equal(t, "Hello World", mPusher.sentBody["content"])
assert.Equal(t, "WARNING", mPusher.sentBody["level"])
assert.Equal(t, float64(42), mPusher.sentBody["extra_val"]) // unmarshaled json numbers are float64 by default
mPusher.mu.Unlock()
// Verify PushHistory recorded
var history model.PushHistory
err = dbConn.First(&history).Error
require.NoError(t, err)
assert.Equal(t, "admin_login", history.EventKey)
assert.Equal(t, "mock_channel", history.Channel)
assert.Equal(t, "success", history.Status)
assert.Equal(t, "Structured Alert", history.Title)
})
}
func TestPushRouters(t *testing.T) {
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
r := setupTestRouter(adminUser)
// Sync events to populate db
err := SyncEvents(context.Background())
require.NoError(t, err)
t.Run("list events", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/push/events", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
err = json.Unmarshal(w.Body.Bytes(), &resp)
require.NoError(t, err)
dataBytes, _ := json.Marshal(resp.Data)
var events []model.PushEvent
err = json.Unmarshal(dataBytes, &events)
require.NoError(t, err)
assert.Len(t, events, 1)
assert.Equal(t, "admin_login", events[0].EventKey)
})
t.Run("toggle event status", func(t *testing.T) {
var event model.PushEvent
dbConn.First(&event)
// 1. 未配置任何渠道时开启,应该被拒绝
req, _ := http.NewRequest("POST", "/api/v1/admin/push/events/"+strconv.FormatUint(event.ID, 10)+"/toggle", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
// 2. 为该事件关联渠道后,再切换开启,应当成功
event.Channels = []string{"email"}
dbConn.Save(&event)
repository.DeleteActivePushEventCache(context.Background(), event.EventKey)
req2, _ := http.NewRequest("POST", "/api/v1/admin/push/events/"+strconv.FormatUint(event.ID, 10)+"/toggle", nil)
w2 := httptest.NewRecorder()
r.ServeHTTP(w2, req2)
assert.Equal(t, http.StatusOK, w2.Code)
var updated model.PushEvent
dbConn.First(&updated)
assert.True(t, updated.Enabled)
})
t.Run("update event", func(t *testing.T) {
var event model.PushEvent
dbConn.First(&event)
updateReq := UpdateEventRequest{
Channels: []string{"email"},
Targets: []string{"user@test.com"},
Template: `{"title": "Custom Login Alert", "content": "Alert", "level": "WARNING"}`,
Enabled: true,
}
bodyBytes, _ := json.Marshal(updateReq)
req, _ := http.NewRequest("PUT", "/api/v1/admin/push/events/"+strconv.FormatUint(event.ID, 10), bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var updated model.PushEvent
dbConn.First(&updated)
assert.Equal(t, []string{"email"}, updated.Channels)
assert.Equal(t, []string{"user@test.com"}, updated.Targets)
assert.Contains(t, updated.Template, "Custom Login Alert")
})
t.Run("list push histories", func(t *testing.T) {
// Populate history record
hist := model.PushHistory{
EventKey: "admin_login",
Channel: "email",
Target: "user@test.com",
Title: "Custom Login Alert",
Content: "Alert",
Level: "WARNING",
Status: "success",
}
dbConn.Create(&hist)
req, _ := http.NewRequest("GET", "/api/v1/admin/push/histories?page=1&page_size=10", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataMap, ok := resp.Data.(map[string]any)
assert.True(t, ok)
assert.Equal(t, float64(1), dataMap["total"])
})
t.Run("test push endpoint", func(t *testing.T) {
mPusher := &mockPusher{}
pkgpush.Register("test_channel", mPusher)
testReq := TestPushRequest{
Config: pkgpush.Config{
Channel: "test_channel",
URL: "http://test-url",
},
Target: "test_target",
}
bodyBytes, _ := json.Marshal(testReq)
req, _ := http.NewRequest("POST", "/api/v1/admin/push/test", bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
})
t.Run("list built-in events", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/push/events/builtin", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
err := json.Unmarshal(w.Body.Bytes(), &resp)
require.NoError(t, err)
builtins, ok := resp.Data.([]any)
assert.True(t, ok)
assert.NotEmpty(t, builtins)
})
t.Run("create and delete push event", func(t *testing.T) {
// Clean up any existing admin_login event first
dbConn.Where("event_key = ?", "admin_login").Delete(&model.PushEvent{})
// 1. Create event
createReq := CreateEventRequest{
EventKey: "admin_login",
Channels: []string{"email"},
Targets: []string{"admin@test.com"},
Enabled: true,
}
bodyBytes, _ := json.Marshal(createReq)
req, _ := http.NewRequest("POST", "/api/v1/admin/push/events", bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
// Verify created in DB
var event model.PushEvent
err := dbConn.Where("event_key = ?", "admin_login").First(&event).Error
require.NoError(t, err)
assert.Equal(t, "admin_login", event.EventKey)
assert.Equal(t, "管理员登录", event.Name)
assert.True(t, event.Enabled)
// 2. Try creating again (should fail)
w2 := httptest.NewRecorder()
req2, _ := http.NewRequest("POST", "/api/v1/admin/push/events", bytes.NewBuffer(bodyBytes))
req2.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w2, req2)
assert.Equal(t, http.StatusBadRequest, w2.Code)
// 3. Delete event
w3 := httptest.NewRecorder()
req3, _ := http.NewRequest("DELETE", "/api/v1/admin/push/events/"+strconv.FormatUint(event.ID, 10), nil)
r.ServeHTTP(w3, req3)
assert.Equal(t, http.StatusOK, w3.Code)
// Verify deleted from DB
var count int64
dbConn.Model(&model.PushEvent{}).Where("event_key = ?", "admin_login").Count(&count)
assert.Equal(t, int64(0), count)
})
}
func TestResolveTarget(t *testing.T) {
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
// 1. 创建测试用户与管理员用户
testUser := &model.User{
ID: 9999,
Username: "target_user",
Email: "target@test.com",
IsAdmin: false,
}
err := dbConn.Create(testUser).Error
require.NoError(t, err)
adminUser := &model.User{
ID: 8888,
Username: "admin_user",
Email: "admin@test.com",
IsAdmin: true,
}
err = dbConn.Create(adminUser).Error
require.NoError(t, err)
flatBody := map[string]any{
"user.id": float64(9999), // JSON 反序列化后一般是 float64
"user.username": "target_user",
"user.email": "target@test.com",
}
ctx := context.Background()
t.Run("dynamic user.id resolved and converted for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "user.id", flatBody, "email")
assert.Equal(t, "target@test.com", res)
})
t.Run("dynamic user.username resolved and converted for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "user.username", flatBody, "email")
assert.Equal(t, "target@test.com", res)
})
t.Run("dynamic user.email resolved directly for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "user.email", flatBody, "email")
assert.Equal(t, "target@test.com", res)
})
t.Run("fixed user id resolved and converted for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "9999", flatBody, "email")
assert.Equal(t, "target@test.com", res)
})
t.Run("fixed username resolved and converted for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "target_user", flatBody, "email")
assert.Equal(t, "target@test.com", res)
})
t.Run("fixed email address resolved directly for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "fixed@example.com", flatBody, "email")
assert.Equal(t, "fixed@example.com", res)
})
t.Run("fixed username resolved for non-email channel", func(t *testing.T) {
res := resolveTarget(ctx, "target_user", flatBody, "lark")
assert.Equal(t, "target_user", res)
})
t.Run("non-exist user resolved as fallback", func(t *testing.T) {
res := resolveTarget(ctx, "non_exist_user", flatBody, "email")
assert.Equal(t, "non_exist_user", res)
})
t.Run("system target resolves to admin email for email channel", func(t *testing.T) {
res := resolveTarget(ctx, "系统", flatBody, "email")
assert.Equal(t, "admin@test.com", res)
res2 := resolveTarget(ctx, "system", flatBody, "email")
assert.Equal(t, "admin@test.com", res2)
res3 := resolveTarget(ctx, "0", flatBody, "email")
assert.Equal(t, "admin@test.com", res3)
})
t.Run("system target resolves to admin username for lark channel", func(t *testing.T) {
res := resolveTarget(ctx, "系统", flatBody, "lark")
assert.Equal(t, "admin_user", res)
})
}
func TestPushChannelAPI(t *testing.T) {
// 1. 模型校验测试
t.Run("validate push channel model constraints", func(t *testing.T) {
// 校验名称合法性
c1 := &model.PushChannel{Name: "invalid-name!", URL: "https://hook.com", Other: "{}"}
assert.Error(t, c1.Validate())
// 校验 URL 安全前缀 HTTPS
c2 := &model.PushChannel{Name: "custom_channel", URL: "http://insecure-hook.com", Other: "{}"}
assert.Error(t, c2.Validate())
// 校验 JSON 格式
c3 := &model.PushChannel{Name: "custom_channel", URL: "https://hook.com", Other: "{invalid-json}"}
assert.Error(t, c3.Validate())
// 正确配置
c4 := &model.PushChannel{Name: "custom_channel", URL: "https://hook.com", Other: "{\"content\":\"$content\"}"}
assert.NoError(t, c4.Validate())
// 飞书渠道校验:非 HTTPS 地址报错
c5 := &model.PushChannel{Name: "lark_channel", Type: "lark", URL: "http://open.feishu.cn", Other: ""}
assert.Error(t, c5.Validate())
// 飞书正确配置
c6 := &model.PushChannel{Name: "lark_channel", Type: "lark", URL: "https://open.feishu.cn", Other: ""}
assert.NoError(t, c6.Validate())
// Telegram 渠道校验
cTelegramErr := &model.PushChannel{Name: "tg_channel", Type: "telegram", URL: "https://api.telegram.org", Token: "", Other: ""}
assert.Error(t, cTelegramErr.Validate())
cTelegramErr2 := &model.PushChannel{Name: "tg_channel", Type: "telegram", URL: "http://api.telegram.org", Token: "123:abc", Other: ""}
assert.Error(t, cTelegramErr2.Validate())
cTelegramOk := &model.PushChannel{Name: "tg_channel", Type: "telegram", URL: "", Token: "123:abc", Other: "-100123"}
assert.NoError(t, cTelegramOk.Validate())
assert.Equal(t, "https://api.telegram.org", cTelegramOk.URL)
// 邮件配置校验:允许空配置以复用系统全局设置
c7 := &model.PushChannel{Name: "email_channel", Type: "email", URL: "", Token: "", Other: ""}
assert.NoError(t, c7.Validate())
// 邮件正确配置 (非 HTTPS 协议 URL 允许)
c8 := &model.PushChannel{Name: "email_channel", Type: "email", URL: "smtp.exmail.qq.com:465", Token: "user@example.com", Other: "authcode"}
assert.NoError(t, c8.Validate())
})
// 2. HTTP CRUD & 触发鉴权测试
dbConn, _, cleanup := setupPushTest(t)
defer cleanup()
// 构建路由以进行 HTTP 模拟请求
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
{
adminGroup.GET("/push/channels", ListChannels)
adminGroup.POST("/push/channels", CreateChannel)
adminGroup.PUT("/push/channels/:id", UpdateChannel)
adminGroup.DELETE("/push/channels/:id", DeleteChannel)
adminGroup.POST("/push/channels/test", TestChannel)
}
var createdID uint64
t.Run("admin create channel", func(t *testing.T) {
reqBody := CreateChannelRequest{
Name: "my_custom_channel",
Description: "My custom channel webhook",
Type: "custom",
Token: "my_chan_token",
URL: "https://webhook.site/test",
Other: `{"title": "$title", "body": "$content"}`,
Enabled: true,
}
bodyBytes, _ := json.Marshal(reqBody)
req, _ := http.NewRequest("POST", "/api/v1/admin/push/channels", bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataMap, ok := resp.Data.(map[string]any)
assert.True(t, ok)
assert.Equal(t, "my_custom_channel", dataMap["name"])
createdID = uint64(dataMap["id"].(float64))
})
t.Run("admin list channels", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/push/channels", nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
list, ok := resp.Data.([]any)
assert.True(t, ok)
assert.Len(t, list, 1)
})
t.Run("admin update channel", func(t *testing.T) {
updateReq := UpdateChannelRequest{
Description: "Updated remark",
Type: "custom",
Token: "new_chan_token",
URL: "https://webhook.site/updated",
Other: `{"text": "$content"}`,
Enabled: true,
}
bodyBytes, _ := json.Marshal(updateReq)
req, _ := http.NewRequest("PUT", "/api/v1/admin/push/channels/"+strconv.FormatUint(createdID, 10), bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var updated model.PushChannel
dbConn.First(&updated, createdID)
assert.Equal(t, "Updated remark", updated.Description)
assert.Equal(t, "new_chan_token", updated.Token)
assert.Equal(t, `{"text": "$content"}`, updated.Other)
})
t.Run("admin test channel endpoint", func(t *testing.T) {
testReq := TestChannelRequest{
Name: "my_custom_channel",
Target: "test_target",
}
bodyBytes, _ := json.Marshal(testReq)
req, _ := http.NewRequest("POST", "/api/v1/admin/push/channels/test", bytes.NewBuffer(bodyBytes))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
})
t.Run("admin delete channel", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/push/channels/"+strconv.FormatUint(createdID, 10), nil)
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var count int64
dbConn.Model(&model.PushChannel{}).Where("id = ?", createdID).Count(&count)
assert.Equal(t, int64(0), count)
})
}
+292
View File
@@ -0,0 +1,292 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package push defines push notification HTTP routes.
package push
import (
"context"
"errors"
"fmt"
"net/http"
"strconv"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/pkg/push"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// UpdateEventRequest 更新事件请求参数
type UpdateEventRequest struct {
Channels []string `json:"channels"`
Targets []string `json:"targets"`
Template string `json:"template" binding:"required"`
Enabled bool `json:"enabled"`
}
// TestPushRequest 测试推送通道请求参数
type TestPushRequest struct {
Config push.Config `json:"config" binding:"required"`
Target string `json:"target"`
}
// SyncEvents automatically registers/updates built-in events in the database.
func SyncEvents(ctx context.Context) error {
return syncBuiltInEvents(ctx)
}
// ListEvents 获取通知事件列表
// @Summary 获取所有通知事件
// @Description 返回系统配置的通知事件列表,包括预置和自定义事件,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]model.PushEvent} "通知事件列表"
// @Router /api/v1/admin/push/events [get]
func ListEvents(c *gin.Context) {
ctx := c.Request.Context()
events, err := listPushEvents(ctx)
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(events))
}
// CreateEventRequest 创建事件请求参数
type CreateEventRequest struct {
EventKey string `json:"event_key"`
TaskType string `json:"task_type"` // 关联的异步任务类型
Channels []string `json:"channels"`
Targets []string `json:"targets"`
Template string `json:"template"`
Enabled bool `json:"enabled"`
}
func findBuiltInEvent(key string) (EventMetadata, bool) {
for _, meta := range BuiltInEvents {
if meta.Key == key {
return meta, true
}
}
return EventMetadata{}, false
}
// ListBuiltInEvents 获取内置通知事件列表
// @Summary 获取所有内置通知事件
// @Description 返回系统定义的所有内置通知事件元数据,供前端下拉框选择,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]EventMetadata} "内置通知事件列表"
// @Router /api/v1/admin/push/events/builtin [get]
func ListBuiltInEvents(c *gin.Context) {
c.JSON(http.StatusOK, response.OK(BuiltInEvents))
}
// CreateEvent 创建通知事件
// @Summary 创建通知事件
// @Description 绑定系统内置事件或异步任务、推送渠道、接收目标并创建通知事件配置,需要管理员权限
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body CreateEventRequest true "创建参数"
// @Success 200 {object} response.Any{data=model.PushEvent} "创建成功"
// @Router /api/v1/admin/push/events [post]
func CreateEvent(c *gin.Context) {
var req CreateEventRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
event, err := createPushEvent(c.Request.Context(), req)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(event))
}
// DeleteEvent 删除通知事件配置
// @Summary 删除通知事件配置
// @Description 删除数据库中的特定通知事件配置,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Param id path int true "事件 ID"
// @Success 200 {object} response.Any{data=string} "删除成功"
// @Router /api/v1/admin/push/events/{id} [delete]
func DeleteEvent(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "invalid event id")
return
}
if err := deletePushEvent(c.Request.Context(), id); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, "notification event not found")
return
}
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// UpdateEvent 更新通知事件
// @Summary 更新通知事件
// @Description 更新已有通知事件的推送渠道、接收目标和内容模板,需要管理员权限
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path int true "事件 ID"
// @Param request body push.UpdateEventRequest true "更新参数"
// @Success 200 {object} response.Any{data=string} "修改成功"
// @Router /api/v1/admin/push/events/{id} [put]
func UpdateEvent(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "invalid event id")
return
}
var req UpdateEventRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if err := updatePushEvent(c.Request.Context(), id, req); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, "notification event not found")
return
}
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// ToggleEvent 快捷切换通知事件启用状态
// @Summary 快捷切换通知事件启用状态
// @Description 启用或禁用指定的通知事件
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Param id path int true "事件 ID"
// @Success 200 {object} response.Any{data=string} "切换成功"
// @Router /api/v1/admin/push/events/{id}/toggle [post]
func ToggleEvent(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "invalid event id")
return
}
enabled, err := togglePushEvent(c.Request.Context(), id)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, "notification event not found")
return
}
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(enabled))
}
// pushHistoriesResponse 推送历史分页响应
//
//nolint:unused
type pushHistoriesResponse struct {
Total int64 `json:"total"`
Results []model.PushHistory `json:"results"`
}
// ListHistories 分页获取通知推送历史
// @Summary 分页获取通知推送历史
// @Description 返回分页的通知历史日志数据,需要管理员权限
// @Tags admin-push
// @Produce json
// @Security SessionCookie
// @Param page query int false "当前页码"
// @Param page_size query int false "分页大小"
// @Param event_key query string false "过滤事件名称"
// @Param status query string false "过滤发送状态"
// @Success 200 {object} response.Any{data=pushHistoriesResponse} "推送历史列表"
// @Router /api/v1/admin/push/histories [get]
func ListHistories(c *gin.Context) {
page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20"))
if page < 1 {
page = 1
}
if pageSize < 1 {
pageSize = 20
}
total, results, err := listPushHistories(c.Request.Context(), repository.PushHistoryListFilter{
EventKey: c.Query("event_key"),
Status: c.Query("status"),
Page: page,
PageSize: pageSize,
})
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(map[string]any{
"total": total,
"results": results,
}))
}
// TestPush 测试推送通道发送
// @Summary 测试推送通道发送
// @Description 接收临时通知渠道配置并在本地同步调用 Pusher.Send 发送测试消息
// @Tags admin-push
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body push.TestPushRequest true "测试请求体"
// @Success 200 {object} response.Any{data=string} "测试成功"
// @Router /api/v1/admin/push/test [post]
func TestPush(c *gin.Context) {
var req TestPushRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
pusher, err := push.GetPusher(req.Config.Channel)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if err := pusher.ValidateConfig(req.Config); err != nil {
response.AbortBadRequest(c, fmt.Sprintf("validation failed: %v", err))
return
}
applySMTPFallbackToPushConfig(c.Request.Context(), &req.Config)
testBody := map[string]any{
keyTitle: "测试通道推送",
keyContent: "当您收到这条消息,说明当前渠道连通性测试通过。",
keyLevel: defaultLevelInfo,
}
if err := pusher.Send(c.Request.Context(), req.Config, req.Target, testBody, "", nil); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
+124
View File
@@ -0,0 +1,124 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package push
import (
"context"
"encoding/json"
"strconv"
"time"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/pkg/logger"
)
// RegisterTaskListeners subscribes push notification handlers to task completion events.
func RegisterTaskListeners() {
task.OnTaskCompleted(handleTaskCompleted)
}
func handleTaskCompleted(ctx context.Context, execution *model.TaskExecution, result *task.TaskResult, execErr error) {
events, err := listActivePushEventsByTaskType(ctx, execution.TaskType)
if err != nil {
logger.ErrorF(ctx, "push_task_completed_listener: failed to query push events for task type %s: %v", execution.TaskType, err)
return
}
if len(events) == 0 {
return
}
body := map[string]any{
"task_id": execution.TaskID,
"task_name": execution.TaskName,
"task_type": execution.TaskType,
"task_status": string(execution.Status),
"task_duration": execution.Duration,
"time": time.Now().Format("2006-01-02 15:04:05"),
}
if execErr != nil {
body["task_error"] = execErr.Error()
} else {
body["task_error"] = ""
}
if result != nil {
body["task_result"] = result.Message
} else {
body["task_result"] = ""
}
var payloadMap map[string]any
if execution.Payload != "" {
if err := json.Unmarshal([]byte(execution.Payload), &payloadMap); err == nil {
body["payload"] = payloadMap
extractUserFromMap(ctx, payloadMap, body)
}
}
if result != nil && result.Detail != "" {
var detailMap map[string]any
if err := json.Unmarshal([]byte(result.Detail), &detailMap); err == nil {
body["detail"] = detailMap
extractUserFromMap(ctx, detailMap, body)
}
}
for _, event := range events {
meta := EventMetadata{
Key: event.EventKey,
Name: event.Name,
Description: "异步任务执行完毕触发的自动通知",
}
DefaultTrigger.Trigger(ctx, meta, body)
}
}
func extractUserFromMap(ctx context.Context, data map[string]any, body map[string]any) {
if u, exists := body["user"]; exists && u != nil {
return
}
if user := loadUserFromPayload(ctx, data); user != nil {
body["user"] = user
}
}
func extractUserID(data map[string]any) (uint64, bool) {
for _, k := range []string{"user_id", "userId", "uid"} {
val, ok := data[k]
if !ok || val == nil {
continue
}
switch v := val.(type) {
case float64:
if v >= 0 {
return uint64(v), true
}
case int:
if v >= 0 {
return uint64(v), true
}
case int64:
if v >= 0 {
return uint64(v), true
}
case uint64:
return v, true
case string:
if id, err := strconv.ParseUint(v, 10, 64); err == nil {
return id, true
}
}
}
return 0, false
}
func extractUsername(data map[string]any) string {
for _, k := range []string{"username", "user_name"} {
if val, ok := data[k]; ok && val != nil {
if s, ok := val.(string); ok && s != "" {
return s
}
}
}
return ""
}
+121
View File
@@ -0,0 +1,121 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package push defines push notification HTTP routes and background tasks.
package push
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/pkg/push"
)
const (
// SendNotificationTask 发送推送通知任务标识
SendNotificationTask = "push:send"
// TaskTypeSendNotification 推送通知管理类型
TaskTypeSendNotification = "send_notification"
)
// SendNotificationMeta represents the task metadata.
var SendNotificationMeta = task.TaskMeta{
Type: TaskTypeSendNotification,
AsynqTask: SendNotificationTask,
Name: "推送通知",
Description: "异步执行系统通知的多渠道派发与推送",
SupportsTime: false,
MaxRetry: task.DefaultMaxRetry,
Queue: task.QueueDefault,
Retryable: true,
Params: []task.TaskParam{
{
Name: "event_key",
Label: "事件标识",
Type: "string",
Required: true,
Placeholder: "admin_login",
},
{
Name: "target",
Label: "目标接收者",
Type: "string",
Required: false,
},
},
}
// PushHandler 通知推送异步任务处理器
//
//nolint:revive
type PushHandler struct{}
// ValidatePayload 校验并标准化推送参数
func (h *PushHandler) ValidatePayload(payload []byte) ([]byte, error) {
if len(payload) == 0 {
return nil, errors.New("payload is required")
}
var req SendPayload
if err := json.Unmarshal(payload, &req); err != nil {
return nil, fmt.Errorf("invalid json format: %w", err)
}
if req.Config.Channel == "" {
return nil, errors.New("channel type is required")
}
return json.Marshal(req)
}
// Execute 异步执行推送操作并记录推送历史审计
func (h *PushHandler) Execute(ctx context.Context, payload []byte) (*task.TaskResult, error) {
var req SendPayload
if err := json.Unmarshal(payload, &req); err != nil {
task.AppendLog(ctx, "解析推送参数失败: %v", err)
return nil, fmt.Errorf("parse payload failed: %w", err)
}
task.AppendLog(ctx, "开始推送通知: 事件 = %s, 渠道 = %s, 接收目标 = %s", req.EventKey, req.Config.Channel, req.Target)
pusher, err := push.GetPusher(req.Config.Channel)
if err != nil {
errWrap := fmt.Errorf("get pusher failed: %w", err)
task.AppendLog(ctx, "推送失败: %v", errWrap)
if task.IsFinalAttempt(ctx) {
h.recordHistory(ctx, req, "failed", errWrap.Error())
}
return nil, errWrap
}
// 执行真正的消息推送,扁平化为原始 json 格式
flatBody := req.Body.Flatten()
err = pusher.Send(ctx, req.Config, req.Target, flatBody, req.Template, nil)
title := req.Body.Title
content := req.Body.Content
if err != nil {
task.AppendLog(ctx, "消息推送失败 (标题: %s): %v", title, err)
if task.IsFinalAttempt(ctx) {
h.recordHistory(ctx, req, "failed", err.Error())
}
return nil, fmt.Errorf("pusher.Send failed: %w", err)
}
task.AppendLog(ctx, "消息推送成功 (标题: %s, 内容摘要: %s)", title, content)
h.recordHistory(ctx, req, "success", "")
return &task.TaskResult{
Message: fmt.Sprintf("推送成功: [%s] -> %s", req.Config.Channel, req.Target),
}, nil
}
func (h *PushHandler) recordHistory(ctx context.Context, req SendPayload, status string, errMsg string) {
if dbErr := recordPushHistory(ctx, req, status, errMsg); dbErr != nil {
task.AppendLog(ctx, "写入推送历史审计记录失败: %v", dbErr)
}
}
+355
View File
@@ -0,0 +1,355 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package status 提供系统状态查询接口
package status
import (
"context"
"fmt"
"log"
"math"
"net/http"
"os"
"os/exec"
"runtime"
"time"
"github.com/Rain-kl/Wavelet/internal/config"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// startTime 记录服务启动时间
var startTime = time.Now()
const (
hoursInDay = 24
minutesInHour = 60
secondsInMinute = 60
nanosPerSecond = 1e9
binaryKB = 0
binaryMB = 1
binaryGB = 2
valueThreshold = 10 // 格式化时区分整数显示的阈值
)
// SystemStatusResponse 系统状态响应结构体
type SystemStatusResponse struct {
Uptime string `json:"uptime"`
NumGoroutine int `json:"num_goroutine"`
Alloc string `json:"alloc"`
TotalAlloc string `json:"total_alloc"`
Sys string `json:"sys"`
Lookups uint64 `json:"lookups"`
Mallocs uint64 `json:"mallocs"`
Frees uint64 `json:"frees"`
HeapAlloc string `json:"heap_alloc"`
HeapSys string `json:"heap_sys"`
HeapIdle string `json:"heap_idle"`
HeapInuse string `json:"heap_inuse"`
HeapReleased string `json:"heap_released"`
HeapObjects uint64 `json:"heap_objects"`
StackInuse string `json:"stack_inuse"`
StackSys string `json:"stack_sys"`
MSpanInuse string `json:"mspan_inuse"`
MSpanSys string `json:"mspan_sys"`
MCacheInuse string `json:"mcache_inuse"`
MCacheSys string `json:"mcache_sys"`
BuckHashSys string `json:"buck_hash_sys"`
GCSys string `json:"gc_sys"`
OtherSys string `json:"other_sys"`
NextGC string `json:"next_gc"`
LastGCTime string `json:"last_gc_time"`
PauseTotalNs string `json:"pause_total_ns"`
LastPause string `json:"last_pause"`
NumGC uint32 `json:"num_gc"`
}
// formatBytes 格式化字节大小
func formatBytes(bytes uint64) string {
const unit = 1024
if bytes < unit {
return fmt.Sprintf("%d B", bytes)
}
div, exp := int64(unit), 0
for n := bytes / unit; n >= unit; n /= unit {
div *= unit
exp++
}
value := float64(bytes) / float64(div)
var suffix string
switch exp {
case binaryKB:
suffix = "KiB"
case binaryMB:
suffix = "MiB"
case binaryGB:
suffix = "GiB"
default:
suffix = "TiB"
}
// 格式化规则:
// - 如果是整数(如 16, 73, 105, 986, 112):
// - 如果 >= 10,则格式化为 "%.0f" (e.g. "16 KiB")
// - 如果 < 10,则格式化为 "%.1f" (e.g. "9.0 KiB")
// - 如果不是整数(如 5.8, 9.1, 7.6, 4.8):格式化为 "%.1f"
if value == math.Trunc(value) {
if value >= valueThreshold {
return fmt.Sprintf("%.0f %s", value, suffix)
}
return fmt.Sprintf("%.1f %s", value, suffix)
}
return fmt.Sprintf("%.1f %s", value, suffix)
}
// formatDuration 格式化时间持续时间
func formatDuration(d time.Duration) string {
days := int(d.Hours()) / hoursInDay
hours := int(d.Hours()) % hoursInDay
minutes := int(d.Minutes()) % minutesInHour
seconds := int(d.Seconds()) % secondsInMinute
var res string
if days > 0 {
res += fmt.Sprintf("%d天", days)
}
if hours > 0 {
res += fmt.Sprintf("%d小时", hours)
}
if minutes > 0 {
res += fmt.Sprintf("%d分钟", minutes)
}
if seconds > 0 || res == "" {
res += fmt.Sprintf("%d秒钟", seconds)
}
return res
}
// GetSystemStatus 获取系统状态信息
// @Summary 获取系统状态信息
// @Description 获取后端服务运行状态、Goroutine、内存指标等详细统计数据,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=status.SystemStatusResponse} "获取成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/status [get]
func GetSystemStatus(c *gin.Context) {
var m runtime.MemStats
runtime.ReadMemStats(&m)
uptime := formatDuration(time.Since(startTime))
numGoroutine := runtime.NumGoroutine()
var lastGCTime string
switch {
case m.LastGC > 0 && m.LastGC <= math.MaxInt64:
lastGCTime = formatDuration(time.Since(time.Unix(0, int64(m.LastGC))))
case m.LastGC > 0:
lastGCTime = "未知"
default:
lastGCTime = "无"
}
var lastPause string
if m.NumGC > 0 {
lastPause = fmt.Sprintf("%.3fs", float64(m.PauseNs[(m.NumGC-1)%256])/nanosPerSecond)
} else {
lastPause = "0.000s"
}
res := SystemStatusResponse{
Uptime: uptime,
NumGoroutine: numGoroutine,
Alloc: formatBytes(m.Alloc),
TotalAlloc: formatBytes(m.TotalAlloc),
Sys: formatBytes(m.Sys),
Lookups: m.Lookups,
Mallocs: m.Mallocs,
Frees: m.Frees,
HeapAlloc: formatBytes(m.HeapAlloc),
HeapSys: formatBytes(m.HeapSys),
HeapIdle: formatBytes(m.HeapIdle),
HeapInuse: formatBytes(m.HeapInuse),
HeapReleased: formatBytes(m.HeapReleased),
HeapObjects: m.HeapObjects,
StackInuse: formatBytes(m.StackInuse),
StackSys: formatBytes(m.StackSys),
MSpanInuse: formatBytes(m.MSpanInuse),
MSpanSys: formatBytes(m.MSpanSys),
MCacheInuse: formatBytes(m.MCacheInuse),
MCacheSys: formatBytes(m.MCacheSys),
BuckHashSys: formatBytes(m.BuckHashSys),
GCSys: formatBytes(m.GCSys),
OtherSys: formatBytes(m.OtherSys),
NextGC: formatBytes(m.NextGC),
LastGCTime: lastGCTime,
PauseTotalNs: fmt.Sprintf("%.1fs", float64(m.PauseTotalNs)/nanosPerSecond),
LastPause: lastPause,
NumGC: m.NumGC,
}
c.JSON(http.StatusOK, response.OK(res))
}
// DatabaseInfoResponse 数据库信息响应结构体
type DatabaseInfoResponse struct {
Type string `json:"type"`
Name string `json:"name"`
Version string `json:"version"`
}
// getSQLiteInfo 返回 SQLite 数据库信息
func getSQLiteInfo(ctx context.Context) DatabaseInfoResponse {
info := DatabaseInfoResponse{
Type: "sqlite",
Name: config.Config.Database.SQLitePath,
Version: "SQLite",
}
if info.Name == "" {
info.Name = "./data/openflare.db"
}
gormDB := db.DB(ctx)
if gormDB == nil {
return info
}
var ver string
if err := gormDB.Raw("SELECT sqlite_version()").Scan(&ver).Error; err == nil && ver != "" {
info.Version = "SQLite " + ver
}
return info
}
// getPostgresInfo 返回 PostgreSQL 数据库信息
func getPostgresInfo(ctx context.Context) DatabaseInfoResponse {
info := DatabaseInfoResponse{
Type: "postgres",
Name: config.Config.Database.Database,
Version: "PostgreSQL",
}
gormDB := db.DB(ctx)
if gormDB == nil {
return info
}
var ver string
if err := gormDB.Raw("SELECT version()").Scan(&ver).Error; err == nil && ver != "" {
info.Version = ver
}
return info
}
// GetDatabaseInfo 获取当前数据库类型及版本信息
// @Summary 获取数据库信息
// @Description 返回当前使用的数据库类型(sqlite/postgres)、名称/路径及版本字符串,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=status.DatabaseInfoResponse} "获取成功"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/db-info [get]
func GetDatabaseInfo(c *gin.Context) {
var info DatabaseInfoResponse
if !config.Config.Database.Enabled {
info = getSQLiteInfo(c.Request.Context())
} else {
info = getPostgresInfo(c.Request.Context())
}
c.JSON(http.StatusOK, response.OK(info))
}
// ExportDatabase 导出数据库
// @Summary 导出数据库
// @Description SQLite 时直接下载 .db 文件;PostgreSQL 时执行 pg_dump 并流式下载 .sql 文件,需要管理员权限
// @Tags admin
// @Produce application/octet-stream
// @Security SessionCookie
// @Success 200 {file} binary "数据库文件"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "导出失败"
// @Router /api/v1/admin/db-export [get]
func ExportDatabase(c *gin.Context) {
if !config.Config.Database.Enabled {
exportSQLite(c)
} else {
exportPostgres(c)
}
}
// exportSQLite 以 HTTP 附件方式下载 SQLite .db 文件
func exportSQLite(c *gin.Context) {
path := config.Config.Database.SQLitePath
if path == "" {
path = "./data/openflare.db"
}
f, err := os.Open(path) //nolint:gosec // path is loaded from server startup configuration, not user input
if err != nil {
response.AbortInternal(c, "无法打开数据库文件: "+err.Error())
return
}
defer func() {
if closeErr := f.Close(); closeErr != nil {
_ = closeErr
}
}()
fi, err := f.Stat()
if err != nil {
response.AbortInternal(c, "无法读取数据库文件信息: "+err.Error())
return
}
c.Header("Content-Disposition", `attachment; filename="openflare.db"`)
c.Header("Content-Type", "application/octet-stream")
c.Header("Content-Length", fmt.Sprintf("%d", fi.Size()))
c.Status(http.StatusOK)
http.ServeContent(c.Writer, c.Request, "openflare.db", fi.ModTime(), f)
}
// exportPostgres 执行 pg_dump 并将输出流式传输给客户端
func exportPostgres(c *gin.Context) {
dbCfg := config.Config.Database
// 检查 pg_dump 是否可用
pgDumpPath, err := exec.LookPath("pg_dump")
if err != nil {
response.AbortInternal(c, "pg_dump 不可用,请确保服务器已安装 PostgreSQL 客户端工具")
return
}
args := []string{
"--no-password",
"-h", dbCfg.Host,
"-p", fmt.Sprintf("%d", dbCfg.Port),
"-U", dbCfg.Username,
dbCfg.Database,
}
cmd := exec.CommandContext(c.Request.Context(), pgDumpPath, args...) //nolint:gosec // pgDumpPath is a looked up command path, args are from database configuration
if dbCfg.Password != "" {
cmd.Env = append(os.Environ(), "PGPASSWORD="+dbCfg.Password)
} else {
cmd.Env = os.Environ()
}
fileName := fmt.Sprintf("openflare_%s.sql", time.Now().Format("20060102_150405"))
c.Header("Content-Disposition", `attachment; filename="`+fileName+`"`)
c.Header("Content-Type", "application/octet-stream")
c.Status(http.StatusOK)
cmd.Stdout = c.Writer
cmd.Stderr = nil // 忽略 stderr 以避免污染输出流
if err := cmd.Run(); err != nil {
// 响应头已发出,无法再写 JSON 错误,记录到服务器日志
log.Printf("[db-export] pg_dump failed: %v\n", err)
}
}
+15
View File
@@ -0,0 +1,15 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package system_config 提供系统配置管理功能
package system_config
// 系统配置错误消息常量
const (
SystemConfigNotFound = "系统配置不存在"
ConfigKeyRequired = "配置键不能为空"
ConfigValueRequired = "配置值不能为空"
ConfigKeyExists = "配置键已存在"
StorageDriverSwitchRequiresMigration = "存在存量文件,请通过存储迁移任务切换存储引擎"
)
+128
View File
@@ -0,0 +1,128 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package system_config
import (
"context"
"encoding/json"
"errors"
"time"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/storage"
"github.com/Rain-kl/Wavelet/pkg/logger"
"gorm.io/gorm"
)
func createSystemConfig(ctx context.Context, req CreateSystemConfigRequest) error {
exists, err := repository.SystemConfigExists(ctx, req.Key)
if err != nil {
return err
}
if exists {
return errors.New(ConfigKeyExists)
}
config := model.SystemConfig{
Key: req.Key,
Value: req.Value,
Type: req.Type,
Visibility: req.Visibility,
Description: req.Description,
}
if err := repository.CreateSystemConfig(ctx, &config); err != nil {
return err
}
invalidateSystemConfigCaches(ctx, req.Key)
if err := repository.InvalidateVisibleSystemConfigsCache(ctx); err != nil {
logger.WarnF(ctx, "清理公共配置列表缓存失败: %v", err)
}
return nil
}
func listSystemConfigs(ctx context.Context, configType string) ([]model.SystemConfig, error) {
return repository.ListAdminSystemConfigs(ctx, configType)
}
func getSystemConfig(ctx context.Context, key string) (model.SystemConfig, error) {
return repository.GetAdminSystemConfigByKey(ctx, key)
}
func updateSystemConfig(ctx context.Context, key string, req UpdateSystemConfigRequest) error {
config, err := repository.GetAdminSystemConfigByKey(ctx, key)
if err != nil {
return err
}
var originalDriver storage.Driver
if key == model.ConfigKeyStorageConfig {
var currentCfg storage.Config
if err := json.Unmarshal([]byte(config.Value), &currentCfg); err == nil {
originalDriver = currentCfg.Driver
}
validatedVal, err := validateAndMergeStorageConfig(ctx, req.Value, config.Value)
if err != nil {
return err
}
req.Value = validatedVal
}
if err := db.DB(ctx).Transaction(func(tx *gorm.DB) error {
updates := map[string]any{
"description": req.Description,
}
if req.Visibility != nil {
updates["visibility"] = *req.Visibility
config.Visibility = *req.Visibility
}
if key != model.ConfigKeySMTPPassword || req.Value != maskedConfigValue {
updates["value"] = req.Value
config.Value = req.Value
}
if err := tx.Model(&config).Updates(updates).Error; err != nil {
return err
}
resolveStorageMigrationTasksOnDirectDriverUpdate(ctx, tx, key, originalDriver, req.Value)
return nil
}); err != nil {
return err
}
invalidateCachesAfterConfigUpdate(ctx, key)
return nil
}
func resolveStorageMigrationTasksOnDirectDriverUpdate(
ctx context.Context,
tx *gorm.DB,
key string,
originalDriver storage.Driver,
newValue string,
) {
if key != model.ConfigKeyStorageConfig || originalDriver == "" {
return
}
var newCfg storage.Config
if err := json.Unmarshal([]byte(newValue), &newCfg); err != nil {
return
}
if newCfg.Driver != originalDriver {
return
}
if err := tx.Model(&model.TaskExecution{}).
Where("task_type = ? AND status = ?", "storage:migrate", model.TaskExecutionStatusFailed).
Updates(map[string]any{
"status": model.TaskExecutionStatusSucceeded,
"result": "存储配置直接更新,故障迁移任务自动标记为已解决",
"finished_at": time.Now(),
}).Error; err != nil {
logger.ErrorF(ctx, "自动更新迁移任务状态失败: %v", err)
}
}
@@ -0,0 +1,365 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package system_config
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/apps/cap"
"github.com/Rain-kl/Wavelet/internal/apps/upload"
"github.com/Rain-kl/Wavelet/internal/common/response"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/storage"
"github.com/Rain-kl/Wavelet/pkg/logger"
mail "github.com/Rain-kl/Wavelet/pkg/mail"
)
const maskedConfigValue = "******"
// CreateSystemConfigRequest 创建系统配置请求
type CreateSystemConfigRequest struct {
Key string `json:"key" binding:"required,max=64"`
Value string `json:"value" binding:"required"`
Type string `json:"type" binding:"required,oneof=system business"`
Visibility int `json:"visibility" binding:"oneof=0 1"`
Description string `json:"description" binding:"max=255"`
}
// UpdateSystemConfigRequest 更新系统配置请求
type UpdateSystemConfigRequest struct {
Value string `json:"value" binding:"required"`
Visibility *int `json:"visibility" binding:"omitempty,oneof=0 1"`
Description string `json:"description" binding:"max=255"`
}
// CreateSystemConfig 创建系统配置
// @Summary 创建系统配置
// @Description 创建一条新的系统配置项,配置键不可重复,同时将新配置同步到 Redis,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body system_config.CreateSystemConfigRequest true "创建请求参数"
// @Success 200 {object} response.Any{data=string} "创建成功"
// @Failure 400 {object} response.Any "参数错误或配置键已存在"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/system-configs [post]
func CreateSystemConfig(c *gin.Context) {
var req CreateSystemConfigRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if err := createSystemConfig(c.Request.Context(), req); err != nil {
if err.Error() == ConfigKeyExists {
response.AbortBadRequest(c, ConfigKeyExists)
return
}
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// ListSystemConfigs 获取系统配置列表
// @Summary 获取系统配置列表
// @Description 返回所有系统配置列表,支持按配置类型(system/business)过滤,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param type query string false "配置类型(system/business)"
// @Success 200 {object} response.Any{data=[]model.SystemConfig} "系统配置列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/system-configs [get]
func ListSystemConfigs(c *gin.Context) {
configs, err := listSystemConfigs(c.Request.Context(), c.Query("type"))
if err != nil {
response.AbortInternal(c, err.Error())
return
}
for i := range configs {
configs[i].Value = maskSensitiveConfig(configs[i].Key, configs[i].Value)
}
c.JSON(http.StatusOK, response.OK(configs))
}
// GetSystemConfig 获取单个系统配置
// @Summary 获取单个系统配置
// @Description 根据配置键获取对应的系统配置详情,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param key path string true "配置键"
// @Success 200 {object} response.Any{data=model.SystemConfig} "系统配置详情"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "配置不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/system-configs/{key} [get]
func GetSystemConfig(c *gin.Context) {
config, err := getSystemConfig(c.Request.Context(), c.Param("key"))
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, SystemConfigNotFound)
} else {
response.AbortInternal(c, err.Error())
}
return
}
config.Value = maskSensitiveConfig(config.Key, config.Value)
c.JSON(http.StatusOK, response.OK(config))
}
// UpdateSystemConfig 更新系统配置
// @Summary 更新系统配置
// @Description 根据配置键更新对应的配置内容,同时将更新同步到 Redis,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param key path string true "配置键"
// @Param request body system_config.UpdateSystemConfigRequest true "更新请求参数"
// @Success 200 {object} response.Any{data=string} "更新成功"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "配置不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/system-configs/{key} [put]
func UpdateSystemConfig(c *gin.Context) {
var req UpdateSystemConfigRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
key := c.Param("key")
if err := updateSystemConfig(c.Request.Context(), key, req); err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, SystemConfigNotFound)
return
}
if isStorageConfigValidationError(err) {
response.AbortBadRequest(c, err.Error())
return
}
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OKNil())
}
func invalidateSystemConfigCaches(ctx context.Context, key string) {
if err := repository.InvalidateSystemConfigCache(ctx, key); err != nil {
logger.WarnF(ctx, "清理系统配置缓存失败: %v", err)
}
if cap.IsRuntimeConfigKey(key) {
cap.InvalidateRuntimeSettings()
}
}
func invalidateCachesAfterConfigUpdate(ctx context.Context, key string) {
invalidateSystemConfigCaches(ctx, key)
if key == model.ConfigKeyStorageConfig {
upload.ResetAccessCaches()
upload.PublishAccessCacheInvalidation(ctx)
storage.ResetCache()
storage.PublishCacheInvalidation(ctx)
}
if key == model.ConfigKeyFileAccessWhitelist {
upload.ResetAccessCaches()
upload.PublishAccessCacheInvalidation(ctx)
}
if err := repository.InvalidateVisibleSystemConfigsCache(ctx); err != nil {
logger.WarnF(ctx, "清理公共配置列表缓存失败: %v", err)
}
}
// TestSMTPRequest 测试 SMTP 配置请求
type TestSMTPRequest struct {
SMTPHost string `json:"smtp_host" binding:"required,max=255"`
SMTPPort int `json:"smtp_port" binding:"required"`
SMTPUsername string `json:"smtp_username" binding:"required,max=255"`
SMTPPassword string `json:"smtp_password" binding:"required,max=255"`
To string `json:"to" binding:"required,email"`
}
// TestSMTPResponse 测试 SMTP 配置响应
type TestSMTPResponse struct {
Success bool `json:"success"`
Log string `json:"log"`
Error string `json:"error"`
}
// TestSMTP 测试 SMTP 邮件发送
// @Summary 测试 SMTP 邮件发送
// @Description 使用传入的配置进行 SMTP 邮件发送测试,支持使用 ****** 占位符使用保存的数据库密码
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body system_config.TestSMTPRequest true "测试请求参数"
// @Success 200 {object} response.Any{data=system_config.TestSMTPResponse} "测试执行完毕"
// @Failure 400 {object} response.Any "参数错误"
// @Router /api/v1/admin/system-configs/smtp/test [post]
func TestSMTP(c *gin.Context) {
var req TestSMTPRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
password := req.SMTPPassword
if password == maskedConfigValue {
if sc, err := repository.GetSystemConfigByKey(c.Request.Context(), model.ConfigKeySMTPPassword); err == nil {
password = sc.Value
}
}
cfg := mail.Config{
Host: req.SMTPHost,
Port: req.SMTPPort,
Username: req.SMTPUsername,
Password: password,
}
subject := "OpenFlare SMTP Test Mail"
body := `<h3>SMTP Mail Connection Test</h3>
<p>If you received this message, your SMTP configuration is correct and mail sending is working properly.</p>
<p>Sent from OpenFlare.</p>`
logs, err := mail.SendMailWithLog(c.Request.Context(), cfg, req.To, subject, body)
resp := TestSMTPResponse{
Success: err == nil,
Log: logs,
}
if err != nil {
resp.Error = err.Error()
}
c.JSON(http.StatusOK, response.OK(resp))
}
func isStorageConfigValidationError(err error) bool {
msg := err.Error()
return msg == StorageDriverSwitchRequiresMigration ||
strings.HasPrefix(msg, "解析") ||
strings.HasPrefix(msg, "验证") ||
strings.HasPrefix(msg, "初始化测试") ||
strings.HasPrefix(msg, "存储连通性") ||
strings.HasPrefix(msg, "序列化") ||
strings.HasPrefix(msg, "检查存量文件")
}
func maskSensitiveConfig(key, value string) string {
if value == "" {
return value
}
switch key {
case model.ConfigKeySMTPPassword:
return maskedConfigValue
case model.ConfigKeyStorageConfig:
var cfg storage.Config
if err := json.Unmarshal([]byte(value), &cfg); err == nil {
masked := storage.MaskSecrets(cfg)
if val, err := json.Marshal(masked); err == nil {
return string(val)
}
}
}
return value
}
// validateAndMergeStorageConfig parses, merges unmasked secrets, validates parameter values,
// and tests connectivity of the new storage configuration.
func validateAndMergeStorageConfig(ctx context.Context, value string, currentConfig string) (string, error) {
var currentCfg storage.Config
if err := json.Unmarshal([]byte(currentConfig), &currentCfg); err != nil {
return "", fmt.Errorf("解析当前存储配置失败: %w", err)
}
var newCfg storage.Config
if err := json.Unmarshal([]byte(value), &newCfg); err != nil {
return "", fmt.Errorf("解析目标存储配置失败: %w", err)
}
// 合并被掩码屏蔽的敏感信息,获取完整的真实配置
targetCfg := storage.MergeMaskedSecrets(newCfg, currentCfg)
if err := validateMergedStorageConfig(ctx, currentCfg, newCfg, targetCfg); err != nil {
return "", err
}
// 序列化为最终保存的真实明文配置,防止保存屏蔽的 ****** 字符
unmaskedVal, err := json.Marshal(targetCfg)
if err != nil {
return "", fmt.Errorf("序列化存储配置失败: %w", err)
}
return string(unmaskedVal), nil
}
func validateMergedStorageConfig(ctx context.Context, currentCfg, newCfg, targetCfg storage.Config) error {
if newCfg.Driver != "" && newCfg.Driver != currentCfg.Driver {
var uploadCount int64
if err := db.DB(ctx).Model(&model.Upload{}).
Where("status != ?", model.UploadStatusDeleted).
Count(&uploadCount).Error; err != nil {
return fmt.Errorf("检查存量文件失败: %w", err)
}
if uploadCount > 0 {
return errors.New(StorageDriverSwitchRequiresMigration)
}
if err := validateDriverConfig(targetCfg, newCfg.Driver); err != nil {
return fmt.Errorf("验证目标存储配置参数失败: %w", err)
}
pendingCfg := targetCfg
pendingCfg.Driver = newCfg.Driver
return testStorageBackend(ctx, pendingCfg, newCfg.Driver)
}
if err := storage.ValidateConfig(targetCfg); err != nil {
return fmt.Errorf("验证存储配置参数失败: %w", err)
}
return testStorageBackend(ctx, targetCfg, targetCfg.Driver)
}
func validateDriverConfig(cfg storage.Config, driver storage.Driver) error {
cfg.Driver = driver
return storage.ValidateConfig(cfg)
}
func testStorageBackend(ctx context.Context, cfg storage.Config, driver storage.Driver) error {
testBackend, err := storage.NewBackend(ctx, cfg, driver)
if err != nil {
return fmt.Errorf("初始化测试存储实例失败: %w", err)
}
if err := testBackend.Test(ctx); err != nil {
return fmt.Errorf("存储连通性测试失败: %w", err)
}
return nil
}
@@ -0,0 +1,543 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package system_config
import (
"bufio"
"bytes"
"context"
"encoding/json"
"net"
"net/http"
"net/http/httptest"
"net/textproto"
"strings"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/db"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/internal/storage"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
const expectedDefaultConfigsCount = 30
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
// Mock authentication middleware
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, oauth.UserObjKey, authUser)
}
c.Next()
})
adminGroup.POST("/system-configs", CreateSystemConfig)
adminGroup.GET("/system-configs", ListSystemConfigs)
systemConfigRouter := adminGroup.Group("/system-configs/:key")
{
systemConfigRouter.GET("", GetSystemConfig)
systemConfigRouter.PUT("", UpdateSystemConfig)
}
return r
}
func TestCreateSystemConfig(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("create successfully", func(t *testing.T) {
payload := CreateSystemConfigRequest{
Key: "custom_key",
Value: "custom_value",
Type: "system",
Visibility: model.ConfigVisibilityVisible,
Description: "desc",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/system-configs", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
// Verify database
var cfg model.SystemConfig
err := dbConn.Where("key = ?", "custom_key").First(&cfg).Error
if err != nil {
t.Fatalf("failed to find system config in DB: %v", err)
}
// Verify caches are invalidated after create and repopulate on read
_, err = db.Redis.HGet(
context.Background(),
db.PrefixedKey(repository.SystemConfigRedisHashKey),
"custom_key",
).Result()
if err == nil {
t.Fatal("expected redis cache miss immediately after create")
}
loaded, err := repository.GetSystemConfigByKey(context.Background(), "custom_key")
if err != nil {
t.Fatalf("GetSystemConfigByKey(custom_key) error = %v", err)
}
if loaded.Value != "custom_value" {
t.Errorf("GetSystemConfigByKey(custom_key).Value = %q, want %q", loaded.Value, "custom_value")
}
if loaded.Visibility != model.ConfigVisibilityVisible {
t.Errorf("GetSystemConfigByKey(custom_key).Visibility = %d, want %d", loaded.Visibility, model.ConfigVisibilityVisible)
}
})
t.Run("create duplicate key error", func(t *testing.T) {
// Key "custom_key" already exists from previous test
payload := CreateSystemConfigRequest{
Key: "custom_key",
Value: "another_value",
Type: "system",
Description: "desc",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/system-configs", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request on duplicate key, got %d", w.Code)
}
})
}
func TestListSystemConfigs(t *testing.T) {
_, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("list all seeded configurations", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/system-configs", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d", w.Code)
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var configs []model.SystemConfig
_ = json.Unmarshal(dataBytes, &configs)
// Defaults seed configurations
if len(configs) != expectedDefaultConfigsCount {
t.Errorf("expected %d default configs, got %d", expectedDefaultConfigsCount, len(configs))
}
})
t.Run("filter by type business", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/system-configs?type=business", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var configs []model.SystemConfig
_ = json.Unmarshal(dataBytes, &configs)
if len(configs) != 1 || configs[0].Key != model.ConfigKeyMaxAPIKeysPerUser {
t.Errorf("expected 1 business config (max_api_keys_per_user), got %d: %v", len(configs), configs)
}
})
}
func TestGetSystemConfig(t *testing.T) {
_, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("get existing configuration", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/system-configs/"+model.ConfigKeySiteName, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d", w.Code)
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var cfg model.SystemConfig
_ = json.Unmarshal(dataBytes, &cfg)
if cfg.Value != "OpenFlare" {
t.Errorf("expected 'OpenFlare', got '%s'", cfg.Value)
}
})
t.Run("get non-existent config", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/system-configs/non_existent_key", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d", w.Code)
}
})
}
func TestUpdateSystemConfig(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("update successfully", func(t *testing.T) {
hidden := model.ConfigVisibilityHidden
payload := UpdateSystemConfigRequest{
Value: "Super Site Name",
Visibility: &hidden,
Description: "Updated Description",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/"+model.ConfigKeySiteName, bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
// Verify database
var cfg model.SystemConfig
dbConn.Where("key = ?", model.ConfigKeySiteName).First(&cfg)
if cfg.Value != "Super Site Name" || cfg.Description != "Updated Description" || cfg.Visibility != model.ConfigVisibilityHidden {
t.Errorf("database values not updated: %+v", cfg)
}
// Verify caches are invalidated after update and repopulate on read
_, err := db.Redis.HGet(
context.Background(),
db.PrefixedKey(repository.SystemConfigRedisHashKey),
model.ConfigKeySiteName,
).Result()
if err == nil {
t.Fatal("expected redis cache miss immediately after update")
}
loaded, err := repository.GetSystemConfigByKey(context.Background(), model.ConfigKeySiteName)
if err != nil {
t.Fatalf("GetSystemConfigByKey(site_name) error = %v", err)
}
if loaded.Value != "Super Site Name" {
t.Errorf("GetSystemConfigByKey(site_name).Value = %q, want %q", loaded.Value, "Super Site Name")
}
if loaded.Visibility != model.ConfigVisibilityHidden {
t.Errorf("GetSystemConfigByKey(site_name).Visibility = %d, want %d", loaded.Visibility, model.ConfigVisibilityHidden)
}
})
t.Run("update non-existent config", func(t *testing.T) {
payload := UpdateSystemConfigRequest{
Value: "New Value",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/invalid_key", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d", w.Code)
}
})
}
func TestTestSMTP(t *testing.T) {
_, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
r := setupTestRouter(adminUser)
r.POST("/api/v1/admin/system-configs/smtp/test", TestSMTP)
// Start a mock SMTP server
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("failed to start mock smtp server: %v", err)
}
defer func() { _ = l.Close() }()
port := l.Addr().(*net.TCPAddr).Port
go func() {
conn, err := l.Accept()
if err != nil {
return
}
defer func() { _ = conn.Close() }()
writer := bufio.NewWriter(conn)
reader := bufio.NewReader(conn)
tp := textproto.NewReader(reader)
// 220 Ready
_, _ = writer.WriteString("220 mock.smtp.com SMTP Ready\r\n")
_ = writer.Flush()
// Read HELO/EHLO
_, _ = tp.ReadLine()
_, _ = writer.WriteString("250-mock.smtp.com\r\n250 AUTH PLAIN\r\n")
_ = writer.Flush()
// Read AUTH PLAIN
_, _ = tp.ReadLine()
_, _ = writer.WriteString("235 Authentication successful\r\n")
_ = writer.Flush()
// Read MAIL FROM
_, _ = tp.ReadLine()
_, _ = writer.WriteString("250 OK\r\n")
_ = writer.Flush()
// Read RCPT TO
_, _ = tp.ReadLine()
_, _ = writer.WriteString("250 OK\r\n")
_ = writer.Flush()
// Read DATA
_, _ = tp.ReadLine()
_, _ = writer.WriteString("354 Start mail input\r\n")
_ = writer.Flush()
// Read body lines until dot
for {
line, err := tp.ReadLine()
if err != nil || line == "." {
break
}
}
_, _ = writer.WriteString("250 OK\r\n")
_ = writer.Flush()
// Read QUIT
_, _ = tp.ReadLine()
_, _ = writer.WriteString("221 Bye\r\n")
_ = writer.Flush()
}()
payload := TestSMTPRequest{
SMTPHost: "127.0.0.1",
SMTPPort: port,
SMTPUsername: "sender@example.com",
SMTPPassword: "password",
To: "recipient@example.com",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/system-configs/smtp/test", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var testResp TestSMTPResponse
json.Unmarshal(dataBytes, &testResp)
if !testResp.Success {
t.Errorf("expected test success, got failed: %s. Log: %s", testResp.Error, testResp.Log)
}
}
func TestUpdateStorageConfigValidation(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("update storage config successfully", func(t *testing.T) {
tempDir := t.TempDir()
cfg := storage.DefaultConfig()
cfg.Local.Root = tempDir
cfgBytes, _ := json.Marshal(cfg)
payload := UpdateSystemConfigRequest{
Value: string(cfgBytes),
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/storage_config", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
// Verify database
var dbCfg model.SystemConfig
dbConn.Where("key = ?", "storage_config").First(&dbCfg)
var savedCfg storage.Config
_ = json.Unmarshal([]byte(dbCfg.Value), &savedCfg)
if savedCfg.Local.Root != tempDir {
t.Errorf("expected local root to be updated to %s, got %s", tempDir, savedCfg.Local.Root)
}
})
t.Run("update storage config failed connectivity check", func(t *testing.T) {
cfg := storage.DefaultConfig()
cfg.Driver = storage.DriverS3
cfg.S3.Bucket = "non-existent-bucket"
cfg.S3.Endpoint = "http://127.0.0.1:9999" // Will fail connectivity check
cfgBytes, _ := json.Marshal(cfg)
payload := UpdateSystemConfigRequest{
Value: string(cfgBytes),
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/storage_config", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
})
t.Run("reject driver switch when uploads exist", func(t *testing.T) {
upload := model.Upload{
ID: 88001,
UserID: 1,
FileName: "keep.txt",
FilePath: "uploads/keep.txt",
FileSize: 4,
MimeType: "text/plain",
Extension: "txt",
Type: "attachment",
Status: model.UploadStatusUsed,
}
if err := dbConn.Create(&upload).Error; err != nil {
t.Fatalf("seed upload failed: %v", err)
}
tempDir := t.TempDir()
cfg := storage.DefaultConfig()
cfg.Driver = storage.DriverS3
cfg.S3.Endpoint = "http://127.0.0.1:19998"
cfg.S3.Region = "us-east-1"
cfg.S3.Bucket = "wavelet"
cfg.S3.AccessKeyID = "test"
cfg.S3.SecretAccessKey = "test"
cfg.Local.Root = tempDir
cfgBytes, _ := json.Marshal(cfg)
payload := UpdateSystemConfigRequest{Value: string(cfgBytes)}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/storage_config", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Fatalf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
if !strings.Contains(w.Body.String(), StorageDriverSwitchRequiresMigration) {
t.Fatalf("expected migration-required error, got: %s", w.Body.String())
}
})
t.Run("switch to local while active s3 is unreachable", func(t *testing.T) {
if err := dbConn.Where("1 = 1").Delete(&model.Upload{}).Error; err != nil {
t.Fatalf("clear uploads failed: %v", err)
}
activeCfg := storage.DefaultConfig()
activeCfg.Driver = storage.DriverS3
activeCfg.S3.Endpoint = "http://127.0.0.1:9999"
activeCfg.S3.Region = "us-east-1"
activeCfg.S3.Bucket = "wavelet"
activeCfg.S3.AccessKeyID = "test"
activeCfg.S3.SecretAccessKey = "test"
activeBytes, _ := json.Marshal(activeCfg)
seedCfg := model.SystemConfig{
Key: "storage_config",
Value: string(activeBytes),
Type: "system",
}
if err := dbConn.Where("key = ?", "storage_config").
Assign(map[string]any{"value": seedCfg.Value, "type": seedCfg.Type}).
FirstOrCreate(&seedCfg).Error; err != nil {
t.Fatalf("seed active storage config failed: %v", err)
}
tempDir := t.TempDir()
stagedCfg := activeCfg
stagedCfg.Driver = storage.DriverLocal
stagedCfg.Local.Root = tempDir
cfgBytes, _ := json.Marshal(stagedCfg)
payload := UpdateSystemConfigRequest{Value: string(cfgBytes)}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/system-configs/storage_config", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var dbCfg model.SystemConfig
if err := dbConn.Where("key = ?", "storage_config").First(&dbCfg).Error; err != nil {
t.Fatalf("load saved storage config failed: %v", err)
}
var savedCfg storage.Config
if err := json.Unmarshal([]byte(dbCfg.Value), &savedCfg); err != nil {
t.Fatalf("parse saved storage config failed: %v", err)
}
if savedCfg.Driver != storage.DriverLocal {
t.Fatalf("active driver = %q, want %q after save", savedCfg.Driver, storage.DriverLocal)
}
if savedCfg.Local.Root != tempDir {
t.Fatalf("staged local root = %q, want %q", savedCfg.Local.Root, tempDir)
}
})
}
+23
View File
@@ -0,0 +1,23 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package task 提供任务管理接口
package task
// 任务管理相关错误消息
const (
InvalidTaskType = "无效的任务类型"
InvalidTimeRange = "无效的时间范围"
TaskDispatchFailed = "任务下发失败"
UserIDRequired = "用户ID必填"
TaskNotFound = "任务执行记录不存在"
TaskNotRetryable = "该任务不支持重试"
TaskNotFailed = "只有失败的任务才能重试"
TaskMaxRetryExceeded = "已达到最大重试次数"
TaskRetryFailed = "任务重试失败"
InvalidCronExpression = "无效的 Cron 表达式"
ScheduleNotFound = "定时任务不存在"
ScheduleSaveFailed = "保存定时任务失败"
ScheduleDeleteFailed = "删除定时任务失败"
)
+416
View File
@@ -0,0 +1,416 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package task
import (
"fmt"
"net/http"
"strconv"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/admin"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/internal/task/scheduler"
"github.com/Rain-kl/Wavelet/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/robfig/cron/v3"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// ListTaskTypes 获取支持的任务类型列表
// @Summary 获取支持的任务类型
// @Description 返回系统支持的所有可调度任务类型列表,包括任务名称、描述、是否支持时间范围等元数据,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]task.TaskMeta} "任务类型列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/tasks/types [get]
func ListTaskTypes(c *gin.Context) {
c.JSON(http.StatusOK, response.OK(task.GetDispatchableTasks()))
}
// DispatchTaskRequest 下发任务请求
type DispatchTaskRequest struct {
TaskType string `json:"task_type" binding:"required"`
StartTime *time.Time `json:"start_time"`
EndTime *time.Time `json:"end_time"`
UserID *uint64 `json:"user_id"`
Payload string `json:"payload"`
}
// DispatchTask 下发任务
// @Summary 下发异步任务
// @Description 手动触发指定类型的异步任务,支持指定时间范围和用户,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body DispatchTaskRequest true "任务请求参数"
// @Success 200 {object} response.Any{data=string} "任务已入队"
// @Failure 400 {object} response.Any "任务类型不存在或参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "任务入队失败"
// @Router /api/v1/admin/tasks/dispatch [post]
func DispatchTask(c *gin.Context) {
var req DispatchTaskRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
meta := task.GetTaskMeta(req.TaskType)
if meta == nil {
response.AbortBadRequest(c, InvalidTaskType)
return
}
var payloadBytes []byte
if strings.TrimSpace(req.Payload) != "" {
payloadBytes = []byte(req.Payload)
}
validated, err := task.ValidateAndNormalizePayload(meta.AsynqTask, payloadBytes)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
taskID, err := task.DispatchTask(c.Request.Context(), req.TaskType, validated, "manual")
if err != nil {
response.AbortInternal(c, fmt.Sprintf("%s: %v", TaskDispatchFailed, err))
return
}
c.JSON(http.StatusOK, response.OK(taskID))
}
// ListTaskExecutions 查询任务执行记录列表
// @Summary 查询任务执行记录
// @Description 分页查询任务执行记录,支持按状态和任务类型筛选,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param status query string false "状态筛选 (pending/running/succeeded/failed)"
// @Param task_type query string false "任务类型筛选"
// @Param page query int false "页码" default(1)
// @Param page_size query int false "每页条数" default(20)
// @Success 200 {object} response.Any{data=object} "任务执行记录列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/tasks/executions [get]
func ListTaskExecutions(c *gin.Context) {
var req model.ListTaskExecutionsRequest
if err := c.ShouldBindQuery(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
if req.TaskType != "" {
if meta := task.GetTaskMeta(req.TaskType); meta != nil {
req.TaskType = meta.AsynqTask
}
}
executions, total, err := model.ListTaskExecutions(c.Request.Context(), req)
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(gin.H{
"items": executions,
"total": total,
"page": req.Page,
"page_size": req.PageSize,
}))
}
// GetTaskExecution 查询单条任务执行详情
// @Summary 查询任务执行详情
// @Description 根据 ID 查询任务执行记录详情,包含完整执行日志,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path int true "任务执行记录 ID"
// @Success 200 {object} response.Any{data=model.TaskExecution} "任务执行详情"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "记录不存在"
// @Router /api/v1/admin/tasks/executions/{id} [get]
func GetTaskExecution(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, admin.InvalidTaskExecutionID)
return
}
execution, err := model.GetTaskExecutionByID(c.Request.Context(), id)
if err != nil {
response.AbortNotFound(c, TaskNotFound)
return
}
c.JSON(http.StatusOK, response.OK(execution))
}
// RetryTask 重试失败的任务
// @Summary 重试失败任务
// @Description 重新下发一条失败的任务,创建新的执行记录,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path int true "任务执行记录 ID"
// @Success 200 {object} response.Any{data=string} "新任务的 TaskID"
// @Failure 400 {object} response.Any "任务不支持重试或参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "记录不存在"
// @Failure 500 {object} response.Any "重试失败"
// @Router /api/v1/admin/tasks/executions/{id}/retry [post]
func RetryTask(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, admin.InvalidTaskExecutionID)
return
}
newTaskID, err := task.RetryTask(c.Request.Context(), id)
if err != nil {
errMsg := err.Error()
switch {
case strings.Contains(errMsg, "不存在"):
response.AbortNotFound(c, errMsg)
case strings.Contains(errMsg, "只有失败的任务") || strings.Contains(errMsg, "不支持重试") || strings.Contains(errMsg, "已达到最大重试"):
response.AbortBadRequest(c, errMsg)
default:
response.AbortInternal(c, fmt.Sprintf("%s: %v", TaskRetryFailed, err))
}
return
}
c.JSON(http.StatusOK, response.OK(newTaskID))
}
// ListSchedules 获取定时任务列表
// @Summary 获取定时任务列表
// @Description 返回系统所有的定时任务配置列表,包括名称、关联的异步任务类型、Cron 表达式和启用状态,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]model.Schedule} "定时任务列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Router /api/v1/admin/tasks/schedules [get]
func ListSchedules(c *gin.Context) {
schedules, err := model.ListSchedules(c.Request.Context())
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(schedules))
}
// CreateScheduleRequest 创建定时任务请求
type CreateScheduleRequest struct {
Name string `json:"name" binding:"required"`
TaskType string `json:"task_type" binding:"required"`
Cron string `json:"cron" binding:"required"`
Payload string `json:"payload"`
IsActive *bool `json:"is_active" binding:"required"`
}
// CreateSchedule 创建定时任务
// @Summary 创建定时任务
// @Description 新增一个动态定时任务配置,关联已有的异步任务,配置 Cron 表达式和执行参数,并触发调度器热加载,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body CreateScheduleRequest true "创建定时任务请求参数"
// @Success 200 {object} response.Any{data=model.Schedule} "创建成功的定时任务信息"
// @Failure 400 {object} response.Any "Cron 表达式无效、异步任务类型不存在或参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "保存定时任务失败"
// @Router /api/v1/admin/tasks/schedules [post]
func CreateSchedule(c *gin.Context) {
var req CreateScheduleRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
// 校验 Cron 表达式
if _, err := cron.ParseStandard(req.Cron); err != nil {
response.AbortBadRequest(c, InvalidCronExpression)
return
}
// 校验关联的异步任务类型
meta := task.GetTaskMeta(req.TaskType)
if meta == nil {
response.AbortBadRequest(c, InvalidTaskType)
return
}
// 校验并规范化 Payload
var payloadBytes []byte
if strings.TrimSpace(req.Payload) != "" {
payloadBytes = []byte(req.Payload)
}
validated, err := task.ValidateAndNormalizePayload(meta.AsynqTask, payloadBytes)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
schedule := &model.Schedule{
Name: req.Name,
TaskType: req.TaskType,
Cron: req.Cron,
Payload: string(validated),
IsActive: *req.IsActive,
}
if err := model.CreateSchedule(c.Request.Context(), schedule); err != nil {
response.AbortInternal(c, fmt.Sprintf("%s: %v", ScheduleSaveFailed, err))
return
}
// 触发调度服务重载
if err := scheduler.ReloadScheduler(); err != nil {
logger.ErrorF(c.Request.Context(), "[TaskAdmin] 重载调度器失败: %v", err)
}
c.JSON(http.StatusOK, response.OK(schedule))
}
// UpdateScheduleRequest 修改定时任务请求
type UpdateScheduleRequest struct {
Name string `json:"name" binding:"required"`
TaskType string `json:"task_type" binding:"required"`
Cron string `json:"cron" binding:"required"`
Payload string `json:"payload"`
IsActive *bool `json:"is_active" binding:"required"`
}
// UpdateSchedule 修改定时任务
// @Summary 修改定时任务
// @Description 修改一个定时任务的配置(名称、Cron 表达式、异步任务参数和是否启用等),并触发调度器热加载,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path int true "定时任务 ID"
// @Param request body UpdateScheduleRequest true "修改定时任务请求参数"
// @Success 200 {object} response.Any{data=model.Schedule} "修改后的定时任务信息"
// @Failure 400 {object} response.Any "Cron 表达式无效、参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "定时任务不存在"
// @Failure 500 {object} response.Any "修改定时任务失败"
// @Router /api/v1/admin/tasks/schedules/{id} [put]
func UpdateSchedule(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "无效的定时任务ID")
return
}
var req UpdateScheduleRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
// 检查定时任务是否存在
schedule, err := model.GetScheduleByID(c.Request.Context(), id)
if err != nil {
response.AbortNotFound(c, ScheduleNotFound)
return
}
// 校验 Cron 表达式
if _, err := cron.ParseStandard(req.Cron); err != nil {
response.AbortBadRequest(c, InvalidCronExpression)
return
}
// 校验关联的异步任务类型
meta := task.GetTaskMeta(req.TaskType)
if meta == nil {
response.AbortBadRequest(c, InvalidTaskType)
return
}
// 校验并规范化 Payload
var payloadBytes []byte
if strings.TrimSpace(req.Payload) != "" {
payloadBytes = []byte(req.Payload)
}
validated, err := task.ValidateAndNormalizePayload(meta.AsynqTask, payloadBytes)
if err != nil {
response.AbortBadRequest(c, err.Error())
return
}
schedule.Name = req.Name
schedule.TaskType = req.TaskType
schedule.Cron = req.Cron
schedule.Payload = string(validated)
schedule.IsActive = *req.IsActive
if err := model.UpdateSchedule(c.Request.Context(), schedule); err != nil {
response.AbortInternal(c, fmt.Sprintf("%s: %v", ScheduleSaveFailed, err))
return
}
// 触发调度服务重载
if err := scheduler.ReloadScheduler(); err != nil {
logger.ErrorF(c.Request.Context(), "[TaskAdmin] 重载调度器失败: %v", err)
}
c.JSON(http.StatusOK, response.OK(schedule))
}
// DeleteSchedule 删除定时任务
// @Summary 删除定时任务
// @Description 删除指定的定时任务配置,并触发调度器热加载,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path int true "定时任务 ID"
// @Success 200 {object} response.Any{data=string} "删除结果"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "删除定时任务失败"
// @Router /api/v1/admin/tasks/schedules/{id} [delete]
func DeleteSchedule(c *gin.Context) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil {
response.AbortBadRequest(c, "无效的定时任务ID")
return
}
if err := model.DeleteSchedule(c.Request.Context(), id); err != nil {
response.AbortInternal(c, fmt.Sprintf("%s: %v", ScheduleDeleteFailed, err))
return
}
// 触发调度服务重载
if err := scheduler.ReloadScheduler(); err != nil {
logger.ErrorF(c.Request.Context(), "[TaskAdmin] 重载调度器失败: %v", err)
}
c.JSON(http.StatusOK, response.OKNil())
}
+495
View File
@@ -0,0 +1,495 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package task
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
uploadtask "github.com/Rain-kl/Wavelet/internal/apps/upload/task"
"github.com/Rain-kl/Wavelet/internal/apps/user"
"github.com/Rain-kl/Wavelet/internal/bootstrap"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/task"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-gonic/gin"
"github.com/hibiken/asynq"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
func setupTaskTestEnvironment(t *testing.T) func() {
_, mr, cleanup := testhelper.SetupTestEnvironment(t)
bootstrap.RegisterTasks()
task.AsynqClient = asynq.NewClient(asynq.RedisClientOpt{
Addr: mr.Addr(),
})
return func() {
if task.AsynqClient != nil {
_ = task.AsynqClient.Close()
task.AsynqClient = nil
}
cleanup()
}
}
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
// Mock authentication middleware
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, oauth.UserObjKey, authUser)
}
c.Next()
})
adminGroup.GET("/tasks/types", ListTaskTypes)
adminGroup.POST("/tasks/dispatch", DispatchTask)
adminGroup.GET("/tasks/executions", ListTaskExecutions)
adminGroup.GET("/tasks/executions/:id", GetTaskExecution)
adminGroup.POST("/tasks/executions/:id/retry", RetryTask)
return r
}
func TestListTaskTypes(t *testing.T) {
cleanup := setupTaskTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/types", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d", w.Code)
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var taskMetas []task.TaskMeta
_ = json.Unmarshal(dataBytes, &taskMetas)
if len(taskMetas) == 0 {
t.Error("expected at least one dispatchable task type")
}
foundCleanup := false
foundWarmImageCache := false
for _, m := range taskMetas {
if m.Type == uploadtask.TaskTypeSystemCleanup {
foundCleanup = true
}
if m.Type == uploadtask.TaskTypeWarmImageCache {
foundWarmImageCache = true
}
}
if !foundCleanup {
t.Errorf("expected task type %s to be listed", uploadtask.TaskTypeSystemCleanup)
}
if !foundWarmImageCache {
t.Errorf("expected task type %s to be listed", uploadtask.TaskTypeWarmImageCache)
}
}
func TestDispatchTask(t *testing.T) {
cleanup := setupTaskTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("dispatch valid task successfully", func(t *testing.T) {
payload := DispatchTaskRequest{
TaskType: uploadtask.TaskTypeSystemCleanup,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code, "Body: %s", w.Body.String())
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
assert.Empty(t, resp.ErrorMsg)
assert.NotNil(t, resp.Data)
// 返回的 data 应该是 taskID
taskID, ok := resp.Data.(string)
assert.True(t, ok)
assert.NotEmpty(t, taskID)
})
t.Run("dispatch send_email task successfully with valid payload", func(t *testing.T) {
payload := DispatchTaskRequest{
TaskType: user.TaskTypeSendEmail,
Payload: `{"to":"receiver@example.com","subject":"Test Subject","body":"Test Body"}`,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code, "Body: %s", w.Body.String())
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
assert.Empty(t, resp.ErrorMsg)
assert.NotNil(t, resp.Data)
})
t.Run("dispatch send_email task failure with invalid payload json", func(t *testing.T) {
payload := DispatchTaskRequest{
TaskType: user.TaskTypeSendEmail,
Payload: `{"to":`,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
assert.Contains(t, resp.ErrorMsg, "无效的 JSON 格式")
})
t.Run("dispatch send_email task failure with missing fields", func(t *testing.T) {
payload := DispatchTaskRequest{
TaskType: user.TaskTypeSendEmail,
Payload: `{"to":"","subject":"Test","body":"Test"}`,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
assert.Contains(t, resp.ErrorMsg, "不能为空")
})
t.Run("dispatch invalid task type failure", func(t *testing.T) {
payload := DispatchTaskRequest{
TaskType: "invalid_task_type",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
assert.Equal(t, InvalidTaskType, resp.ErrorMsg)
})
t.Run("dispatch with empty body failure", func(t *testing.T) {
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/dispatch", bytes.NewBuffer([]byte("{}")))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
})
}
func TestListTaskExecutions(t *testing.T) {
cleanup := setupTaskTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
ctx := context.Background()
// 准备测试数据
now := time.Now()
records := []*model.TaskExecution{
{TaskID: "exec_001", TaskType: "system:cleanup", TaskName: "系统垃圾清理", Status: model.TaskExecutionStatusSucceeded, TriggeredBy: "manual", Duration: 1500, Result: "清理完成", StartedAt: &now, FinishedAt: &now},
{TaskID: "exec_002", TaskType: "system:cleanup", TaskName: "系统垃圾清理", Status: model.TaskExecutionStatusFailed, TriggeredBy: "system", ErrorMessage: "连接超时", StartedAt: &now, FinishedAt: &now},
{TaskID: "exec_003", TaskType: "system:cleanup", TaskName: "系统垃圾清理", Status: model.TaskExecutionStatusPending, TriggeredBy: "manual", Retryable: true, MaxRetry: 3},
}
for _, r := range records {
err := model.CreateTaskExecution(ctx, r)
require.NoError(t, err)
}
t.Run("list all executions", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var data map[string]interface{}
json.Unmarshal(dataBytes, &data)
assert.Equal(t, float64(3), data["total"])
})
t.Run("filter by status", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions?status=failed", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var data map[string]interface{}
json.Unmarshal(dataBytes, &data)
assert.Equal(t, float64(1), data["total"])
})
t.Run("filter by task_type (asynq task name)", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions?task_type=system:cleanup", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var data map[string]interface{}
json.Unmarshal(dataBytes, &data)
assert.Equal(t, float64(3), data["total"])
})
t.Run("filter by task_type (management task type)", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions?task_type=system_cleanup", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var data map[string]interface{}
json.Unmarshal(dataBytes, &data)
assert.Equal(t, float64(3), data["total"])
})
t.Run("pagination", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions?page=1&page_size=2", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var data map[string]interface{}
json.Unmarshal(dataBytes, &data)
assert.Equal(t, float64(3), data["total"])
})
}
func TestGetTaskExecution(t *testing.T) {
cleanup := setupTaskTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
ctx := context.Background()
// 创建测试记录
execution := &model.TaskExecution{
TaskID: "detail_001",
TaskType: "system:cleanup",
TaskName: "系统垃圾清理",
Status: model.TaskExecutionStatusSucceeded,
Log: "[10:00:01] 开始扫描\n[10:00:02] 找到 50 个文件\n[10:00:03] 清理完成",
Result: "共清理 50 个文件",
Duration: 2000,
Retryable: true,
MaxRetry: 3,
TriggeredBy: "manual",
}
err := model.CreateTaskExecution(ctx, execution)
require.NoError(t, err)
t.Run("get existing execution", func(t *testing.T) {
url := fmt.Sprintf("/api/v1/admin/tasks/executions/%d", execution.ID)
req, _ := http.NewRequest("GET", url, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var detail model.TaskExecution
json.Unmarshal(dataBytes, &detail)
assert.Equal(t, "detail_001", detail.TaskID)
assert.Equal(t, model.TaskExecutionStatusSucceeded, detail.Status)
assert.Contains(t, detail.Log, "开始扫描")
assert.Contains(t, detail.Log, "清理完成")
assert.Equal(t, int64(2000), detail.Duration)
})
t.Run("get non-existent execution", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions/99999999", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusNotFound, w.Code)
})
t.Run("invalid ID format", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/tasks/executions/invalid", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
})
}
func TestRetryTask(t *testing.T) {
cleanup := setupTaskTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
ctx := context.Background()
t.Run("retry failed task successfully", func(t *testing.T) {
now := time.Now()
execution := &model.TaskExecution{
TaskID: "retry_api_001",
TaskType: "system:cleanup",
TaskName: "系统垃圾清理",
Status: model.TaskExecutionStatusFailed,
ErrorMessage: "S3 连接超时",
Retryable: true,
MaxRetry: 3,
RetryCount: 0,
TriggeredBy: "manual",
StartedAt: &now,
FinishedAt: &now,
}
err := model.CreateTaskExecution(ctx, execution)
require.NoError(t, err)
url := fmt.Sprintf("/api/v1/admin/tasks/executions/%d/retry", execution.ID)
req, _ := http.NewRequest("POST", url, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusOK, w.Code)
var resp response.Any
json.Unmarshal(w.Body.Bytes(), &resp)
assert.Empty(t, resp.ErrorMsg)
assert.NotNil(t, resp.Data)
// 验证新记录
newTaskID, ok := resp.Data.(string)
assert.True(t, ok)
assert.NotEmpty(t, newTaskID)
newExecution, err := model.GetTaskExecutionByTaskID(ctx, newTaskID)
require.NoError(t, err)
assert.Equal(t, 1, newExecution.RetryCount)
assert.Equal(t, "retry", newExecution.TriggeredBy)
})
t.Run("retry succeeded task fails", func(t *testing.T) {
execution := &model.TaskExecution{
TaskID: "retry_succeeded_001",
TaskType: "system:cleanup",
TaskName: "系统垃圾清理",
Status: model.TaskExecutionStatusSucceeded,
Retryable: true,
MaxRetry: 3,
TriggeredBy: "manual",
}
err := model.CreateTaskExecution(ctx, execution)
require.NoError(t, err)
url := fmt.Sprintf("/api/v1/admin/tasks/executions/%d/retry", execution.ID)
req, _ := http.NewRequest("POST", url, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
})
t.Run("retry non-retryable task fails", func(t *testing.T) {
execution := &model.TaskExecution{
TaskID: "retry_not_allowed_001",
TaskType: "system:cleanup",
TaskName: "系统垃圾清理",
Status: model.TaskExecutionStatusFailed,
Retryable: false,
TriggeredBy: "manual",
}
err := model.CreateTaskExecution(ctx, execution)
require.NoError(t, err)
url := fmt.Sprintf("/api/v1/admin/tasks/executions/%d/retry", execution.ID)
req, _ := http.NewRequest("POST", url, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
})
t.Run("retry non-existent task", func(t *testing.T) {
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/executions/99999999/retry", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusNotFound, w.Code)
})
t.Run("retry with invalid ID", func(t *testing.T) {
req, _ := http.NewRequest("POST", "/api/v1/admin/tasks/executions/invalid/retry", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
assert.Equal(t, http.StatusBadRequest, w.Code)
})
}
+16
View File
@@ -0,0 +1,16 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package template 提供模板管理功能
package template
// 模板管理相关错误消息
const (
TemplateNotFound = "模板不存在"
TemplateKeyRequired = "模板标识符不能为空"
TemplateNameRequired = "模板名称不能为空"
TemplateContentRequired = "模板内容不能为空"
TemplateKeyExists = "模板标识符已存在"
SystemTemplateCannotDelete = "系统预置模板不可删除"
SystemTemplateCannotModifyKey = "系统预置模板不可修改标识符"
)
+78
View File
@@ -0,0 +1,78 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package template
import (
"context"
"errors"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
)
func createTemplate(ctx context.Context, req CreateTemplateRequest) (model.Template, error) {
exists, err := repository.TemplateExistsByKey(ctx, req.Key)
if err != nil {
return model.Template{}, err
}
if exists {
return model.Template{}, errors.New(TemplateKeyExists)
}
tmpl := model.Template{
Key: req.Key,
Name: req.Name,
Type: req.Type,
Subject: req.Subject,
Content: req.Content,
Description: req.Description,
IsSystem: false,
}
if err := tmpl.Validate(); err != nil {
return model.Template{}, err
}
if err := repository.CreateTemplate(ctx, &tmpl); err != nil {
return model.Template{}, err
}
return tmpl, nil
}
func listTemplates(ctx context.Context) ([]model.Template, error) {
return repository.ListTemplates(ctx)
}
func getTemplate(ctx context.Context, key string) (model.Template, error) {
return repository.GetTemplateByKey(ctx, key)
}
func updateTemplate(ctx context.Context, key string, req UpdateTemplateRequest) (model.Template, error) {
tmpl, err := repository.GetTemplateByKey(ctx, key)
if err != nil {
return model.Template{}, err
}
tmpl.Name = req.Name
tmpl.Type = req.Type
tmpl.Subject = req.Subject
tmpl.Content = req.Content
tmpl.Description = req.Description
if err := tmpl.Validate(); err != nil {
return model.Template{}, err
}
if err := repository.SaveTemplate(ctx, &tmpl); err != nil {
return model.Template{}, err
}
return tmpl, nil
}
func deleteTemplate(ctx context.Context, key string) error {
tmpl, err := repository.GetTemplateByKey(ctx, key)
if err != nil {
return err
}
if tmpl.IsSystem {
return errors.New(SystemTemplateCannotDelete)
}
return repository.DeleteTemplate(ctx, &tmpl)
}
+176
View File
@@ -0,0 +1,176 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package template
import (
"errors"
"net/http"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// CreateTemplateRequest 创建模板请求
type CreateTemplateRequest struct {
Key string `json:"key" binding:"required,max=80"`
Name string `json:"name" binding:"required,max=100"`
Type string `json:"type" binding:"required,max=20"`
Subject string `json:"subject" binding:"max=255"`
Content string `json:"content" binding:"required"`
Description string `json:"description" binding:"max=255"`
}
// UpdateTemplateRequest 更新模板请求
type UpdateTemplateRequest struct {
Name string `json:"name" binding:"required,max=100"`
Type string `json:"type" binding:"required,max=20"`
Subject string `json:"subject" binding:"max=255"`
Content string `json:"content" binding:"required"`
Description string `json:"description" binding:"max=255"`
}
func abortTemplateLogicError(c *gin.Context, err error) bool {
if err == nil {
return false
}
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, TemplateNotFound)
return true
}
msg := err.Error()
switch msg {
case TemplateKeyExists, SystemTemplateCannotDelete:
response.AbortBadRequest(c, msg)
return true
}
response.AbortInternal(c, msg)
return true
}
// CreateTemplate 创建模板
// @Summary 创建模板
// @Description 创建一条新的自定义通知模板,模板标识符(Key)不可重复,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body template.CreateTemplateRequest true "创建请求参数"
// @Success 200 {object} response.Any{data=string} "创建成功"
// @Failure 400 {object} response.Any "参数错误或模板标识符已存在"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/templates [post]
func CreateTemplate(c *gin.Context) {
var req CreateTemplateRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
tmpl, err := createTemplate(c.Request.Context(), req)
if abortTemplateLogicError(c, err) {
return
}
c.JSON(http.StatusOK, response.OK(tmpl))
}
// ListTemplates 获取模板列表
// @Summary 获取模板列表
// @Description 返回所有通知模板列表,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=[]model.Template} "模板列表"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/templates [get]
func ListTemplates(c *gin.Context) {
templates, err := listTemplates(c.Request.Context())
if err != nil {
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(templates))
}
// GetTemplate 获取单个模板
// @Summary 获取单个模板
// @Description 根据模板标识符获取对应的模板详情,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param key path string true "模板标识符"
// @Success 200 {object} response.Any{data=model.Template} "模板详情"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "模板不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/templates/{key} [get]
func GetTemplate(c *gin.Context) {
tmpl, err := getTemplate(c.Request.Context(), c.Param("key"))
if abortTemplateLogicError(c, err) {
return
}
c.JSON(http.StatusOK, response.OK(tmpl))
}
// UpdateTemplate 更新模板
// @Summary 更新模板
// @Description 根据模板标识符更新对应的模板内容,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param key path string true "模板标识符"
// @Param request body template.UpdateTemplateRequest true "更新请求参数"
// @Success 200 {object} response.Any{data=model.Template} "更新成功"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "模板不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/templates/{key} [put]
func UpdateTemplate(c *gin.Context) {
var req UpdateTemplateRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
tmpl, err := updateTemplate(c.Request.Context(), c.Param("key"), req)
if abortTemplateLogicError(c, err) {
return
}
c.JSON(http.StatusOK, response.OK(tmpl))
}
// DeleteTemplate 删除模板
// @Summary 删除模板
// @Description 根据模板标识符删除对应模板,系统预置模板不可删除,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param key path string true "模板标识符"
// @Success 200 {object} response.Any{data=string} "删除成功"
// @Failure 400 {object} response.Any "不可删除系统模板"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "模板不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/templates/{key} [delete]
func DeleteTemplate(c *gin.Context) {
if err := deleteTemplate(c.Request.Context(), c.Param("key")); abortTemplateLogicError(c, err) {
return
}
c.JSON(http.StatusOK, response.OKNil())
}
@@ -0,0 +1,242 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package template
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
// Mock authentication middleware
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, oauth.UserObjKey, authUser)
}
c.Next()
})
adminGroup.GET("/templates", ListTemplates)
adminGroup.POST("/templates", CreateTemplate)
templateRouter := adminGroup.Group("/templates/:key")
{
templateRouter.GET("", GetTemplate)
templateRouter.PUT("", UpdateTemplate)
templateRouter.DELETE("", DeleteTemplate)
}
return r
}
func TestCreateTemplate(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("create successfully", func(t *testing.T) {
payload := CreateTemplateRequest{
Key: "test_template",
Name: "Test Template",
Type: "email",
Subject: "Test Subject",
Content: "Hello {{.Name}}",
Description: "Test Desc",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/templates", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var tmpl model.Template
err := dbConn.Where("key = ?", "test_template").First(&tmpl).Error
if err != nil {
t.Fatalf("failed to find template in DB: %v", err)
}
if tmpl.Name != "Test Template" {
t.Errorf("expected Name 'Test Template', got '%s'", tmpl.Name)
}
})
t.Run("create duplicate key error", func(t *testing.T) {
payload := CreateTemplateRequest{
Key: "test_template",
Name: "Another Name",
Type: "email",
Subject: "Another Subject",
Content: "Hello",
Description: "desc",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/templates", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request on duplicate key, got %d", w.Code)
}
})
}
func TestListTemplates(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
// Seed system templates manually for testing
t1 := model.Template{Key: "login_email", Name: "Login Code", Type: "email", Content: "code {{.Code}}", IsSystem: true}
t2 := model.Template{Key: "register_email", Name: "Register Code", Type: "email", Content: "code {{.Code}}", IsSystem: true}
dbConn.Create(&t1)
dbConn.Create(&t2)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("list templates", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/templates", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d", w.Code)
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var templates []model.Template
_ = json.Unmarshal(dataBytes, &templates)
if len(templates) != 2 {
t.Errorf("expected 2 templates, got %d", len(templates))
}
})
}
func TestGetTemplate(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
t1 := model.Template{Key: "login_email", Name: "Login Code", Type: "email", Content: "code {{.Code}}", IsSystem: true}
dbConn.Create(&t1)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("get existing", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/templates/login_email", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d", w.Code)
}
})
t.Run("get non-existent", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/templates/non_existent", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d", w.Code)
}
})
}
func TestUpdateTemplate(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
t1 := model.Template{Key: "login_email", Name: "Login Code", Type: "email", Content: "code {{.Code}}", IsSystem: true}
dbConn.Create(&t1)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("update successfully", func(t *testing.T) {
payload := UpdateTemplateRequest{
Name: "Updated Login Code",
Type: "email",
Subject: "New Subject",
Content: "new code {{.Code}}",
Description: "new desc",
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/templates/login_email", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var tmpl model.Template
dbConn.Where("key = ?", "login_email").First(&tmpl)
if tmpl.Name != "Updated Login Code" || tmpl.Subject != "New Subject" {
t.Errorf("database values not updated: %+v", tmpl)
}
})
}
func TestDeleteTemplate(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
t1 := model.Template{Key: "login_email", Name: "Login Code", Type: "email", Content: "code {{.Code}}", IsSystem: true}
t2 := model.Template{Key: "custom_tmpl", Name: "Custom", Type: "email", Content: "hi", IsSystem: false}
dbConn.Create(&t1)
dbConn.Create(&t2)
adminUser := &model.User{ID: 1001, Username: "admin", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("delete system template should fail", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/templates/login_email", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request when deleting system template, got %d", w.Code)
}
})
t.Run("delete custom template should succeed", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/templates/custom_tmpl", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d", w.Code)
}
var count int64
dbConn.Model(&model.Template{}).Where("key = ?", "custom_tmpl").Count(&count)
if count != 0 {
t.Error("custom template was not deleted from DB")
}
})
}
+17
View File
@@ -0,0 +1,17 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package updater manages GitHub Release checks and in-place application upgrades.
package updater
const (
errInvalidRepository = "上游仓库地址无效"
errReleaseRequestFailed = "获取上游版本失败"
errReleaseResponseInvalid = "上游版本响应无效"
errNoCompatibleRelease = "未找到兼容的 Release"
errNoCompatibleAsset = "未找到当前系统对应的 Release 资产"
errDevelopmentBuild = "开发版本无法执行自动升级"
errAlreadyUpToDate = "当前已是最新版本"
errUpgradeAlreadyRunning = "已有升级任务正在执行"
errAutomaticUpgradeBlocked = "当前平台暂不支持自动替换二进制"
)
+40
View File
@@ -0,0 +1,40 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import "context"
// GetStatus returns the current build and the newest compatible upstream release.
func GetStatus(ctx context.Context) (Status, error) {
status, _, err := defaultManager.status(ctx)
return status, err
}
// PrepareUpgrade downloads and stages the upgrade binary for the current platform.
func PrepareUpgrade(ctx context.Context) (executable string, stagedBinary string, status Status, err error) {
status, _, err = defaultManager.status(ctx)
if err != nil {
return "", "", Status{}, err
}
executable, stagedBinary, err = defaultManager.prepareUpgrade(ctx)
return executable, stagedBinary, status, err
}
// ApplyPreparedUpgrade replaces the running binary and restarts the process.
func ApplyPreparedUpgrade(executable, stagedBinary string) error {
return replaceAndRestart(executable, stagedBinary)
}
// FinishUpgrade clears the in-progress upgrade flag after a failed restart.
func FinishUpgrade() {
defaultManager.finishUpgrade()
}
// IsUpgrading reports whether an upgrade task is currently running.
func IsUpgrading() bool {
defaultManager.mu.Lock()
defer defaultManager.mu.Unlock()
return defaultManager.upgrading
}
+647
View File
@@ -0,0 +1,647 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import (
"archive/tar"
"archive/zip"
"compress/gzip"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
"time"
"github.com/Rain-kl/Wavelet/internal/buildinfo"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
"github.com/Rain-kl/Wavelet/pkg/logger"
"golang.org/x/mod/semver"
)
const (
githubAPIBaseURL = "https://api.github.com"
maxArchiveSize = int64(1024 * 1024 * 1024)
maxReleaseSize = int64(4 * 1024 * 1024)
repositoryParts = 2
windowsOS = "windows"
archiveFileMode = 0o600
stagedBinaryMode = 0o700
)
type releaseAsset struct {
Name string `json:"name"`
BrowserDownloadURL string `json:"browser_download_url"`
Size int64 `json:"size"`
State string `json:"state"`
}
type githubRelease struct {
TagName string `json:"tag_name"`
Name string `json:"name"`
Body string `json:"body"`
HTMLURL string `json:"html_url"`
Draft bool `json:"draft"`
Prerelease bool `json:"prerelease"`
Published time.Time `json:"published_at"`
Assets []releaseAsset `json:"assets"`
}
// Status describes the current build and the newest compatible upstream release.
type Status struct {
CurrentVersion string `json:"current_version"`
BuildTime string `json:"build_time"`
LatestVersion string `json:"latest_version"`
UpdateAvailable bool `json:"update_available"`
CanUpgrade bool `json:"can_upgrade"`
Prerelease bool `json:"prerelease"`
ReleaseName string `json:"release_name"`
ReleaseNotes string `json:"release_notes"`
ReleaseURL string `json:"release_url"`
PublishedAt string `json:"published_at"`
UpstreamRepository string `json:"upstream_repository"`
AssetName string `json:"asset_name"`
Platform string `json:"platform"`
}
type releaseClient interface {
Do(req *http.Request) (*http.Response, error)
}
type manager struct {
client releaseClient
mu sync.Mutex
upgrading bool
}
var defaultManager = &manager{
client: &http.Client{Timeout: 10 * time.Minute},
}
func normalizeVersion(version string) string {
version = strings.TrimSpace(version)
if version == "" || version == "dev" {
return ""
}
if !strings.HasPrefix(version, "v") {
version = "v" + version
}
if !semver.IsValid(version) {
return ""
}
return version
}
func parseRepository(raw string) (string, error) {
raw = strings.TrimSpace(raw)
if raw == "" {
return "", errors.New(errInvalidRepository)
}
if !strings.Contains(raw, "://") {
repo := strings.TrimSuffix(strings.Trim(raw, "/"), ".git")
if len(strings.Split(repo, "/")) == repositoryParts {
return repo, nil
}
return "", errors.New(errInvalidRepository)
}
parsed, err := url.Parse(raw)
if err != nil || !strings.EqualFold(parsed.Hostname(), "github.com") {
return "", errors.New(errInvalidRepository)
}
repo := strings.TrimSuffix(strings.Trim(parsed.Path, "/"), ".git")
if len(strings.Split(repo, "/")) != repositoryParts {
return "", errors.New(errInvalidRepository)
}
return repo, nil
}
func expectedAssetName(tag string) string {
extension := "tar.gz"
if runtime.GOOS == windowsOS {
extension = "zip"
}
return fmt.Sprintf("wavelet_%s_%s_%s.%s", tag, runtime.GOOS, runtime.GOARCH, extension)
}
func expectedAssetNames(repository, tag string) []string {
names := []string{expectedAssetName(tag)}
if parts := strings.Split(repository, "/"); len(parts) == repositoryParts {
repoName := parts[1]
if repoName != "wavelet" {
extension := "tar.gz"
if runtime.GOOS == windowsOS {
extension = "zip"
}
names = append(names, fmt.Sprintf("%s_%s_%s_%s.%s", repoName, tag, runtime.GOOS, runtime.GOARCH, extension))
}
}
return names
}
func selectLatestRelease(repository string, releases []githubRelease) (githubRelease, releaseAsset, error) {
var selected githubRelease
var selectedAsset releaseAsset
selectedVersion := ""
for _, release := range releases {
version := normalizeVersion(release.TagName)
if release.Draft || version == "" {
continue
}
expectedNames := expectedAssetNames(repository, release.TagName)
for _, asset := range release.Assets {
matched := false
for _, name := range expectedNames {
if asset.Name == name {
matched = true
break
}
}
if !matched || asset.BrowserDownloadURL == "" || asset.State != "uploaded" {
continue
}
if selectedVersion == "" || semver.Compare(version, selectedVersion) > 0 {
selected = release
selectedAsset = asset
selectedVersion = version
}
}
}
if selectedVersion == "" {
return githubRelease{}, releaseAsset{}, errors.New(errNoCompatibleRelease)
}
return selected, selectedAsset, nil
}
func (m *manager) fetchRelease(ctx context.Context, repository string) (githubRelease, releaseAsset, error) {
req, err := http.NewRequestWithContext(
ctx,
http.MethodGet,
fmt.Sprintf("%s/repos/%s/releases?per_page=30", githubAPIBaseURL, repository),
nil,
)
if err != nil {
return githubRelease{}, releaseAsset{}, fmt.Errorf("%s: %w", errReleaseRequestFailed, err)
}
req.Header.Set("Accept", "application/vnd.github+json")
req.Header.Set("User-Agent", "OpenFlare-Updater")
req.Header.Set("X-GitHub-Api-Version", "2022-11-28")
resp, err := m.client.Do(req)
if err != nil {
return githubRelease{}, releaseAsset{}, fmt.Errorf("%s: %w", errReleaseRequestFailed, err)
}
defer func() {
// The response body is read-only; close errors cannot affect the parsed result.
_ = resp.Body.Close()
}()
if resp.StatusCode != http.StatusOK {
return githubRelease{}, releaseAsset{}, fmt.Errorf("%s: HTTP %d", errReleaseRequestFailed, resp.StatusCode)
}
var releases []githubRelease
decoder := json.NewDecoder(io.LimitReader(resp.Body, maxReleaseSize))
if err := decoder.Decode(&releases); err != nil {
return githubRelease{}, releaseAsset{}, fmt.Errorf("%s: %w", errReleaseResponseInvalid, err)
}
release, asset, err := selectLatestRelease(repository, releases)
if err != nil {
return githubRelease{}, releaseAsset{}, err
}
logger.InfoF(ctx, "[Updater] Selected latest compatible release: %s (Asset: %s)", release.TagName, asset.Name)
return release, asset, nil
}
func loadRepository(ctx context.Context) (string, error) {
config, err := repository.GetSystemConfigByKey(ctx, model.ConfigKeyUpdateUpstreamRepository)
if err != nil {
return "", fmt.Errorf("%s: %w", errInvalidRepository, err)
}
return parseRepository(config.Value)
}
func (m *manager) status(ctx context.Context) (Status, releaseAsset, error) {
upstreamRepo, err := loadRepository(ctx)
if err != nil {
return Status{}, releaseAsset{}, err
}
release, asset, err := m.fetchRelease(ctx, upstreamRepo)
if err != nil {
return Status{}, releaseAsset{}, err
}
currentVersion := normalizeVersion(buildinfo.Version)
latestVersion := normalizeVersion(release.TagName)
updateAvailable := currentVersion != "" && semver.Compare(latestVersion, currentVersion) > 0
logger.InfoF(ctx, "[Updater] Check update complete. current: %s, latest: %s, update_available: %t", buildinfo.Version, release.TagName, updateAvailable)
return Status{
CurrentVersion: buildinfo.Version,
BuildTime: buildinfo.BuildTime,
LatestVersion: release.TagName,
UpdateAvailable: updateAvailable,
CanUpgrade: updateAvailable && runtime.GOOS != windowsOS,
Prerelease: release.Prerelease,
ReleaseName: release.Name,
ReleaseNotes: release.Body,
ReleaseURL: release.HTMLURL,
PublishedAt: release.Published.Format(time.RFC3339),
UpstreamRepository: upstreamRepo,
AssetName: asset.Name,
Platform: runtime.GOOS + "/" + runtime.GOARCH,
}, asset, nil
}
func downloadArchive(ctx context.Context, client releaseClient, asset releaseAsset, destination string) error {
if asset.Size <= 0 || asset.Size > maxArchiveSize {
return fmt.Errorf("release 资产大小无效: %d", asset.Size)
}
logger.InfoF(ctx, "[Updater] Downloading release asset: %s", asset.Name)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, asset.BrowserDownloadURL, nil)
if err != nil {
return fmt.Errorf("创建升级下载请求失败: %w", err)
}
req.Header.Set("User-Agent", "OpenFlare-Updater")
resp, err := client.Do(req)
if err != nil {
return fmt.Errorf("下载升级资产失败: %w", err)
}
defer func() {
// The downloaded body has already been validated by size before use.
_ = resp.Body.Close()
}()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("下载升级资产失败: HTTP %d", resp.StatusCode)
}
file, err := os.OpenFile(destination, os.O_CREATE|os.O_EXCL|os.O_WRONLY, archiveFileMode) //nolint:gosec // destination is created inside the verified executable directory.
if err != nil {
return fmt.Errorf("创建升级归档失败: %w", err)
}
written, err := io.Copy(file, io.LimitReader(resp.Body, maxArchiveSize+1))
if err != nil {
_ = file.Close()
return fmt.Errorf("写入升级归档失败: %w", err)
}
if err := file.Close(); err != nil {
return fmt.Errorf("关闭升级归档失败: %w", err)
}
if written > maxArchiveSize || written != asset.Size {
return fmt.Errorf("升级归档大小不匹配: got %d, want %d", written, asset.Size)
}
logger.InfoF(ctx, "[Updater] Successfully downloaded release asset to %s", destination)
return nil
}
func safeArchivePath(destination, name string) (string, error) {
cleanName := filepath.Clean(name)
if filepath.IsAbs(cleanName) || cleanName == "." || strings.HasPrefix(cleanName, ".."+string(filepath.Separator)) {
return "", fmt.Errorf("归档包含非法路径: %s", name)
}
target := filepath.Join(destination, cleanName)
relative, err := filepath.Rel(destination, target)
if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return "", fmt.Errorf("归档路径越界: %s", name)
}
return target, nil
}
func matchBinaryName(name string, candidates []string) bool {
for _, candidate := range candidates {
if runtime.GOOS == windowsOS {
if strings.EqualFold(name, candidate) {
return true
}
} else {
if name == candidate {
return true
}
}
}
return false
}
func getCandidateBinaryNames(executable string, repository string) []string {
execName := filepath.Base(executable)
names := []string{execName}
addName := func(base string) {
name := base
if runtime.GOOS == windowsOS && !strings.HasSuffix(strings.ToLower(name), ".exe") {
name += ".exe"
}
for _, existing := range names {
if existing == name {
return
}
}
names = append(names, name)
}
if parts := strings.Split(repository, "/"); len(parts) == repositoryParts {
addName(parts[1])
}
addName("wavelet")
return names
}
func isLikelyBinary(name string, isDir bool, mode os.FileMode) bool {
if isDir {
return false
}
base := strings.ToLower(filepath.Base(name))
// Exclude typical non-binary metadata files
exclusions := []string{
"license", "licence", "copying", "notice", "readme", "changelog",
}
for _, excl := range exclusions {
if strings.HasPrefix(base, excl) {
return false
}
}
if runtime.GOOS == windowsOS {
return filepath.Ext(base) == ".exe"
}
// On Unix, it should either have the executable permission bit set, OR have no extension
return (mode.Perm()&0111 != 0) || (filepath.Ext(base) == "")
}
func findBinaryInTarGz(archivePath string, candidates []string) (string, error) {
file, err := os.Open(archivePath) //nolint:gosec // archivePath is created by prepareUpgrade in the executable directory.
if err != nil {
return "", err
}
defer func() {
_ = file.Close()
}()
gzipReader, err := gzip.NewReader(file)
if err != nil {
return "", err
}
defer func() {
_ = gzipReader.Close()
}()
reader := tar.NewReader(gzipReader)
var binaries []string
for {
header, err := reader.Next()
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return "", err
}
if header.Typeflag == tar.TypeReg && isLikelyBinary(header.Name, false, header.FileInfo().Mode()) {
binaries = append(binaries, header.Name)
}
}
if len(binaries) == 1 {
return binaries[0], nil
}
// Fallback to candidate match if multiple or zero likely binaries found
for _, name := range binaries {
if matchBinaryName(filepath.Base(name), candidates) {
return name, nil
}
}
return "", errors.New(errNoCompatibleAsset)
}
func findBinaryInZip(archivePath string, candidates []string) (string, error) {
reader, err := zip.OpenReader(archivePath)
if err != nil {
return "", err
}
defer func() {
_ = reader.Close()
}()
var binaries []string
for _, file := range reader.File {
if !file.FileInfo().IsDir() && isLikelyBinary(file.Name, false, file.FileInfo().Mode()) {
binaries = append(binaries, file.Name)
}
}
if len(binaries) == 1 {
return binaries[0], nil
}
// Fallback to candidate match if multiple or zero likely binaries found
for _, name := range binaries {
if matchBinaryName(filepath.Base(name), candidates) {
return name, nil
}
}
return "", errors.New(errNoCompatibleAsset)
}
func extractTarGz(ctx context.Context, archivePath, destination, targetName string, candidates []string) (string, error) {
binaryPathInArchive, err := findBinaryInTarGz(archivePath, candidates)
if err != nil {
return "", err
}
logger.InfoF(ctx, "[Updater] Extracting tar.gz archive: %s (extracting: %s)", archivePath, binaryPathInArchive)
file, err := os.Open(archivePath) //nolint:gosec // archivePath is created by prepareUpgrade in the executable directory.
if err != nil {
return "", err
}
defer func() {
// Read-only archive close errors do not change extraction validity.
_ = file.Close()
}()
gzipReader, err := gzip.NewReader(file)
if err != nil {
return "", err
}
defer func() {
// The gzip checksum is verified while reading the selected file.
_ = gzipReader.Close()
}()
reader := tar.NewReader(gzipReader)
for {
header, err := reader.Next()
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return "", err
}
if header.Name != binaryPathInArchive {
continue
}
target, err := safeArchivePath(destination, targetName)
if err != nil {
return "", err
}
output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, stagedBinaryMode) //nolint:gosec // target is constrained by safeArchivePath.
if err != nil {
return "", err
}
written, copyErr := io.Copy(output, io.LimitReader(reader, maxArchiveSize+1))
closeErr := output.Close()
if copyErr != nil {
return "", copyErr
}
if closeErr != nil {
return "", closeErr
}
if written > maxArchiveSize {
return "", errors.New("解压后的程序文件超过大小限制")
}
logger.InfoF(ctx, "[Updater] Successfully extracted binary to %s", target)
return target, nil
}
return "", errors.New(errNoCompatibleAsset)
}
func extractZip(ctx context.Context, archivePath, destination, targetName string, candidates []string) (string, error) {
binaryPathInArchive, err := findBinaryInZip(archivePath, candidates)
if err != nil {
return "", err
}
logger.InfoF(ctx, "[Updater] Extracting zip archive: %s (extracting: %s)", archivePath, binaryPathInArchive)
reader, err := zip.OpenReader(archivePath)
if err != nil {
return "", err
}
defer func() {
// Read-only archive close errors do not change extraction validity.
_ = reader.Close()
}()
for _, file := range reader.File {
if file.Name != binaryPathInArchive {
continue
}
target, err := safeArchivePath(destination, targetName)
if err != nil {
return "", err
}
input, err := file.Open()
if err != nil {
return "", err
}
output, err := os.OpenFile(target, os.O_CREATE|os.O_EXCL|os.O_WRONLY, stagedBinaryMode) //nolint:gosec // target is constrained by safeArchivePath.
if err != nil {
_ = input.Close()
return "", err
}
written, copyErr := io.Copy(output, io.LimitReader(input, maxArchiveSize+1))
inputCloseErr := input.Close()
outputCloseErr := output.Close()
if copyErr != nil {
return "", copyErr
}
if inputCloseErr != nil {
return "", inputCloseErr
}
if outputCloseErr != nil {
return "", outputCloseErr
}
if written > maxArchiveSize {
return "", errors.New("解压后的程序文件超过大小限制")
}
logger.InfoF(ctx, "[Updater] Successfully extracted binary to %s", target)
return target, nil
}
return "", errors.New(errNoCompatibleAsset)
}
func (m *manager) prepareUpgrade(ctx context.Context) (string, string, error) {
if runtime.GOOS == windowsOS {
return "", "", errors.New(errAutomaticUpgradeBlocked)
}
if normalizeVersion(buildinfo.Version) == "" {
return "", "", errors.New(errDevelopmentBuild)
}
m.mu.Lock()
defer m.mu.Unlock()
if m.upgrading {
return "", "", errors.New(errUpgradeAlreadyRunning)
}
status, asset, err := m.status(ctx)
if err != nil {
return "", "", err
}
if !status.UpdateAvailable {
return "", "", errors.New(errAlreadyUpToDate)
}
logger.InfoF(ctx, "[Updater] Preparing upgrade. current: %s, latest: %s", status.CurrentVersion, status.LatestVersion)
executable, err := os.Executable()
if err != nil {
return "", "", fmt.Errorf("定位当前程序失败: %w", err)
}
executable, err = filepath.EvalSymlinks(executable)
if err != nil {
return "", "", fmt.Errorf("解析当前程序路径失败: %w", err)
}
tempDir, err := os.MkdirTemp(filepath.Dir(executable), ".wavelet-update-*")
if err != nil {
return "", "", fmt.Errorf("创建升级目录失败: %w", err)
}
archivePath := filepath.Join(tempDir, asset.Name)
if err := downloadArchive(ctx, m.client, asset, archivePath); err != nil {
// Cleanup is best effort because the download error is the actionable failure.
_ = os.RemoveAll(tempDir)
return "", "", err
}
targetName := filepath.Base(executable)
candidates := getCandidateBinaryNames(executable, status.UpstreamRepository)
var stagedBinary string
if strings.HasSuffix(asset.Name, ".zip") {
stagedBinary, err = extractZip(ctx, archivePath, tempDir, targetName, candidates)
} else {
stagedBinary, err = extractTarGz(ctx, archivePath, tempDir, targetName, candidates)
}
if err != nil {
// Cleanup is best effort because the extraction error is the actionable failure.
_ = os.RemoveAll(tempDir)
return "", "", fmt.Errorf("解压升级资产失败: %w", err)
}
logger.InfoF(ctx, "[Updater] Staged binary successfully prepared: %s", stagedBinary)
m.upgrading = true
return executable, stagedBinary, nil
}
func (m *manager) finishUpgrade() {
m.mu.Lock()
defer m.mu.Unlock()
m.upgrading = false
}
+130
View File
@@ -0,0 +1,130 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import (
"runtime"
"testing"
"time"
)
func TestParseRepository(t *testing.T) {
tests := []struct {
name string
input string
want string
wantErr bool
}{
{name: "short form", input: "Rain-kl/OpenFlare", want: "Rain-kl/OpenFlare"},
{name: "GitHub URL", input: "https://github.com/Rain-kl/OpenFlare.git", want: "Rain-kl/OpenFlare"},
{name: "unsupported host", input: "https://example.com/Rain-kl/OpenFlare", wantErr: true},
{name: "missing owner", input: "OpenFlare", wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got, err := parseRepository(tt.input)
if gotErr := err != nil; gotErr != tt.wantErr {
t.Errorf("parseRepository(%q) error = %v, want error presence = %t", tt.input, err, tt.wantErr)
}
if got != tt.want {
t.Errorf("parseRepository(%q) = %q, want %q", tt.input, got, tt.want)
}
})
}
}
func TestSelectLatestRelease(t *testing.T) {
assetNameV1 := expectedAssetName("v1.0.0")
assetNameV2 := expectedAssetName("v2.0.0")
releases := []githubRelease{
{
TagName: "v1.0.0",
Published: time.Date(2026, time.June, 1, 0, 0, 0, 0, time.UTC),
Assets: []releaseAsset{{
Name: assetNameV1,
BrowserDownloadURL: "https://example.com/v1",
State: "uploaded",
}},
},
{
TagName: "v2.0.0",
Published: time.Date(2026, time.June, 2, 0, 0, 0, 0, time.UTC),
Assets: []releaseAsset{{
Name: assetNameV2,
BrowserDownloadURL: "https://example.com/v2",
State: "uploaded",
}},
},
{
TagName: "v3.0.0",
Assets: []releaseAsset{{
Name: "wavelet_v3.0.0_other_platform.tar.gz",
BrowserDownloadURL: "https://example.com/v3",
State: "uploaded",
}},
},
}
release, asset, err := selectLatestRelease("Rain-kl/OpenFlare", releases)
if err != nil {
t.Fatalf("selectLatestRelease() error = %v", err)
}
if release.TagName != "v2.0.0" {
t.Errorf("selectLatestRelease() tag = %q, want %q", release.TagName, "v2.0.0")
}
if asset.Name != assetNameV2 {
t.Errorf("selectLatestRelease() asset = %q, want %q", asset.Name, assetNameV2)
}
}
func TestSelectLatestReleaseWithCustomRepo(t *testing.T) {
extension := "tar.gz"
if runtime.GOOS == "windows" {
extension = "zip"
}
releases := []githubRelease{
{
TagName: "v1.0.0",
Published: time.Date(2026, time.June, 1, 0, 0, 0, 0, time.UTC),
Assets: []releaseAsset{{
Name: "wavelet_v1.0.0_" + runtime.GOOS + "_" + runtime.GOARCH + "." + extension,
BrowserDownloadURL: "https://example.com/v1",
State: "uploaded",
}},
},
{
TagName: "v2.0.0",
Published: time.Date(2026, time.June, 2, 0, 0, 0, 0, time.UTC),
Assets: []releaseAsset{{
Name: "PixezSync_v2.0.0_" + runtime.GOOS + "_" + runtime.GOARCH + "." + extension,
BrowserDownloadURL: "https://example.com/v2",
State: "uploaded",
}},
},
}
release, asset, err := selectLatestRelease("Rain-kl/PixezSync", releases)
if err != nil {
t.Fatalf("selectLatestRelease() error = %v", err)
}
if release.TagName != "v2.0.0" {
t.Errorf("selectLatestRelease() tag = %q, want %q", release.TagName, "v2.0.0")
}
expectedName := "PixezSync_v2.0.0_" + runtime.GOOS + "_" + runtime.GOARCH + "." + extension
if asset.Name != expectedName {
t.Errorf("selectLatestRelease() asset = %q, want %q", asset.Name, expectedName)
}
}
func TestExpectedAssetName(t *testing.T) {
extension := "tar.gz"
if runtime.GOOS == "windows" {
extension = "zip"
}
want := "wavelet_v1.2.3_" + runtime.GOOS + "_" + runtime.GOARCH + "." + extension
if got := expectedAssetName("v1.2.3"); got != want {
t.Errorf("expectedAssetName(%q) = %q, want %q", "v1.2.3", got, want)
}
}
@@ -0,0 +1,50 @@
//go:build !windows
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import (
"context"
"fmt"
"os"
"path/filepath"
"syscall"
"github.com/Rain-kl/Wavelet/pkg/logger"
)
const installedBinaryMode = 0o755
func replaceAndRestart(executable, stagedBinary string) error {
ctx := context.Background()
logger.InfoF(ctx, "[Updater] Swapping executable: %s -> %s", executable, stagedBinary)
backup := executable + ".old"
if err := os.Remove(backup); err != nil && !os.IsNotExist(err) {
return fmt.Errorf("删除旧备份失败: %w", err)
}
if err := os.Rename(executable, backup); err != nil {
return fmt.Errorf("备份当前程序失败: %w", err)
}
if err := os.Rename(stagedBinary, executable); err != nil {
_ = os.Rename(backup, executable)
return fmt.Errorf("替换当前程序失败: %w", err)
}
if err := os.Chmod(executable, installedBinaryMode); err != nil { //nolint:gosec // the installed application binary must be executable.
_ = os.Remove(executable)
_ = os.Rename(backup, executable)
return fmt.Errorf("设置程序执行权限失败: %w", err)
}
stagingDir := filepath.Dir(stagedBinary)
// Cleanup is best effort; a leftover staging directory must not block restart.
_ = os.RemoveAll(stagingDir)
logger.InfoF(ctx, "[Updater] Executing syscall.Exec to restart service: %s %v", executable, os.Args)
return syscall.Exec(executable, os.Args, os.Environ()) //nolint:gosec // executable is resolved from os.Executable and never supplied by a request.
}
@@ -0,0 +1,12 @@
//go:build windows
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import "errors"
func replaceAndRestart(_, _ string) error {
return errors.New(errAutomaticUpgradeBlocked)
}
+68
View File
@@ -0,0 +1,68 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package updater
import (
"context"
"net/http"
"time"
"github.com/Rain-kl/Wavelet/pkg/logger"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// GetUpdateStatus 获取应用更新状态
// @Summary 获取应用更新状态
// @Description 从系统配置指定的 GitHub 上游仓库查询最新兼容 Release,并与当前服务版本比较
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any{data=updater.Status} "更新状态"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "查询失败"
// @Router /api/v1/admin/update [get]
func GetUpdateStatus(c *gin.Context) {
status, _, err := defaultManager.status(c.Request.Context())
if err != nil {
logger.ErrorF(c.Request.Context(), "[Updater] check release failed: %v", err)
response.AbortInternal(c, err.Error())
return
}
c.JSON(http.StatusOK, response.OK(status))
}
// ApplyUpdate 下载并应用应用更新
// @Summary 下载并应用应用更新
// @Description 下载当前平台对应的 GitHub Actions Release 资产,替换当前二进制并重启进程
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Success 200 {object} response.Any "升级已准备并即将重启"
// @Failure 400 {object} response.Any "当前版本不可升级"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "升级准备失败"
// @Router /api/v1/admin/update/apply [post]
func ApplyUpdate(c *gin.Context) {
executable, stagedBinary, err := defaultManager.prepareUpgrade(c.Request.Context())
if err != nil {
logger.ErrorF(c.Request.Context(), "[Updater] prepare upgrade failed: %v", err)
response.AbortBadRequest(c, err.Error())
return
}
logger.InfoF(c.Request.Context(), "[Updater] upgrade prepared; restarting with %s", stagedBinary)
c.JSON(http.StatusOK, response.OKNil())
go func() {
time.Sleep(time.Second)
if err := replaceAndRestart(executable, stagedBinary); err != nil {
defaultManager.finishUpgrade()
logger.ErrorF(context.Background(), "[Updater] replace and restart failed: %v", err)
}
}()
}
+21
View File
@@ -0,0 +1,21 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package user 提供用户管理功能
package user
const (
userNotFound = "用户不存在"
cannotDisable = "不能禁用管理员用户"
cannotDelete = "不能删除管理员用户"
cannotDeleteSelf = "不能删除当前登录用户"
updateUserFailed = "更新用户状态失败"
deleteUserFailed = "删除用户失败"
usernameExists = "用户名已存在"
usernameRequired = "用户名不能为空"
passwordTooShort = "密码长度不能少于 8 位" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
createUserFailed = "创建用户失败"
emailRequired = "邮箱不能为空"
emailExists = "邮箱已被注册"
)
+106
View File
@@ -0,0 +1,106 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package user
import (
"context"
"errors"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/db/idgen"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/repository"
)
func listUsers(ctx context.Context, req listUsersRequest) (int64, []model.User, error) {
return repository.ListAdminUsers(ctx, repository.AdminUserListFilter{
UserID: req.UserID,
Username: strings.TrimSpace(req.Username),
Page: req.Page,
PageSize: req.PageSize,
})
}
func getUserDetail(ctx context.Context, id uint64) (model.User, error) {
return repository.GetAdminUserDetail(ctx, id)
}
func updateUserStatus(ctx context.Context, id uint64, active bool) error {
flags, err := repository.GetUserAdminFlags(ctx, id)
if err != nil {
return err
}
if !active && flags.IsAdmin {
return errors.New(cannotDisable)
}
return repository.UpdateUserActive(ctx, id, active)
}
func deleteUser(ctx context.Context, currentUserID, targetID uint64) error {
if currentUserID == targetID {
return errors.New(cannotDeleteSelf)
}
flags, err := repository.GetUserAdminFlags(ctx, targetID)
if err != nil {
return err
}
if flags.IsAdmin {
return errors.New(cannotDelete)
}
return repository.DeleteUserWithRelations(ctx, targetID)
}
func createUser(ctx context.Context, req createUserRequest) (model.User, error) {
req.Username = strings.TrimSpace(req.Username)
req.Nickname = strings.TrimSpace(req.Nickname)
req.Password = strings.TrimSpace(req.Password)
req.Email = strings.TrimSpace(req.Email)
if req.Username == "" {
return model.User{}, errors.New(usernameRequired)
}
if req.Email == "" {
return model.User{}, errors.New(emailRequired)
}
if len(req.Password) < minPasswordLength {
return model.User{}, errors.New(passwordTooShort)
}
count, err := repository.CountUsersByUsername(ctx, req.Username)
if err != nil {
return model.User{}, err
}
if count > 0 {
return model.User{}, errors.New(usernameExists)
}
emailCount, err := repository.CountUsersByEmail(ctx, req.Email)
if err != nil {
return model.User{}, err
}
if emailCount > 0 {
return model.User{}, errors.New(emailExists)
}
newUser := model.User{
ID: idgen.NextUint64ID(),
Username: req.Username,
Nickname: req.Nickname,
Email: req.Email,
IsActive: req.IsActive,
IsAdmin: req.IsAdmin,
LastLoginAt: time.Time{},
}
if newUser.Nickname == "" {
newUser.Nickname = req.Username
}
if err := newUser.SetEncryptedPassword(req.Password); err != nil {
return model.User{}, err
}
if err := repository.CreateUser(ctx, &newUser); err != nil {
return model.User{}, err
}
return newUser, nil
}
+288
View File
@@ -0,0 +1,288 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package user
import (
"errors"
"net/http"
"strconv"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/gin-gonic/gin"
"gorm.io/gorm"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// minPasswordLength 密码最小长度
const minPasswordLength = 8
// listUsersRequest 用户列表查询请求
type listUsersRequest struct {
Page int `form:"page" binding:"min=1"`
PageSize int `form:"page_size" binding:"min=1,max=100"`
UserID *uint64 `form:"user_id" binding:"omitempty,gt=0"`
Username string `form:"username"`
}
type user struct {
ID uint64 `json:"id,string"`
Username string `json:"username"`
Nickname string `json:"nickname"`
Email string `json:"email"`
AvatarURL string `json:"avatar_url"`
IsActive bool `json:"is_active"`
IsAdmin bool `json:"is_admin"`
Bio string `json:"bio"`
Phone string `json:"phone"`
Gender string `json:"gender"`
Website string `json:"website"`
Location string `json:"location"`
LastLoginAt time.Time `json:"last_login_at"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// listUsersResponse 用户列表响应
type listUsersResponse struct {
Users []user `json:"users"`
Total int64 `json:"total"`
}
func parseUserID(c *gin.Context) (uint64, bool) {
id, err := strconv.ParseUint(c.Param("id"), 10, 64)
if err != nil || id == 0 {
response.AbortBadRequest(c, userNotFound)
return 0, false
}
return id, true
}
func toUser(u model.User) user {
return user{
ID: u.ID,
Username: u.Username,
Nickname: u.Nickname,
Email: u.Email,
AvatarURL: u.AvatarURL,
IsActive: u.IsActive,
IsAdmin: u.IsAdmin,
Bio: u.Bio,
Phone: u.Phone,
Gender: u.Gender,
Website: u.Website,
Location: u.Location,
LastLoginAt: u.LastLoginAt,
CreatedAt: u.CreatedAt,
UpdatedAt: u.UpdatedAt,
}
}
func abortUserLogicError(c *gin.Context, err error, notFoundMsg string, forbiddenMsgs, badRequestMsgs []string) bool {
if err == nil {
return false
}
if errors.Is(err, gorm.ErrRecordNotFound) {
response.AbortNotFound(c, notFoundMsg)
return true
}
msg := err.Error()
for _, m := range badRequestMsgs {
if msg == m {
response.AbortBadRequest(c, msg)
return true
}
}
for _, m := range forbiddenMsgs {
if msg == m {
response.AbortForbidden(c, msg)
return true
}
}
response.AbortInternal(c, msg)
return true
}
// ListUsers 获取用户列表
// @Summary 获取用户列表
// @Description 分页返回用户列表,支持按用户 ID 和用户名筛选,需要管理员权限
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param request query listUsersRequest true "查询参数"
// @Success 200 {object} response.Any{data=user.listUsersResponse} "用户列表"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/users [get]
func ListUsers(c *gin.Context) {
var req listUsersRequest
if err := c.ShouldBindQuery(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
total, modelUsers, err := listUsers(c.Request.Context(), req)
if err != nil {
response.AbortInternal(c, err.Error())
return
}
users := make([]user, 0, len(modelUsers))
for _, modelUser := range modelUsers {
users = append(users, toUser(modelUser))
}
c.JSON(http.StatusOK, response.OK(listUsersResponse{
Users: users,
Total: total,
}))
}
// GetUser 获取用户详情
// @Summary 获取用户详情
// @Description 返回指定用户的完整个人资料和系统状态,需要管理员权限,不返回密码等敏感字段
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path int true "用户 ID"
// @Success 200 {object} response.Any{data=user.user} "用户详情"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 404 {object} response.Any "用户不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/users/{id} [get]
func GetUser(c *gin.Context) {
id, ok := parseUserID(c)
if !ok {
return
}
targetUser, err := getUserDetail(c.Request.Context(), id)
if abortUserLogicError(c, err, userNotFound, nil, nil) {
return
}
c.JSON(http.StatusOK, response.OK(toUser(targetUser)))
}
// updateUserStatusRequest 更新用户状态请求
type updateUserStatusRequest struct {
IsActive bool `json:"is_active"`
}
// UpdateUserStatus 更新用户状态(启用/禁用)
// @Summary 更新用户状态
// @Description 启用或禁用指定用户,管理员账号无法被禁用,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param id path int true "用户 ID"
// @Param request body updateUserStatusRequest true "状态参数"
// @Success 200 {object} response.Any{data=string} "更新成功"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限或尝试禁用管理员"
// @Failure 404 {object} response.Any "用户不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/users/{id}/status [put]
func UpdateUserStatus(c *gin.Context) {
var req updateUserStatusRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
id, ok := parseUserID(c)
if !ok {
return
}
if err := updateUserStatus(c.Request.Context(), id, req.IsActive); err != nil {
if abortUserLogicError(c, err, userNotFound, []string{cannotDisable}, nil) {
return
}
response.AbortInternal(c, updateUserFailed)
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// DeleteUser 删除用户
// @Summary 删除用户
// @Description 删除指定非管理员用户,需要管理员权限,不能删除当前登录用户
// @Tags admin
// @Produce json
// @Security SessionCookie
// @Param id path int true "用户 ID"
// @Success 200 {object} response.Any{data=string} "删除成功"
// @Failure 400 {object} response.Any "参数错误"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限、尝试删除管理员或当前用户"
// @Failure 404 {object} response.Any "用户不存在"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/users/{id} [delete]
func DeleteUser(c *gin.Context) {
id, ok := parseUserID(c)
if !ok {
return
}
currUser, _ := oauth.GetFromContext[*model.User](c, oauth.UserObjKey)
if err := deleteUser(c.Request.Context(), currUser.ID, id); err != nil {
if abortUserLogicError(c, err, userNotFound, []string{cannotDelete, cannotDeleteSelf}, nil) {
return
}
response.AbortInternal(c, deleteUserFailed)
return
}
c.JSON(http.StatusOK, response.OKNil())
}
// createUserRequest 创建用户请求
type createUserRequest struct {
Username string `json:"username" binding:"required,min=3,max=64"`
Password string `json:"password" binding:"required,min=8,max=64"`
Nickname string `json:"nickname" binding:"omitempty,max=64"`
Email string `json:"email" binding:"required,email,max=255"`
IsActive bool `json:"is_active"`
IsAdmin bool `json:"is_admin"`
}
// CreateUser 创建用户
// @Summary 创建用户
// @Description 创建一个本地密码登录的新用户,需要管理员权限
// @Tags admin
// @Accept json
// @Produce json
// @Security SessionCookie
// @Param request body user.createUserRequest true "创建用户参数"
// @Success 200 {object} response.Any{data=user.user} "创建成功"
// @Failure 400 {object} response.Any "参数错误或用户名已存在"
// @Failure 401 {object} response.Any "未登录"
// @Failure 403 {object} response.Any "无管理员权限"
// @Failure 500 {object} response.Any "内部错误"
// @Router /api/v1/admin/users [post]
func CreateUser(c *gin.Context) {
var req createUserRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.AbortBadRequest(c, err.Error())
return
}
newUser, err := createUser(c.Request.Context(), req)
if abortUserLogicError(c, err, "", nil, []string{usernameRequired, emailRequired, passwordTooShort, usernameExists, emailExists}) {
return
}
c.JSON(http.StatusOK, response.OK(toUser(newUser)))
}
+582
View File
@@ -0,0 +1,582 @@
// Copyright 2025 linux.do
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package user
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/oauth"
"github.com/Rain-kl/Wavelet/internal/model"
"github.com/Rain-kl/Wavelet/internal/testhelper"
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
func setupTestRouter(authUser *model.User) *gin.Engine {
r := testhelper.NewTestGinEngine()
adminGroup := r.Group("/api/v1/admin")
// Mock authentication middleware
adminGroup.Use(func(c *gin.Context) {
if authUser != nil {
oauth.SetToContext(c, oauth.UserObjKey, authUser)
}
c.Next()
})
adminGroup.GET("/users", ListUsers)
adminGroup.POST("/users", CreateUser)
adminGroup.GET("/users/:id", GetUser)
adminGroup.PUT("/users/:id/status", UpdateUserStatus)
adminGroup.DELETE("/users/:id", DeleteUser)
return r
}
func TestListUsers(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
// Seed users
users := []model.User{
{
ID: 1001,
Username: "alice",
Nickname: "Alice Nickname",
IsActive: true,
IsAdmin: false,
LastLoginAt: time.Now(),
},
{
ID: 1002,
Username: "bob",
Nickname: "Bob Nickname",
IsActive: true,
IsAdmin: false,
LastLoginAt: time.Now(),
},
{
ID: 1003,
Username: "charlie",
Nickname: "Charlie Nickname",
IsActive: false,
IsAdmin: true,
LastLoginAt: time.Now(),
},
}
for _, u := range users {
if err := dbConn.Create(&u).Error; err != nil {
t.Fatalf("failed to seed user: %v", err)
}
}
adminUser := &model.User{ID: 1003, Username: "charlie", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("basic pagination list", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users?page=1&page_size=2", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
// Parse data map to our structure
dataBytes, _ := json.Marshal(resp.Data)
var listResp listUsersResponse
if err := json.Unmarshal(dataBytes, &listResp); err != nil {
t.Fatalf("failed to parse list response: %v", err)
}
if len(listResp.Users) != 2 {
t.Errorf("expected 2 users, got %d", len(listResp.Users))
}
if listResp.Total != 3 {
t.Errorf("expected total 3, got %d", listResp.Total)
}
// Ordered by ID ASC
if listResp.Users[0].ID != 1001 || listResp.Users[1].ID != 1002 {
t.Errorf("expected ordered ASC, got first ID %d, second ID %d", listResp.Users[0].ID, listResp.Users[1].ID)
}
})
t.Run("filter by user_id", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users?page=1&page_size=10&user_id=1001", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var listResp listUsersResponse
_ = json.Unmarshal(dataBytes, &listResp)
if len(listResp.Users) != 1 || listResp.Users[0].ID != 1001 {
t.Errorf("expected 1 user with ID 1001, got total %d", len(listResp.Users))
}
})
t.Run("filter by username prefix", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users?page=1&page_size=10&username=bo", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
dataBytes, _ := json.Marshal(resp.Data)
var listResp listUsersResponse
_ = json.Unmarshal(dataBytes, &listResp)
if len(listResp.Users) != 1 || listResp.Users[0].Username != "bob" {
t.Errorf("expected bob, got %v", listResp.Users)
}
})
t.Run("invalid pagination parameter", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users?page=0&page_size=10", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d", w.Code)
}
})
}
func TestGetUser(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
targetUser := model.User{
ID: 1001,
Username: "alice",
Password: "secret-hash",
Nickname: "Alice Nickname",
Email: "alice@example.com",
AvatarURL: "https://example.com/avatar.png",
IsActive: true,
IsAdmin: false,
Bio: "hello",
Phone: "123456",
Gender: "female",
Website: "https://example.com",
Location: "Shanghai",
}
if err := dbConn.Create(&targetUser).Error; err != nil {
t.Fatalf("failed to seed user: %v", err)
}
adminUser := &model.User{ID: 1003, Username: "charlie", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("get full user profile", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users/1001", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
dataBytes, _ := json.Marshal(resp.Data)
var resUser user
if err := json.Unmarshal(dataBytes, &resUser); err != nil {
t.Fatalf("failed to parse response data: %v", err)
}
if resUser.Email != targetUser.Email || resUser.Bio != targetUser.Bio || resUser.Phone != targetUser.Phone ||
resUser.Gender != targetUser.Gender || resUser.Website != targetUser.Website || resUser.Location != targetUser.Location {
t.Errorf("profile fields were not returned correctly: %+v", resUser)
}
if bytes.Contains(dataBytes, []byte("secret-hash")) {
t.Error("response should not include password")
}
})
t.Run("get non-existent user", func(t *testing.T) {
req, _ := http.NewRequest("GET", "/api/v1/admin/users/9999", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d. Body: %s", w.Code, w.Body.String())
}
})
}
func TestUpdateUserStatus(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
// Seed users
regularUser := model.User{
ID: 1001,
Username: "alice",
IsActive: true,
IsAdmin: false,
}
adminUser := model.User{
ID: 1002,
Username: "bob",
IsActive: true,
IsAdmin: true,
}
dbConn.Create(&regularUser)
dbConn.Create(&adminUser)
router := setupTestRouter(&adminUser)
t.Run("disable regular user successfully", func(t *testing.T) {
payload := updateUserStatusRequest{IsActive: false}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/users/1001/status", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
// Verify DB status
var u model.User
dbConn.First(&u, 1001)
if u.IsActive {
t.Error("user should be deactivated in the database")
}
})
t.Run("cannot disable admin user", func(t *testing.T) {
payload := updateUserStatusRequest{IsActive: false}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/users/1002/status", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Errorf("expected 403 Forbidden, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
if resp.ErrorMsg != cannotDisable {
t.Errorf("expected error message '%s', got '%s'", cannotDisable, resp.ErrorMsg)
}
})
t.Run("cannot enable/disable non-existent user", func(t *testing.T) {
payload := updateUserStatusRequest{IsActive: false}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("PUT", "/api/v1/admin/users/9999/status", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d. Body: %s", w.Code, w.Body.String())
}
})
}
func TestCreateUser(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
adminUser := &model.User{ID: 1003, Username: "charlie", IsAdmin: true}
router := setupTestRouter(adminUser)
t.Run("create user successfully", func(t *testing.T) {
payload := createUserRequest{
Username: "newuser",
Password: "newpassword123",
Nickname: "New Nickname",
Email: "newuser@example.com",
IsActive: true,
IsAdmin: false,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/users", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Errorf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
t.Fatalf("failed to parse response: %v", err)
}
if resp.ErrorMsg != "" {
t.Errorf("expected empty error message, got '%s'", resp.ErrorMsg)
}
dataBytes, _ := json.Marshal(resp.Data)
var resUser user
if err := json.Unmarshal(dataBytes, &resUser); err != nil {
t.Fatalf("failed to parse response data: %v", err)
}
if resUser.Username != "newuser" || resUser.Nickname != "New Nickname" || !resUser.IsActive || resUser.IsAdmin {
t.Errorf("unexpected user values: %+v", resUser)
}
// Verify in DB
var dbUser model.User
if err := dbConn.Where("username = ?", "newuser").First(&dbUser).Error; err != nil {
t.Fatalf("failed to find user in db: %v", err)
}
if dbUser.Email != "newuser@example.com" {
t.Errorf("expected email 'newuser@example.com', got '%s'", dbUser.Email)
}
if !dbUser.CheckPassword("newpassword123") {
t.Error("password was not hashed correctly")
}
})
t.Run("create user with duplicate username", func(t *testing.T) {
// Create the first user
existing := model.User{
ID: 2001,
Username: "dupuser",
Nickname: "Dup User",
Email: "dupuser@example.com",
}
dbConn.Create(&existing)
payload := createUserRequest{
Username: "dupuser",
Password: "password123",
Nickname: "Another Nick",
Email: "another@example.com",
IsActive: true,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/users", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
if resp.ErrorMsg != usernameExists {
t.Errorf("expected error '%s', got '%s'", usernameExists, resp.ErrorMsg)
}
})
t.Run("create user with duplicate email", func(t *testing.T) {
existing := model.User{
ID: 2002,
Username: "existingemail",
Nickname: "Existing Email",
Email: "dupemail@example.com",
}
dbConn.Create(&existing)
payload := createUserRequest{
Username: "newuser2",
Password: "password123",
Nickname: "New User 2",
Email: "dupemail@example.com",
IsActive: true,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/users", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
var resp response.Any
_ = json.Unmarshal(w.Body.Bytes(), &resp)
if resp.ErrorMsg != emailExists {
t.Errorf("expected error '%s', got '%s'", emailExists, resp.ErrorMsg)
}
})
t.Run("validation error - password too short", func(t *testing.T) {
payload := createUserRequest{
Username: "shortpass",
Password: "123",
Email: "shortpass@example.com",
IsActive: true,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/users", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
})
t.Run("validation error - invalid email format", func(t *testing.T) {
payload := map[string]interface{}{
"username": "bademail",
"password": "password123",
"email": "not-an-email",
"is_active": true,
}
body, _ := json.Marshal(payload)
req, _ := http.NewRequest("POST", "/api/v1/admin/users", bytes.NewBuffer(body))
req.Header.Set("Content-Type", "application/json")
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusBadRequest {
t.Errorf("expected 400 Bad Request, got %d. Body: %s", w.Code, w.Body.String())
}
})
}
func TestDeleteUser(t *testing.T) {
dbConn, _, cleanup := testhelper.SetupTestEnvironment(t)
defer cleanup()
if err := dbConn.AutoMigrate(&model.AccessToken{}, &model.ExternalAccount{}); err != nil {
t.Fatalf("failed to migrate delete-related tables: %v", err)
}
regularUser := model.User{
ID: 1001,
Username: "alice",
IsActive: true,
IsAdmin: false,
}
adminUser := model.User{
ID: 1002,
Username: "bob",
IsActive: true,
IsAdmin: true,
}
selfUser := model.User{
ID: 1003,
Username: "charlie",
IsActive: true,
IsAdmin: false,
}
if err := dbConn.Create(&regularUser).Error; err != nil {
t.Fatalf("failed to seed regular user: %v", err)
}
if err := dbConn.Create(&adminUser).Error; err != nil {
t.Fatalf("failed to seed admin user: %v", err)
}
if err := dbConn.Create(&selfUser).Error; err != nil {
t.Fatalf("failed to seed self user: %v", err)
}
if err := dbConn.Create(&model.AccessToken{
UserID: regularUser.ID,
Name: "api",
TokenHash: "hash-for-delete-user-test",
MaskedToken: "at_****test",
}).Error; err != nil {
t.Fatalf("failed to seed access token: %v", err)
}
if err := dbConn.Create(&model.ExternalAccount{
ID: 5001,
AuthSourceID: 1,
UserID: regularUser.ID,
ExternalID: "external-alice",
}).Error; err != nil {
t.Fatalf("failed to seed external account: %v", err)
}
router := setupTestRouter(&selfUser)
t.Run("delete regular user successfully", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/users/1001", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("expected 200 OK, got %d. Body: %s", w.Code, w.Body.String())
}
var count int64
if err := dbConn.Model(&model.User{}).Where("id = ?", 1001).Count(&count).Error; err != nil {
t.Fatalf("failed to count deleted user: %v", err)
}
if count != 0 {
t.Errorf("expected deleted user count 0, got %d", count)
}
if err := dbConn.Model(&model.AccessToken{}).Where("user_id = ?", 1001).Count(&count).Error; err != nil {
t.Fatalf("failed to count deleted access tokens: %v", err)
}
if count != 0 {
t.Errorf("expected deleted access token count 0, got %d", count)
}
if err := dbConn.Model(&model.ExternalAccount{}).Where("user_id = ?", 1001).Count(&count).Error; err != nil {
t.Fatalf("failed to count deleted external accounts: %v", err)
}
if count != 0 {
t.Errorf("expected deleted external account count 0, got %d", count)
}
})
t.Run("cannot delete admin user", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/users/1002", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Fatalf("expected 403 Forbidden, got %d. Body: %s", w.Code, w.Body.String())
}
})
t.Run("cannot delete current user", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/users/1003", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Fatalf("expected 403 Forbidden, got %d. Body: %s", w.Code, w.Body.String())
}
})
t.Run("delete non-existent user", func(t *testing.T) {
req, _ := http.NewRequest("DELETE", "/api/v1/admin/users/9999", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
if w.Code != http.StatusNotFound {
t.Errorf("expected 404 Not Found, got %d. Body: %s", w.Code, w.Body.String())
}
})
}
+699
View File
@@ -0,0 +1,699 @@
package agent
import (
"context"
"encoding/json"
"errors"
"log/slog"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
"github.com/Rain-kl/Wavelet/internal/apps/agent/observability"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
"github.com/Rain-kl/Wavelet/internal/apps/agent/state"
"github.com/Rain-kl/Wavelet/internal/apps/agent/wsclient"
)
type HeartbeatService interface {
Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error)
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
SetToken(token string)
}
type SyncService interface {
SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error
SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
ForceSyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
WAFIPGroupChecksums() (map[string]string, error)
ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error
}
type Updater interface {
CheckAndUpdate(ctx context.Context, repo string, options UpdateOptions) error
}
type RuntimeManager interface {
CheckHealth(ctx context.Context) error
Restart(ctx context.Context) error
}
type WebSocketService interface {
Connect(ctx context.Context) (protocol.WebSocketConnection, error)
SetToken(token string)
URL() string
}
type UpdateOptions struct {
Channel string
TagName string
Force bool
}
type Runner struct {
Config *config.Config
StateStore *state.Store
ObservabilityBuffer *state.ObservabilityBufferStore
HeartbeatService HeartbeatService
SyncService SyncService
Updater Updater
RuntimeManager RuntimeManager
WebSocketService WebSocketService
autoUpdate bool
updateNow bool
updateRepo string
updateChan string
updateTag string
restartOpenrestyNow bool
websocketUpgradeEnabled bool
}
func (r *Runner) Run(ctx context.Context) error {
nodeID, err := r.StateStore.EnsureNodeID()
if err != nil {
return err
}
slog.Info("agent runner started", "node_id", nodeID, "node", r.Config.NodeName, "ip", r.Config.NodeIP)
if r.hasAccessToken() {
if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, true); hbErr != nil {
slog.Error("agent startup heartbeat failed", "error", hbErr)
}
} else if err = r.tryRegister(ctx, &nodeID); err != nil {
slog.Error("agent initial discovery register failed", "error", err)
}
heartbeatTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
defer heartbeatTicker.Stop()
var wsDone <-chan error
wsBackoff := newWebSocketBackoff()
nextWSAttempt := time.Now()
tryStartWebSocket := func() {
if wsDone != nil || !r.shouldUseWebSocket() || time.Now().Before(nextWSAttempt) {
return
}
done, startErr := r.startWebSocket(ctx, nodeID)
if startErr != nil {
delay := wsBackoff.Next()
nextWSAttempt = time.Now().Add(delay)
slog.Debug("agent ws upgrade failed; falling back to http heartbeat",
"enabled", r.websocketUpgradeEnabled,
"url", r.websocketURL(),
"retry_after", delay,
"error", startErr,
)
return
}
wsBackoff.Reset()
wsDone = done
slog.Debug("agent switched to websocket mode", "url", r.websocketURL())
}
tryStartWebSocket()
for {
select {
case <-ctx.Done():
slog.Info("agent runner shutting down", "error", ctx.Err())
return ctx.Err()
case wsErr := <-wsDone:
wsDone = nil
delay := wsBackoff.Next()
nextWSAttempt = time.Now().Add(delay)
slog.Debug("agent ws disconnected; resuming http heartbeat", "retry_after", delay, "error", wsErr)
if r.hasAccessToken() {
if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
slog.Error("agent heartbeat after ws disconnect failed", "error", hbErr)
}
}
case <-heartbeatTicker.C:
if wsDone != nil {
continue
}
if !r.hasAccessToken() {
if err = r.tryRegister(ctx, &nodeID); err != nil {
slog.Error("agent discovery register failed", "error", err)
}
continue
}
if changed, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
slog.Error("agent heartbeat failed", "error", hbErr)
} else {
if changed {
heartbeatTicker.Reset(r.Config.HeartbeatInterval.Duration())
}
tryStartWebSocket()
}
}
}
}
func (r *Runner) performHeartbeatCycle(ctx context.Context, nodeID string, startup bool) (bool, error) {
r.refreshOpenrestyHealth(ctx)
payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
heartbeatResult, err := r.HeartbeatService.Heartbeat(ctx, payload)
if err != nil {
return false, err
}
r.ackObservabilityWindows(ackWindows)
if heartbeatResult == nil {
heartbeatResult = &protocol.HeartbeatResult{}
}
mode := "periodic"
if startup {
mode = "startup"
}
slog.Debug("agent heartbeat succeeded", "mode", mode, "node_id", nodeID)
changed := r.applySettings(heartbeatResult.AgentSettings)
r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
if startup {
if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent startup sync failed", "error", err)
} else {
slog.Debug("agent startup sync completed")
}
} else if err = r.SyncService.SyncOnce(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent sync failed", "error", err)
}
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
return changed, nil
}
func (r *Runner) shouldUseWebSocket() bool {
enabled := r.WebSocketService != nil && r.websocketUpgradeEnabled && r.hasAccessToken()
slog.Debug("agent ws upgrade eligibility checked", "enabled", enabled, "server_enabled", r.websocketUpgradeEnabled, "url", r.websocketURL())
return enabled
}
func (r *Runner) websocketURL() string {
if r.WebSocketService == nil {
return ""
}
return r.WebSocketService.URL()
}
func (r *Runner) startWebSocket(ctx context.Context, nodeID string) (<-chan error, error) {
if r.WebSocketService == nil {
return nil, errors.New("websocket service is not configured")
}
conn, err := r.WebSocketService.Connect(ctx)
if err != nil {
return nil, err
}
done := make(chan error, 1)
go func() {
defer func() {
_ = conn.Close()
}()
done <- r.runWebSocket(ctx, nodeID, conn)
}()
return done, nil
}
type agentWSHandler struct {
runner *Runner
conn protocol.WebSocketConnection
nodeID string
statusTicker *time.Ticker
}
func (h *agentWSHandler) OnConnect(ctx context.Context) error {
return h.runner.sendWebSocketStatus(ctx, h.nodeID, h.conn)
}
func (h *agentWSHandler) HandleMessage(ctx context.Context, msg wsclient.WSMessage) error {
var payloadBytes []byte
if msg.Payload != nil {
payloadBytes = []byte(msg.Payload)
}
protoMsg := protocol.WSMessage{
Type: msg.Type,
Payload: payloadBytes,
}
changed, err := h.runner.handleWebSocketMessage(ctx, protoMsg, h.conn)
if err != nil {
return err
}
if changed {
h.statusTicker.Reset(h.runner.Config.HeartbeatInterval.Duration())
}
return nil
}
func (h *agentWSHandler) OnClose(err error) {
slog.Error("agent ws receive failed", "error", err)
}
func (r *Runner) runWebSocket(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
slog.Debug("agent ws connected", "url", conn.URL(), "node_id", nodeID)
statusTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
defer statusTicker.Stop()
childCtx, cancel := context.WithCancel(ctx)
defer cancel()
// Start status ticker sender in background
go func() {
for {
select {
case <-childCtx.Done():
return
case <-statusTicker.C:
if err := r.sendWebSocketStatus(childCtx, nodeID, conn); err != nil {
slog.Error("agent ws send status failed", "error", err)
_ = conn.Close()
return
}
}
}
}()
wsConn, ok := conn.(*wsclient.Connection)
if !ok {
return errors.New("invalid websocket connection type")
}
return wsConn.RunReceiveLoop(childCtx, &agentWSHandler{
runner: r,
conn: conn,
nodeID: nodeID,
statusTicker: statusTicker,
})
}
func (r *Runner) sendWebSocketStatus(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
r.refreshOpenrestyHealth(ctx)
payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
if err := conn.SendStatus(payload); err != nil {
return err
}
r.ackObservabilityWindows(ackWindows)
return nil
}
func (r *Runner) handleWebSocketMessage(ctx context.Context, message protocol.WSMessage, conn protocol.WebSocketConnection) (bool, error) {
switch message.Type {
case protocol.WSMessageTypeSettings:
var settings protocol.AgentSettings
if err := json.Unmarshal(message.Payload, &settings); err != nil {
slog.Debug("agent ws settings decode failed", "error", err)
return false, nil
}
changed := r.applySettings(&settings)
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
if !r.websocketUpgradeEnabled {
slog.Debug("agent ws disabled by server settings; falling back to http heartbeat")
return changed, errors.New("websocket upgrade disabled by server")
}
return changed, nil
case protocol.WSMessageTypeActiveConfig:
var target protocol.ActiveConfigMeta
if err := json.Unmarshal(message.Payload, &target); err != nil {
slog.Debug("agent ws active config decode failed", "error", err)
return false, nil
}
slog.Debug("agent ws active config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
if err := r.SyncService.SyncOnce(ctx, &target); err != nil {
r.recordSyncError(err)
slog.Error("agent ws triggered sync failed", "version", target.Version, "error", err)
}
return false, nil
case protocol.WSMessageTypeForceSyncConfig:
var target protocol.ActiveConfigMeta
if err := json.Unmarshal(message.Payload, &target); err != nil {
slog.Debug("agent ws force sync config decode failed", "error", err)
return false, nil
}
slog.Debug("agent ws force sync config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
if err := r.SyncService.ForceSyncOnce(ctx, &target); err != nil {
r.recordSyncError(err)
slog.Error("agent ws triggered force sync failed", "version", target.Version, "error", err)
}
return false, nil
case protocol.WSMessageTypeWAFIPGroups:
var groups []protocol.WAFIPGroup
if err := json.Unmarshal(message.Payload, &groups); err != nil {
slog.Debug("agent ws waf ip groups decode failed", "error", err)
return false, nil
}
r.applyWAFIPGroups(ctx, groups)
return false, nil
case protocol.WSMessageTypePing:
slog.Debug("agent ws ping received")
return false, conn.SendPong()
case protocol.WSMessageTypePong:
slog.Debug("agent ws pong received")
return false, nil
default:
slog.Debug("agent ws unsupported message type", "type", message.Type)
return false, nil
}
}
type webSocketBackoff struct {
delays []time.Duration
index int
}
func newWebSocketBackoff() *webSocketBackoff {
return &webSocketBackoff{
delays: []time.Duration{
time.Second,
2 * time.Second,
5 * time.Second,
10 * time.Second,
30 * time.Second,
},
}
}
func (backoff *webSocketBackoff) Next() time.Duration {
if backoff == nil || len(backoff.delays) == 0 {
return 30 * time.Second
}
if backoff.index >= len(backoff.delays) {
return backoff.delays[len(backoff.delays)-1]
}
delay := backoff.delays[backoff.index]
backoff.index++
return delay
}
func (backoff *webSocketBackoff) Reset() {
if backoff != nil {
backoff.index = 0
}
}
func (r *Runner) hasAccessToken() bool {
return strings.TrimSpace(r.Config.AccessToken) != ""
}
func (r *Runner) applySettings(settings *protocol.AgentSettings) bool {
if settings == nil {
return false
}
changed := false
if settings.HeartbeatInterval > 0 {
newInterval := config.MillisecondDuration(time.Duration(settings.HeartbeatInterval) * time.Millisecond)
if newInterval != r.Config.HeartbeatInterval {
slog.Info("agent heartbeat interval updated", "from", r.Config.HeartbeatInterval, "to", newInterval)
r.Config.HeartbeatInterval = newInterval
changed = true
}
}
if settings.WebsocketUpgradeEnabled != r.websocketUpgradeEnabled {
slog.Debug("agent websocket upgrade setting updated", "from", r.websocketUpgradeEnabled, "to", settings.WebsocketUpgradeEnabled)
}
r.websocketUpgradeEnabled = settings.WebsocketUpgradeEnabled
r.autoUpdate = settings.AutoUpdate
r.updateNow = settings.UpdateNow
r.updateRepo = strings.TrimSpace(settings.UpdateRepo)
r.updateChan = strings.TrimSpace(settings.UpdateChannel)
r.updateTag = strings.TrimSpace(settings.UpdateTag)
r.restartOpenrestyNow = settings.RestartOpenrestyNow
return changed
}
func (r *Runner) tryRestartOpenresty(ctx context.Context) {
if !r.restartOpenrestyNow {
return
}
r.restartOpenrestyNow = false
if r.RuntimeManager == nil {
return
}
slog.Info("agent openresty restart requested by server")
if err := r.RuntimeManager.Restart(ctx); err != nil {
slog.Error("agent openresty restart failed", "error", err)
r.recordOpenrestyUnhealthy(err, false)
return
}
slog.Info("agent openresty restart succeeded")
r.recordOpenrestyHealthy()
}
func (r *Runner) tryAutoUpdate(ctx context.Context) {
force := r.updateNow
shouldCheck := r.autoUpdate || force
r.updateNow = false
r.updateTag = strings.TrimSpace(r.updateTag)
if !shouldCheck || r.Updater == nil || r.updateRepo == "" {
return
}
channel := "stable"
if force && r.updateChan != "" {
channel = r.updateChan
}
if err := r.Updater.CheckAndUpdate(ctx, r.updateRepo, UpdateOptions{
Channel: channel,
TagName: r.updateTag,
Force: force,
}); err != nil {
slog.Error("agent update check failed", "error", err)
}
if force {
r.updateTag = ""
r.updateChan = ""
}
}
func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
if strings.TrimSpace(r.Config.DiscoveryToken) == "" {
return errors.New("agent_token 为空且未配置 discovery_token")
}
slog.Info("agent discovery registration started")
response, err := r.HeartbeatService.Register(ctx, r.nodePayload(*nodeID))
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")
}
snapshot, err := r.StateStore.Load()
if 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 {
return err
}
r.HeartbeatService.SetToken(response.AccessToken)
if r.WebSocketService != nil {
r.WebSocketService.SetToken(response.AccessToken)
}
*nodeID = response.NodeID
slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
r.refreshOpenrestyHealth(ctx)
payload, ackWindows := r.prepareHeartbeatPayload(*nodeID)
heartbeatResult, heartbeatErr := r.HeartbeatService.Heartbeat(ctx, payload)
if heartbeatErr != nil {
slog.Error("agent post-register heartbeat failed", "error", heartbeatErr)
return nil
}
r.ackObservabilityWindows(ackWindows)
if heartbeatResult == nil {
heartbeatResult = &protocol.HeartbeatResult{}
}
r.applySettings(heartbeatResult.AgentSettings)
r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
r.recordSyncError(err)
slog.Error("agent post-register startup sync failed", "error", err)
} else {
slog.Debug("agent post-register startup sync completed")
}
r.tryRestartOpenresty(ctx)
r.tryAutoUpdate(ctx)
return nil
}
func (r *Runner) recordSyncError(err error) {
if err == nil || r.StateStore == nil {
return
}
snapshot, loadErr := r.StateStore.Load()
if loadErr != nil {
slog.Error("load state before recording sync error failed", "error", loadErr)
return
}
snapshot.LastError = err.Error()
slog.Warn("recording sync error into state", "error", snapshot.LastError)
if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
slog.Error("save state after sync error failed", "error", saveErr)
}
}
func (r *Runner) refreshOpenrestyHealth(ctx context.Context) {
if r.RuntimeManager == nil || r.StateStore == nil {
return
}
if err := r.RuntimeManager.CheckHealth(ctx); err != nil {
if strings.Contains(err.Error(), "openresty config not exists") {
return
}
r.recordOpenrestyUnhealthy(err, true)
return
}
r.recordOpenrestyHealthy()
}
func (r *Runner) recordOpenrestyHealthy() {
if r.StateStore == nil {
return
}
snapshot, err := r.StateStore.Load()
if err != nil {
slog.Error("load state before recording openresty health failed", "error", err)
return
}
if snapshot.OpenrestyStatus == protocol.OpenrestyStatusHealthy && strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
return
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = ""
if err = r.StateStore.Save(snapshot); err != nil {
slog.Error("save state after recording openresty health failed", "error", err)
}
}
func (r *Runner) recordOpenrestyUnhealthy(err error, fallbackOnly bool) {
if err == nil || r.StateStore == nil {
return
}
snapshot, loadErr := r.StateStore.Load()
if loadErr != nil {
slog.Error("load state before recording openresty error failed", "error", loadErr)
return
}
message := strings.TrimSpace(err.Error())
if !fallbackOnly || strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
snapshot.OpenrestyMessage = message
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
slog.Error("save state after recording openresty error failed", "error", saveErr)
}
}
func (r *Runner) nodePayload(nodeID string) protocol.NodePayload {
snapshot, _ := r.StateStore.Load()
openrestyStatus := strings.TrimSpace(snapshot.OpenrestyStatus)
if openrestyStatus == "" {
openrestyStatus = protocol.OpenrestyStatusUnknown
}
profile := observability.BuildProfile(r.Config, r.StateStore)
managedOpenRestyMetrics := observability.CollectManagedOpenRestyMetrics(r.Config)
trafficReport, accessLogs, fallbackMetrics := observability.BuildTrafficObservability(r.Config, r.StateStore, managedOpenRestyMetrics)
if managedOpenRestyMetrics == nil {
managedOpenRestyMetrics = fallbackMetrics
}
metricSnapshot := observability.BuildSnapshot(r.Config, r.StateStore)
openrestyObservation := observability.BuildOpenrestyObservation(managedOpenRestyMetrics)
healthEvents := observability.BuildHealthEvents(snapshot)
payload := protocol.NodePayload{
NodeID: nodeID,
Name: r.Config.NodeName,
IP: r.Config.NodeIP,
Version: r.Config.Version,
ExtVersion: r.Config.ExtVersion,
CurrentVersion: snapshot.CurrentVersion,
LastError: snapshot.LastError,
OpenrestyStatus: openrestyStatus,
OpenrestyMessage: snapshot.OpenrestyMessage,
Profile: profile,
Snapshot: metricSnapshot,
OpenrestyObservation: openrestyObservation,
TrafficReport: trafficReport,
AccessLogs: accessLogs,
HealthEvents: healthEvents,
}
if r.SyncService != nil {
checksums, err := r.SyncService.WAFIPGroupChecksums()
if err != nil {
slog.Debug("load local waf ip group checksums failed", "error", err)
} else if len(checksums) > 0 {
payload.WAFIPGroupChecksums = checksums
}
}
return payload
}
func (r *Runner) applyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) {
if len(groups) == 0 || r.SyncService == nil {
return
}
if err := r.SyncService.ApplyWAFIPGroups(ctx, groups); err != nil {
r.recordSyncError(err)
slog.Error("agent apply waf ip groups failed", "error", err)
}
}
func (r *Runner) prepareHeartbeatPayload(nodeID string) (protocol.NodePayload, []int64) {
payload := r.nodePayload(nodeID)
if r.ObservabilityBuffer == nil || (payload.Snapshot == nil && payload.TrafficReport == nil && len(payload.AccessLogs) == 0) {
return payload, nil
}
now := time.Now().UTC()
retainAfterUnix := now.Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
windowStartedAtUnix := state.ObservabilityWindowStartedAt(payload.Snapshot, payload.OpenrestyObservation, payload.TrafficReport)
if windowStartedAtUnix <= 0 {
return payload, nil
}
record := state.ObservabilityBufferRecord{
WindowStartedAtUnix: windowStartedAtUnix,
Snapshot: payload.Snapshot,
OpenrestyObservation: payload.OpenrestyObservation,
TrafficReport: payload.TrafficReport,
AccessLogs: payload.AccessLogs,
QueuedAtUnix: now.Unix(),
}
if err := r.ObservabilityBuffer.Upsert(record, retainAfterUnix); err != nil {
slog.Error("upsert observability buffer failed", "error", err)
return payload, nil
}
records, err := r.ObservabilityBuffer.Replayable(windowStartedAtUnix, retainAfterUnix)
if err != nil {
slog.Error("load replayable observability buffer failed", "error", err)
return payload, []int64{windowStartedAtUnix}
}
ackWindows := make([]int64, 0, len(records)+1)
buffered := make([]protocol.BufferedObservabilityRecord, 0, len(records))
for _, item := range records {
if item.WindowStartedAtUnix <= 0 {
continue
}
buffered = append(buffered, protocol.BufferedObservabilityRecord{
WindowStartedAtUnix: item.WindowStartedAtUnix,
Snapshot: item.Snapshot,
OpenrestyObservation: item.OpenrestyObservation,
TrafficReport: item.TrafficReport,
AccessLogs: item.AccessLogs,
})
ackWindows = append(ackWindows, item.WindowStartedAtUnix)
}
payload.BufferedObservability = buffered
ackWindows = append(ackWindows, windowStartedAtUnix)
return payload, ackWindows
}
func (r *Runner) ackObservabilityWindows(windowStartedAtUnix []int64) {
if r.ObservabilityBuffer == nil || len(windowStartedAtUnix) == 0 {
return
}
retainAfterUnix := time.Now().UTC().Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
if err := r.ObservabilityBuffer.Ack(windowStartedAtUnix, retainAfterUnix); err != nil {
slog.Error("ack observability buffer failed", "error", err)
}
}
+640
View File
@@ -0,0 +1,640 @@
package agent
import (
"context"
"encoding/json"
"errors"
"os"
"path/filepath"
"sync"
"testing"
"time"
"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"
)
type fakeHeartbeatService struct {
mu sync.Mutex
registerCalls int
heartbeatCalls int
registerErr error
registerResp *protocol.RegisterNodeResponse
heartbeatErrs []error
heartbeatResults []*protocol.HeartbeatResult
heartbeatPayloads []protocol.NodePayload
onHeartbeat func(int)
lastToken string
}
func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.registerCalls++
return f.registerResp, f.registerErr
}
func (f *fakeHeartbeatService) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) {
f.mu.Lock()
f.heartbeatCalls++
callIndex := f.heartbeatCalls
f.heartbeatPayloads = append(f.heartbeatPayloads, payload)
var err error
if len(f.heartbeatErrs) >= callIndex {
err = f.heartbeatErrs[callIndex-1]
}
var result *protocol.HeartbeatResult
if len(f.heartbeatResults) >= callIndex {
result = f.heartbeatResults[callIndex-1]
}
onHeartbeat := f.onHeartbeat
f.mu.Unlock()
if onHeartbeat != nil {
onHeartbeat(callIndex)
}
return result, err
}
func (f *fakeHeartbeatService) SetToken(token string) {
f.mu.Lock()
defer f.mu.Unlock()
f.lastToken = token
}
type fakeSyncService struct {
mu sync.Mutex
startupErr error
syncOnceErr error
startupCalls int
syncOnceCalls int
lastTarget *protocol.ActiveConfigMeta
onSyncOnceCall func(int)
wafChecksums map[string]string
wafGroups []protocol.WAFIPGroup
}
type fakeRuntimeManager struct {
mu sync.Mutex
healthErr error
restartErr error
restartCalls int
clearHealthOnRestart bool
}
func (f *fakeRuntimeManager) CheckHealth(ctx context.Context) error {
f.mu.Lock()
defer f.mu.Unlock()
return f.healthErr
}
func (f *fakeRuntimeManager) Restart(ctx context.Context) error {
f.mu.Lock()
defer f.mu.Unlock()
f.restartCalls++
if f.clearHealthOnRestart && f.restartErr == nil {
f.healthErr = nil
}
return f.restartErr
}
func (f *fakeSyncService) SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error {
f.mu.Lock()
defer f.mu.Unlock()
f.startupCalls++
return f.startupErr
}
func (f *fakeSyncService) SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error {
f.mu.Lock()
f.syncOnceCalls++
if target != nil {
copied := *target
f.lastTarget = &copied
}
callIndex := f.syncOnceCalls
callback := f.onSyncOnceCall
f.mu.Unlock()
if callback != nil {
callback(callIndex)
}
return f.syncOnceErr
}
func (f *fakeSyncService) ForceSyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error {
f.mu.Lock()
f.syncOnceCalls++
if target != nil {
copied := *target
f.lastTarget = &copied
}
callIndex := f.syncOnceCalls
callback := f.onSyncOnceCall
f.mu.Unlock()
if callback != nil {
callback(callIndex)
}
return f.syncOnceErr
}
func (f *fakeSyncService) WAFIPGroupChecksums() (map[string]string, error) {
if f.wafChecksums == nil {
return map[string]string{}, nil
}
return f.wafChecksums, nil
}
func (f *fakeSyncService) ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error {
f.mu.Lock()
defer f.mu.Unlock()
f.wafGroups = append(f.wafGroups, groups...)
return nil
}
type fakeWebSocketConnection struct {
pongCalls int
}
func (f *fakeWebSocketConnection) URL() string {
return "ws://127.0.0.1/api/v1/agent/ws"
}
func (f *fakeWebSocketConnection) SendStatus(payload protocol.NodePayload) error {
return nil
}
func (f *fakeWebSocketConnection) SendPong() error {
f.pongCalls++
return nil
}
func (f *fakeWebSocketConnection) Receive() (protocol.WSMessage, error) {
return protocol.WSMessage{}, errors.New("not implemented")
}
func (f *fakeWebSocketConnection) Close() error {
return nil
}
func TestRunnerKeepsHeartbeatWhenStartupSyncFails(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
heartbeatService := &fakeHeartbeatService{
heartbeatResults: []*protocol.HeartbeatResult{{}},
onHeartbeat: func(callCount int) {
if callCount >= 2 {
cancel()
}
},
}
syncService := &fakeSyncService{
startupErr: errors.New("当前没有激活版本,保持当前 OpenResty 配置"),
}
runner := &Runner{
Config: &config.Config{
AccessToken: "agent-token",
NodeName: "edge-01",
NodeIP: "10.0.0.8",
Version: config.Version,
ExtVersion: "1.27.1.2",
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: syncService,
}
err := runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if heartbeatService.registerCalls != 0 {
t.Fatalf("expected no discovery register call, got %d", heartbeatService.registerCalls)
}
if heartbeatService.heartbeatCalls < 2 {
t.Fatalf("expected heartbeat loop to continue, got %d heartbeat calls", heartbeatService.heartbeatCalls)
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.LastError != "当前没有激活版本,保持当前 OpenResty 配置" {
t.Fatalf("expected startup sync error to be recorded, got %q", snapshot.LastError)
}
}
func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
heartbeatService := &fakeHeartbeatService{
registerErr: errors.New("register timeout"),
heartbeatErrs: []error{errors.New("heartbeat timeout")},
heartbeatResults: []*protocol.HeartbeatResult{
{},
},
}
syncService := &fakeSyncService{
syncOnceErr: errors.New("openresty reload failed"),
onSyncOnceCall: func(callCount int) {
if callCount >= 1 {
cancel()
}
},
}
runner := &Runner{
Config: &config.Config{
AccessToken: "agent-token",
NodeName: "edge-01",
NodeIP: "10.0.0.8",
Version: config.Version,
ExtVersion: "1.27.1.2",
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: syncService,
}
err := runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if heartbeatService.registerCalls != 0 {
t.Fatalf("expected no register attempt, got %d", heartbeatService.registerCalls)
}
if syncService.syncOnceCalls == 0 {
t.Fatal("expected sync loop to continue after heartbeat/register errors")
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.LastError != "openresty reload failed" {
t.Fatalf("expected sync error to be recorded, got %q", snapshot.LastError)
}
}
func TestRunnerReportsOpenrestyHealthAndExecutesRestart(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
if err := stateStore.Save(&state.Snapshot{
OpenrestyStatus: protocol.OpenrestyStatusUnhealthy,
OpenrestyMessage: "docker run openresty failed: bind 80 already allocated",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
heartbeatService := &fakeHeartbeatService{
heartbeatResults: []*protocol.HeartbeatResult{{
AgentSettings: &protocol.AgentSettings{RestartOpenrestyNow: true},
}},
onHeartbeat: func(callCount int) {
if callCount >= 1 {
cancel()
}
},
}
runtimeManager := &fakeRuntimeManager{
healthErr: errors.New("docker openresty container is not running"),
clearHealthOnRestart: true,
}
runner := &Runner{
Config: &config.Config{
AccessToken: "agent-token",
NodeName: "edge-01",
NodeIP: "10.0.0.8",
Version: config.Version,
ExtVersion: "1.27.1.2",
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: &fakeSyncService{},
RuntimeManager: runtimeManager,
}
err := runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if len(heartbeatService.heartbeatPayloads) == 0 {
t.Fatal("expected at least one heartbeat payload")
}
payload := heartbeatService.heartbeatPayloads[0]
if payload.OpenrestyStatus != protocol.OpenrestyStatusUnhealthy {
t.Fatalf("expected unhealthy openresty status in heartbeat payload, got %q", payload.OpenrestyStatus)
}
if payload.OpenrestyMessage != "docker run openresty failed: bind 80 already allocated" {
t.Fatalf("unexpected openresty message: %q", payload.OpenrestyMessage)
}
if runtimeManager.restartCalls != 1 {
t.Fatalf("expected one openresty restart attempt, got %d", runtimeManager.restartCalls)
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy || snapshot.OpenrestyMessage != "" {
t.Fatal("expected restart success to mark openresty healthy")
}
}
func TestRunnerHeartbeatPayloadIncludesObservabilityExtensions(t *testing.T) {
tempDir := t.TempDir()
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
if err := stateStore.Save(&state.Snapshot{
NodeID: "node-observe",
CurrentVersion: "20260314-001",
LastError: "sync failed",
OpenrestyStatus: protocol.OpenrestyStatusUnhealthy,
OpenrestyMessage: "reload failed",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
runner := &Runner{
Config: &config.Config{
NodeName: "edge-observe-1",
NodeIP: "10.0.0.51",
Version: config.Version,
ExtVersion: "1.27.1.2",
DataDir: tempDir,
RouteConfigPath: filepath.Join(tempDir, "conf.d", "openflare_routes.conf"),
AccessLogPath: filepath.Join(tempDir, "var", "log", "openflare", "access.log"),
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
},
StateStore: stateStore,
}
if err := os.MkdirAll(filepath.Dir(runner.Config.AccessLogPath), 0o755); err != nil {
t.Fatalf("failed to prepare access log dir: %v", err)
}
if err := os.WriteFile(
runner.Config.AccessLogPath,
[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
0o644,
); err != nil {
t.Fatalf("failed to prepare access log: %v", err)
}
firstPayload := runner.nodePayload("node-observe")
if firstPayload.Profile == nil {
t.Fatal("expected first heartbeat payload to include system profile")
}
if firstPayload.Snapshot == nil {
t.Fatal("expected first heartbeat payload to include metric snapshot")
}
if firstPayload.TrafficReport == nil || firstPayload.TrafficReport.RequestCount != 1 {
t.Fatalf("expected first heartbeat payload to include traffic report, got %+v", firstPayload.TrafficReport)
}
if len(firstPayload.AccessLogs) != 1 || firstPayload.AccessLogs[0].Path != "/" {
t.Fatalf("expected first heartbeat payload to include access logs, got %+v", firstPayload.AccessLogs)
}
if len(firstPayload.HealthEvents) != 2 {
t.Fatalf("expected health events for openresty and sync error, got %+v", firstPayload.HealthEvents)
}
secondPayload := runner.nodePayload("node-observe")
if secondPayload.Profile != nil {
t.Fatal("expected unchanged profile to be omitted on subsequent heartbeat")
}
if secondPayload.Snapshot == nil {
t.Fatal("expected metric snapshot to continue reporting on subsequent heartbeat")
}
if secondPayload.TrafficReport != nil {
t.Fatalf("expected unchanged traffic window to be omitted on subsequent heartbeat, got %+v", secondPayload.TrafficReport)
}
if len(secondPayload.AccessLogs) != 0 {
t.Fatalf("expected unchanged access log delta to be omitted on subsequent heartbeat, got %+v", secondPayload.AccessLogs)
}
}
func TestRunnerReplaysBufferedObservabilityAfterHeartbeatRecovery(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
tempDir := t.TempDir()
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
bufferStore := state.NewObservabilityBufferStore(filepath.Join(tempDir, "observability-buffer.json"))
nowUnix := time.Now().UTC().Unix()
bufferWindow := nowUnix - (nowUnix % 60) - 60
if err := bufferStore.Upsert(state.ObservabilityBufferRecord{
WindowStartedAtUnix: bufferWindow,
Snapshot: &protocol.NodeMetricSnapshot{CapturedAtUnix: bufferWindow + 5, CPUUsagePercent: 30},
TrafficReport: &protocol.NodeTrafficReport{WindowStartedAtUnix: bufferWindow, WindowEndedAtUnix: bufferWindow + 60, RequestCount: 8},
QueuedAtUnix: bufferWindow + 60,
}, 0); err != nil {
t.Fatalf("failed to seed observability buffer: %v", err)
}
heartbeatService := &fakeHeartbeatService{
heartbeatErrs: []error{errors.New("server offline"), nil},
heartbeatResults: []*protocol.HeartbeatResult{{}, {}},
onHeartbeat: func(callCount int) {
if callCount >= 2 {
cancel()
}
},
}
runner := &Runner{
Config: &config.Config{
AccessToken: "agent-token",
NodeName: "edge-buffer-01",
NodeIP: "10.0.0.52",
Version: config.Version,
ExtVersion: "1.27.1.2",
DataDir: tempDir,
RouteConfigPath: filepath.Join(tempDir, "conf.d", "openflare_routes.conf"),
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
ObservabilityReplayMinutes: 15,
},
StateStore: stateStore,
ObservabilityBuffer: bufferStore,
HeartbeatService: heartbeatService,
SyncService: &fakeSyncService{},
}
if err := os.MkdirAll(filepath.Dir(runner.Config.RouteConfigPath), 0o755); err != nil {
t.Fatalf("failed to prepare route config dir: %v", err)
}
if err := os.WriteFile(
filepath.Join(filepath.Dir(runner.Config.RouteConfigPath), "openflare_access.log"),
[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
0o644,
); err != nil {
t.Fatalf("failed to prepare access log: %v", err)
}
runErr := runner.Run(ctx)
if runErr != context.Canceled {
t.Fatalf("expected run to stop by context cancellation, got %v", runErr)
}
if len(heartbeatService.heartbeatPayloads) != 2 {
t.Fatalf("expected two heartbeat payloads, got %d", len(heartbeatService.heartbeatPayloads))
}
secondPayload := heartbeatService.heartbeatPayloads[1]
if len(secondPayload.BufferedObservability) != 1 {
t.Fatalf("expected second heartbeat to replay one buffered observation, got %+v", secondPayload.BufferedObservability)
}
if len(secondPayload.BufferedObservability[0].AccessLogs) != 0 {
t.Fatalf("expected seeded buffered observation to keep empty access logs, got %+v", secondPayload.BufferedObservability[0].AccessLogs)
}
replayable, err := bufferStore.Replayable(0, 0)
if err != nil {
t.Fatalf("Replayable after recovery failed: %v", err)
}
if len(replayable) != 0 {
t.Fatalf("expected buffer to be acked after successful heartbeat, got %+v", replayable)
}
}
func TestRunnerDiscoveryRegisterUpdatesTokenAndNodeID(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
heartbeatService := &fakeHeartbeatService{
registerResp: &protocol.RegisterNodeResponse{
NodeID: "node-server-assigned",
AccessToken: "agent-token-issued",
Name: "edge-01",
},
heartbeatResults: []*protocol.HeartbeatResult{{}},
onHeartbeat: func(callCount int) {
if callCount >= 1 {
cancel()
}
},
}
syncService := &fakeSyncService{}
configPath := filepath.Join(t.TempDir(), "agent.json")
if err := os.WriteFile(configPath, []byte(`{"server_url":"http://127.0.0.1:3000","discovery_token":"discovery-token","node_name":"edge-01","node_ip":"10.0.0.8"}`), 0o644); err != nil {
t.Fatalf("failed to seed config file: %v", err)
}
cfg, err := config.Load(configPath)
if err != nil {
t.Fatalf("failed to load config: %v", err)
}
runner := &Runner{
Config: &config.Config{
ServerURL: cfg.ServerURL,
DiscoveryToken: cfg.DiscoveryToken,
NodeName: cfg.NodeName,
NodeIP: cfg.NodeIP,
Version: config.Version,
ExtVersion: "1.27.1.2",
HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
},
StateStore: stateStore,
HeartbeatService: heartbeatService,
SyncService: syncService,
}
runner.Config = cfg
runner.Config.Version = config.Version
runner.Config.ExtVersion = "1.27.1.2"
runner.Config.HeartbeatInterval = config.MillisecondDuration(10 * time.Millisecond)
err = runner.Run(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("expected context cancellation, got %v", err)
}
if heartbeatService.registerCalls == 0 {
t.Fatal("expected discovery register to be attempted")
}
if heartbeatService.lastToken != "agent-token-issued" {
t.Fatalf("expected client token to be updated, got %q", heartbeatService.lastToken)
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.NodeID != "node-server-assigned" {
t.Fatalf("expected node id to be replaced, got %q", snapshot.NodeID)
}
if runner.Config.AccessToken != "agent-token-issued" || runner.Config.DiscoveryToken != "" {
t.Fatal("expected config token rotation to complete")
}
}
func TestRunnerHandlesWebSocketActiveConfigMessage(t *testing.T) {
syncService := &fakeSyncService{}
runner := &Runner{SyncService: syncService}
payload, err := json.Marshal(protocol.ActiveConfigMeta{
Version: "20260529-001",
Checksum: "checksum-ws",
})
if err != nil {
t.Fatalf("marshal active config: %v", err)
}
changed, err := runner.handleWebSocketMessage(context.Background(), protocol.WSMessage{
Type: protocol.WSMessageTypeActiveConfig,
Payload: payload,
}, &fakeWebSocketConnection{})
if err != nil {
t.Fatalf("handle websocket active config: %v", err)
}
if changed {
t.Fatal("active config message should not change heartbeat interval")
}
if syncService.syncOnceCalls != 1 {
t.Fatalf("expected one sync call, got %d", syncService.syncOnceCalls)
}
if syncService.lastTarget == nil || syncService.lastTarget.Version != "20260529-001" || syncService.lastTarget.Checksum != "checksum-ws" {
t.Fatalf("unexpected sync target: %+v", syncService.lastTarget)
}
}
func TestRunnerHandlesWebSocketSettingsDisabled(t *testing.T) {
runner := &Runner{
Config: &config.Config{
HeartbeatInterval: config.MillisecondDuration(10 * time.Second),
},
websocketUpgradeEnabled: true,
}
payload, err := json.Marshal(protocol.AgentSettings{
HeartbeatInterval: 15000,
WebsocketUpgradeEnabled: false,
})
if err != nil {
t.Fatalf("marshal settings: %v", err)
}
changed, err := runner.handleWebSocketMessage(context.Background(), protocol.WSMessage{
Type: protocol.WSMessageTypeSettings,
Payload: payload,
}, &fakeWebSocketConnection{})
if err == nil {
t.Fatal("expected disabled websocket setting to request fallback")
}
if !changed {
t.Fatal("expected heartbeat interval change to be reported")
}
if runner.websocketUpgradeEnabled {
t.Fatal("expected websocket upgrade to be disabled")
}
}
func TestWebSocketBackoffSequence(t *testing.T) {
backoff := newWebSocketBackoff()
expected := []time.Duration{
time.Second,
2 * time.Second,
5 * time.Second,
10 * time.Second,
30 * time.Second,
30 * time.Second,
}
for _, want := range expected {
if got := backoff.Next(); got != want {
t.Fatalf("unexpected backoff: got %s want %s", got, want)
}
}
backoff.Reset()
if got := backoff.Next(); got != time.Second {
t.Fatalf("expected reset backoff to return 1s, got %s", got)
}
}
+463
View File
@@ -0,0 +1,463 @@
package config
import (
"context"
"encoding/json"
"errors"
"fmt"
"net"
"os"
pathpkg "path"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/Rain-kl/Wavelet/pkg/geoip"
"github.com/Rain-kl/Wavelet/pkg/geoip/iputil"
"github.com/Rain-kl/Wavelet/pkg/utils"
)
const (
defaultMainConfigRelativePath = "etc/nginx/nginx.conf"
defaultRouteConfigRelativePath = "etc/nginx/conf.d/openflare_routes.conf"
defaultCertDirRelativePath = "etc/nginx/certs"
defaultLuaDirRelativePath = "etc/nginx/lua"
defaultRuntimeConfigDirRelativePath = "etc/openflare"
defaultPagesDirRelativePath = "var/lib/openflare/pages"
defaultMMDBRelativePath = "etc/openflare/GeoLite2-Country.mmdb"
defaultAccessLogRelativePath = "var/log/openflare/access.log"
defaultStateRelativePath = "var/lib/openflare/agent-state.json"
defaultObservabilityBufferRelativePath = "var/lib/openflare/observability-buffer.json"
defaultOpenRestyObservabilityPort = 18081
defaultObservabilityReplayMinutes = 15
defaultMMDBUpdateInterval = 24 * time.Hour
defaultMMDBDownloadURL = "https://raw.githubusercontent.com/Loyalsoldier/geoip/release/GeoLite2-Country.mmdb"
)
var (
lookupOutboundIP = geoip.GetOutboundIP
lookupLocalIP = detectLocalNodeIP
)
type Config struct {
ServerURL string `json:"server_url"`
AccessToken string `json:"agent_token"`
DiscoveryToken string `json:"discovery_token"`
NodeName string `json:"node_name"`
NodeIP string `json:"node_ip"`
Version string `json:"-"`
ExtVersion string `json:"-"`
OpenrestyPath string `json:"openresty_path"`
OpenrestyResolvers []string `json:"openresty_resolvers,omitempty"`
DataDir string `json:"data_dir"`
MainConfigPath string `json:"main_config_path"`
RouteConfigPath string `json:"route_config_path"`
AccessLogPath string `json:"access_log_path"`
CertDir string `json:"cert_dir"`
OpenrestyCertDir string `json:"openresty_cert_dir"`
LuaDir string `json:"lua_dir"`
OpenrestyLuaDir string `json:"openresty_lua_dir"`
RuntimeConfigDir string `json:"runtime_config_dir"`
PagesDir string `json:"pages_dir"`
MMDBPath string `json:"mmdb_path"`
MMDBUpdateInterval MillisecondDuration `json:"mmdb_update_interval"`
MMDBDownloadURL string `json:"mmdb_download_url"`
OpenrestyObservabilityPort int `json:"openresty_observability_port"`
ObservabilityBufferPath string `json:"observability_buffer_path"`
ObservabilityReplayMinutes int `json:"observability_replay_minutes"`
StatePath string `json:"state_path"`
HeartbeatInterval MillisecondDuration `json:"heartbeat_interval"`
RequestTimeout MillisecondDuration `json:"request_timeout"`
configPath string
}
type configFile struct {
ServerURL string `json:"server_url"`
AccessToken string `json:"agent_token"`
DiscoveryToken string `json:"discovery_token"`
NodeName string `json:"node_name"`
NodeIP string `json:"node_ip"`
OpenrestyPath string `json:"openresty_path"`
OpenrestyResolvers []string `json:"openresty_resolvers"`
DataDir string `json:"data_dir"`
MainConfigPath string `json:"main_config_path"`
RouteConfigPath string `json:"route_config_path"`
AccessLogPath string `json:"access_log_path"`
CertDir string `json:"cert_dir"`
OpenrestyCertDir string `json:"openresty_cert_dir"`
LuaDir string `json:"lua_dir"`
OpenrestyLuaDir string `json:"openresty_lua_dir"`
RuntimeConfigDir string `json:"runtime_config_dir"`
PagesDir string `json:"pages_dir"`
MMDBPath string `json:"mmdb_path"`
MMDBUpdateInterval MillisecondDuration `json:"mmdb_update_interval"`
MMDBDownloadURL string `json:"mmdb_download_url"`
OpenrestyObservabilityPort int `json:"openresty_observability_port"`
ObservabilityBufferPath string `json:"observability_buffer_path"`
ObservabilityReplayMinutes int `json:"observability_replay_minutes"`
StatePath string `json:"state_path"`
HeartbeatInterval MillisecondDuration `json:"heartbeat_interval"`
RequestTimeout MillisecondDuration `json:"request_timeout"`
}
func Load(path string) (*Config, error) {
data, err := os.ReadFile(path)
if err != nil && !os.IsNotExist(err) {
return nil, err
}
file := &configFile{}
if err == nil {
if err = json.Unmarshal(data, file); err != nil {
return nil, err
}
}
if err != nil && !hasEnvConfig() {
return nil, err
}
cfg := &Config{
ServerURL: file.ServerURL,
AccessToken: file.AccessToken,
DiscoveryToken: file.DiscoveryToken,
NodeName: file.NodeName,
NodeIP: file.NodeIP,
OpenrestyPath: file.OpenrestyPath,
OpenrestyResolvers: append([]string{}, file.OpenrestyResolvers...),
DataDir: file.DataDir,
MainConfigPath: file.MainConfigPath,
RouteConfigPath: file.RouteConfigPath,
AccessLogPath: file.AccessLogPath,
CertDir: file.CertDir,
OpenrestyCertDir: file.OpenrestyCertDir,
LuaDir: file.LuaDir,
OpenrestyLuaDir: file.OpenrestyLuaDir,
RuntimeConfigDir: file.RuntimeConfigDir,
PagesDir: file.PagesDir,
MMDBPath: file.MMDBPath,
MMDBUpdateInterval: file.MMDBUpdateInterval,
MMDBDownloadURL: file.MMDBDownloadURL,
OpenrestyObservabilityPort: file.OpenrestyObservabilityPort,
ObservabilityBufferPath: file.ObservabilityBufferPath,
ObservabilityReplayMinutes: file.ObservabilityReplayMinutes,
StatePath: file.StatePath,
HeartbeatInterval: file.HeartbeatInterval,
RequestTimeout: file.RequestTimeout,
}
cfg.configPath = path
applyEnvOverrides(cfg)
applyDefaults(cfg, filepath.Dir(path))
if err = validate(cfg); err != nil {
return nil, err
}
return cfg, nil
}
func applyDefaults(cfg *Config, baseDir string) {
baseDir = filepath.Clean(baseDir)
cfg.Version = Version
cfg.OpenrestyResolvers = utils.UniqueAndCleanStringSlice(cfg.OpenrestyResolvers)
if cfg.OpenrestyPath == "" {
cfg.OpenrestyPath = "openresty"
}
if cfg.DataDir == "" {
cfg.DataDir = filepath.Join(baseDir, "data")
}
if cfg.NodeName == "" {
cfg.NodeName = detectHostname()
}
if cfg.NodeIP == "" {
cfg.NodeIP = detectNodeIP()
}
if cfg.MainConfigPath == "" {
cfg.MainConfigPath = joinManagedPath(cfg.DataDir, defaultMainConfigRelativePath)
}
if cfg.RouteConfigPath == "" {
cfg.RouteConfigPath = joinManagedPath(cfg.DataDir, defaultRouteConfigRelativePath)
}
if cfg.AccessLogPath == "" {
cfg.AccessLogPath = joinManagedPath(cfg.DataDir, defaultAccessLogRelativePath)
}
if cfg.StatePath == "" {
cfg.StatePath = joinManagedPath(cfg.DataDir, defaultStateRelativePath)
}
if cfg.CertDir == "" {
cfg.CertDir = joinManagedPath(cfg.DataDir, defaultCertDirRelativePath)
}
if cfg.OpenrestyCertDir == "" {
cfg.OpenrestyCertDir = cfg.CertDir
}
if cfg.LuaDir == "" {
cfg.LuaDir = joinManagedPath(cfg.DataDir, defaultLuaDirRelativePath)
}
if cfg.OpenrestyLuaDir == "" {
cfg.OpenrestyLuaDir = cfg.LuaDir
}
if cfg.RuntimeConfigDir == "" {
cfg.RuntimeConfigDir = joinManagedPath(cfg.DataDir, defaultRuntimeConfigDirRelativePath)
}
if cfg.PagesDir == "" {
cfg.PagesDir = joinManagedPath(cfg.DataDir, defaultPagesDirRelativePath)
}
if cfg.MMDBPath == "" {
cfg.MMDBPath = joinManagedPath(cfg.DataDir, defaultMMDBRelativePath)
}
if cfg.MMDBUpdateInterval <= 0 {
cfg.MMDBUpdateInterval = MillisecondDuration(defaultMMDBUpdateInterval)
}
if cfg.MMDBDownloadURL == "" {
cfg.MMDBDownloadURL = defaultMMDBDownloadURL
}
if cfg.OpenrestyObservabilityPort <= 0 {
cfg.OpenrestyObservabilityPort = defaultOpenRestyObservabilityPort
}
if cfg.ObservabilityBufferPath == "" {
cfg.ObservabilityBufferPath = joinManagedPath(cfg.DataDir, defaultObservabilityBufferRelativePath)
}
if cfg.ObservabilityReplayMinutes <= 0 {
cfg.ObservabilityReplayMinutes = defaultObservabilityReplayMinutes
}
if cfg.HeartbeatInterval <= 0 {
cfg.HeartbeatInterval = MillisecondDuration(10 * time.Second)
}
if cfg.RequestTimeout <= 0 {
cfg.RequestTimeout = MillisecondDuration(10 * time.Second)
}
normalizeManagedPaths(cfg)
}
func normalizeManagedPaths(cfg *Config) {
if cfg == nil {
return
}
paths := []*string{
&cfg.DataDir,
&cfg.MainConfigPath,
&cfg.RouteConfigPath,
&cfg.AccessLogPath,
&cfg.CertDir,
&cfg.OpenrestyCertDir,
&cfg.LuaDir,
&cfg.OpenrestyLuaDir,
&cfg.RuntimeConfigDir,
&cfg.PagesDir,
&cfg.StatePath,
&cfg.ObservabilityBufferPath,
&cfg.MMDBPath,
}
for _, p := range paths {
if usesSlashPath(*p) {
*p = filepath.ToSlash(*p)
}
}
}
func hasEnvConfig() bool {
for _, key := range []string{
"OPENFLARE_SERVER_URL",
"OPENFLARE_AGENT_TOKEN",
"OPENFLARE_DISCOVERY_TOKEN",
"OPENFLARE_NODE_NAME",
"OPENFLARE_NODE_IP",
"OPENFLARE_DATA_DIR",
"OPENFLARE_OPENRESTY_PATH",
"OPENFLARE_PAGES_DIR",
"OPENFLARE_HEARTBEAT_INTERVAL",
"OPENFLARE_REQUEST_TIMEOUT",
"OPENFLARE_OPENRESTY_OBSERVABILITY_PORT",
"OPENFLARE_MMDB_PATH",
"OPENFLARE_MMDB_UPDATE_INTERVAL",
"OPENFLARE_MMDB_DOWNLOAD_URL",
} {
if strings.TrimSpace(os.Getenv(key)) != "" {
return true
}
}
return false
}
func applyEnvOverrides(cfg *Config) {
if cfg == nil {
return
}
overrideString := func(key string, target *string) {
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
*target = value
}
}
overrideString("OPENFLARE_SERVER_URL", &cfg.ServerURL)
overrideString("OPENFLARE_AGENT_TOKEN", &cfg.AccessToken)
overrideString("OPENFLARE_DISCOVERY_TOKEN", &cfg.DiscoveryToken)
overrideString("OPENFLARE_NODE_NAME", &cfg.NodeName)
overrideString("OPENFLARE_NODE_IP", &cfg.NodeIP)
overrideString("OPENFLARE_DATA_DIR", &cfg.DataDir)
overrideString("OPENFLARE_OPENRESTY_PATH", &cfg.OpenrestyPath)
overrideString("OPENFLARE_PAGES_DIR", &cfg.PagesDir)
overrideString("OPENFLARE_MMDB_PATH", &cfg.MMDBPath)
overrideString("OPENFLARE_MMDB_DOWNLOAD_URL", &cfg.MMDBDownloadURL)
if value := strings.TrimSpace(os.Getenv("OPENFLARE_HEARTBEAT_INTERVAL")); value != "" {
if duration, err := parseDurationValue(value); err == nil {
cfg.HeartbeatInterval = duration
}
}
if value := strings.TrimSpace(os.Getenv("OPENFLARE_REQUEST_TIMEOUT")); value != "" {
if duration, err := parseDurationValue(value); err == nil {
cfg.RequestTimeout = duration
}
}
if value := strings.TrimSpace(os.Getenv("OPENFLARE_MMDB_UPDATE_INTERVAL")); value != "" {
if duration, err := parseDurationValue(value); err == nil {
cfg.MMDBUpdateInterval = duration
}
}
if value := strings.TrimSpace(os.Getenv("OPENFLARE_OPENRESTY_OBSERVABILITY_PORT")); value != "" {
var port int
if _, err := fmt.Sscanf(value, "%d", &port); err == nil {
cfg.OpenrestyObservabilityPort = port
}
}
}
func parseDurationValue(value string) (MillisecondDuration, error) {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
return 0, nil
}
if parsed, err := time.ParseDuration(trimmed); err == nil {
return MillisecondDuration(parsed), nil
}
ms, err := strconv.ParseInt(trimmed, 10, 64)
if err != nil {
return 0, err
}
return MillisecondDuration(time.Duration(ms) * time.Millisecond), nil
}
func usesSlashPath(path string) bool {
return strings.HasPrefix(path, "/")
}
func joinManagedPath(base string, relative string) string {
if usesSlashPath(base) {
return pathpkg.Join(filepath.ToSlash(base), relative)
}
return filepath.Join(base, relative)
}
func validate(cfg *Config) error {
if cfg.ServerURL == "" {
return errors.New("server_url 不能为空")
}
if strings.TrimSpace(cfg.AccessToken) == "" && strings.TrimSpace(cfg.DiscoveryToken) == "" {
return errors.New("agent_token 和 discovery_token 不能同时为空")
}
if cfg.NodeName == "" {
return errors.New("node_name 不能为空")
}
if cfg.NodeIP == "" {
return errors.New("node_ip 不能为空")
}
if cfg.OpenrestyObservabilityPort <= 0 || cfg.OpenrestyObservabilityPort > 65535 {
return errors.New("openresty_observability_port 必须在 1-65535 之间")
}
if cfg.ObservabilityReplayMinutes <= 0 {
return errors.New("observability_replay_minutes 必须大于 0")
}
if cfg.MMDBUpdateInterval <= 0 {
return errors.New("mmdb_update_interval 必须大于 0")
}
return nil
}
func (cfg *Config) InitialAuthToken() string {
if cfg == nil {
return ""
}
if token := strings.TrimSpace(cfg.AccessToken); token != "" {
return token
}
return strings.TrimSpace(cfg.DiscoveryToken)
}
func (cfg *Config) Save() error {
if cfg == nil {
return errors.New("config 不能为空")
}
if cfg.configPath == "" {
return errors.New("config path 未初始化")
}
data, err := json.MarshalIndent(cfg, "", " ")
if err != nil {
return err
}
return os.WriteFile(cfg.configPath, data, 0o644)
}
func detectHostname() string {
host, err := os.Hostname()
if err != nil {
return ""
}
return strings.TrimSpace(host)
}
func detectNodeIP() string {
if ip := detectOutboundNodeIP(); ip != "" {
return ip
}
return lookupLocalIP()
}
func detectOutboundNodeIP() string {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ip, err := lookupOutboundIP(ctx)
if err != nil || ip == nil {
return ""
}
return ip.String()
}
func detectLocalNodeIP() string {
interfaces, err := net.Interfaces()
if err != nil {
return ""
}
bestIP := ""
bestPriority := -1
for _, iface := range interfaces {
if iface.Flags&net.FlagUp == 0 || iface.Flags&net.FlagLoopback != 0 {
continue
}
addrs, err := iface.Addrs()
if err != nil {
continue
}
for _, addr := range addrs {
ipNet, ok := addr.(*net.IPNet)
if !ok || ipNet.IP == nil || ipNet.IP.IsLoopback() {
continue
}
ipv4 := normalizeIPv4(ipNet.IP)
priority := nodeIPPriority(ipv4)
if priority > bestPriority {
bestIP = ipv4.String()
bestPriority = priority
}
if bestPriority == 2 {
return bestIP
}
}
}
return bestIP
}
func normalizeIPv4(ip net.IP) net.IP {
if ip == nil {
return nil
}
return ip.To4()
}
func nodeIPPriority(ip net.IP) int {
return iputil.Score(ip)
}
+520
View File
@@ -0,0 +1,520 @@
package config
import (
"context"
"encoding/json"
"errors"
"net"
"os"
"path/filepath"
"testing"
"time"
"github.com/Rain-kl/Wavelet/pkg/geoip"
)
func TestLoadDefaultsToManagedBinaryPaths(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
"node_ip": "10.0.0.8",
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.DataDir != filepath.Join(dir, "data") {
t.Fatalf("unexpected data dir: %s", cfg.DataDir)
}
if cfg.OpenrestyPath != "openresty" {
t.Fatalf("unexpected openresty path: %s", cfg.OpenrestyPath)
}
if cfg.MainConfigPath != filepath.Join(dir, "data", defaultMainConfigRelativePath) {
t.Fatalf("unexpected main config path: %s", cfg.MainConfigPath)
}
if cfg.RouteConfigPath != filepath.Join(dir, "data", defaultRouteConfigRelativePath) {
t.Fatalf("unexpected route config path: %s", cfg.RouteConfigPath)
}
if cfg.AccessLogPath != filepath.Join(dir, "data", defaultAccessLogRelativePath) {
t.Fatalf("unexpected access log path: %s", cfg.AccessLogPath)
}
if cfg.CertDir != filepath.Join(dir, "data", defaultCertDirRelativePath) {
t.Fatalf("unexpected cert dir: %s", cfg.CertDir)
}
if cfg.LuaDir != filepath.Join(dir, "data", defaultLuaDirRelativePath) {
t.Fatalf("unexpected lua dir: %s", cfg.LuaDir)
}
if cfg.RuntimeConfigDir != filepath.Join(dir, "data", defaultRuntimeConfigDirRelativePath) {
t.Fatalf("unexpected runtime config dir: %s", cfg.RuntimeConfigDir)
}
if cfg.OpenrestyCertDir != cfg.CertDir {
t.Fatalf("unexpected openresty cert dir: %s", cfg.OpenrestyCertDir)
}
if cfg.OpenrestyLuaDir != cfg.LuaDir {
t.Fatalf("unexpected openresty lua dir: %s", cfg.OpenrestyLuaDir)
}
if cfg.StatePath != filepath.Join(dir, "data", defaultStateRelativePath) {
t.Fatalf("unexpected state path: %s", cfg.StatePath)
}
if cfg.ObservabilityBufferPath != filepath.Join(dir, "data", defaultObservabilityBufferRelativePath) {
t.Fatalf("unexpected observability buffer path: %s", cfg.ObservabilityBufferPath)
}
if cfg.OpenrestyObservabilityPort != defaultOpenRestyObservabilityPort {
t.Fatalf("unexpected openresty observability port: %d", cfg.OpenrestyObservabilityPort)
}
if cfg.ObservabilityReplayMinutes != defaultObservabilityReplayMinutes {
t.Fatalf("unexpected observability replay minutes: %d", cfg.ObservabilityReplayMinutes)
}
}
func TestLoadPathModeKeepsExplicitPaths(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
"node_ip": "10.0.0.8",
"openresty_path": "/usr/local/openresty/nginx/sbin/openresty",
"main_config_path": "/tmp/nginx.conf",
"route_config_path": "/tmp/routes.conf",
"state_path": "/tmp/agent-state.json",
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.MainConfigPath != "/tmp/nginx.conf" {
t.Fatalf("unexpected main config path: %s", cfg.MainConfigPath)
}
if cfg.RouteConfigPath != "/tmp/routes.conf" {
t.Fatalf("unexpected route config path: %s", cfg.RouteConfigPath)
}
if cfg.StatePath != "/tmp/agent-state.json" {
t.Fatalf("unexpected state path: %s", cfg.StatePath)
}
if cfg.ObservabilityBufferPath != filepath.Join(dir, "data", defaultObservabilityBufferRelativePath) {
t.Fatalf("unexpected observability buffer path: %s", cfg.ObservabilityBufferPath)
}
if cfg.OpenrestyCertDir != cfg.CertDir {
t.Fatalf("expected path mode openresty cert dir to equal cert dir, got %s / %s", cfg.OpenrestyCertDir, cfg.CertDir)
}
if cfg.OpenrestyLuaDir != cfg.LuaDir {
t.Fatalf("expected path mode openresty lua dir to equal lua dir, got %s / %s", cfg.OpenrestyLuaDir, cfg.LuaDir)
}
if cfg.OpenrestyObservabilityPort != defaultOpenRestyObservabilityPort {
t.Fatalf("unexpected path mode openresty observability port: %d", cfg.OpenrestyObservabilityPort)
}
}
func TestLoadNormalizesExplicitResolvers(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
"node_ip": "10.0.0.8",
"openresty_resolvers": []string{" 10.0.0.2 ", "10.0.0.2", "", "1.1.1.1"},
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
expected := []string{"10.0.0.2", "1.1.1.1"}
if len(cfg.OpenrestyResolvers) != len(expected) {
t.Fatalf("unexpected resolver count: %#v", cfg.OpenrestyResolvers)
}
for index, value := range expected {
if cfg.OpenrestyResolvers[index] != value {
t.Fatalf("unexpected resolver at %d: got %q want %q", index, cfg.OpenrestyResolvers[index], value)
}
}
}
func TestLoadUsesCustomDataDirForGeneratedFiles(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
"node_ip": "10.0.0.8",
"data_dir": "/srv/openflare",
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.RouteConfigPath != "/srv/openflare/"+defaultRouteConfigRelativePath {
t.Fatalf("unexpected route config path: %s", cfg.RouteConfigPath)
}
if cfg.MainConfigPath != "/srv/openflare/"+defaultMainConfigRelativePath {
t.Fatalf("unexpected main config path: %s", cfg.MainConfigPath)
}
if cfg.AccessLogPath != "/srv/openflare/"+defaultAccessLogRelativePath {
t.Fatalf("unexpected access log path: %s", cfg.AccessLogPath)
}
if cfg.StatePath != "/srv/openflare/"+defaultStateRelativePath {
t.Fatalf("unexpected state path: %s", cfg.StatePath)
}
if cfg.ObservabilityBufferPath != "/srv/openflare/"+defaultObservabilityBufferRelativePath {
t.Fatalf("unexpected observability buffer path: %s", cfg.ObservabilityBufferPath)
}
if cfg.CertDir != "/srv/openflare/"+defaultCertDirRelativePath {
t.Fatalf("unexpected cert dir: %s", cfg.CertDir)
}
if cfg.LuaDir != "/srv/openflare/"+defaultLuaDirRelativePath {
t.Fatalf("unexpected lua dir: %s", cfg.LuaDir)
}
if cfg.RuntimeConfigDir != "/srv/openflare/"+defaultRuntimeConfigDirRelativePath {
t.Fatalf("unexpected runtime config dir: %s", cfg.RuntimeConfigDir)
}
}
func TestLoadUsesEnvConfigWhenFileIsMissing(t *testing.T) {
dir := t.TempDir()
t.Setenv("OPENFLARE_SERVER_URL", "http://127.0.0.1:3000")
t.Setenv("OPENFLARE_AGENT_TOKEN", "token")
t.Setenv("OPENFLARE_NODE_NAME", "edge-env")
t.Setenv("OPENFLARE_NODE_IP", "10.0.0.9")
t.Setenv("OPENFLARE_DATA_DIR", "/srv/openflare-env")
t.Setenv("OPENFLARE_OPENRESTY_PATH", "/usr/bin/openresty")
t.Setenv("OPENFLARE_HEARTBEAT_INTERVAL", "45s")
t.Setenv("OPENFLARE_REQUEST_TIMEOUT", "2500")
t.Setenv("OPENFLARE_OPENRESTY_OBSERVABILITY_PORT", "19091")
cfg, err := Load(filepath.Join(dir, "missing-agent.json"))
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.ServerURL != "http://127.0.0.1:3000" || cfg.AccessToken != "token" {
t.Fatalf("unexpected env auth config: %#v", cfg)
}
if cfg.OpenrestyPath != "/usr/bin/openresty" {
t.Fatalf("unexpected openresty path: %s", cfg.OpenrestyPath)
}
if cfg.DataDir != "/srv/openflare-env" {
t.Fatalf("unexpected data dir: %s", cfg.DataDir)
}
if cfg.HeartbeatInterval.Duration() != 45*time.Second {
t.Fatalf("unexpected heartbeat interval: %s", cfg.HeartbeatInterval)
}
if cfg.RequestTimeout.Duration() != 2500*time.Millisecond {
t.Fatalf("unexpected request timeout: %s", cfg.RequestTimeout)
}
if cfg.OpenrestyObservabilityPort != 19091 {
t.Fatalf("unexpected observability port: %d", cfg.OpenrestyObservabilityPort)
}
}
func TestLoadDetectsOutboundIPWhenNodeIPMissing(t *testing.T) {
previousLookup := lookupOutboundIP
lookupOutboundIP = func(ctx context.Context, strategies ...geoip.OutboundIPStrategy) (net.IP, error) {
return net.ParseIP("8.8.8.8"), nil
}
defer func() {
lookupOutboundIP = previousLookup
}()
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.NodeIP != "8.8.8.8" {
t.Fatalf("expected outbound IP, got %s", cfg.NodeIP)
}
}
func TestLoadFallsBackToLocalIPWhenOutboundLookupFails(t *testing.T) {
previousOutboundLookup := lookupOutboundIP
previousLocalLookup := lookupLocalIP
lookupOutboundIP = func(ctx context.Context, strategies ...geoip.OutboundIPStrategy) (net.IP, error) {
return nil, errors.New("realip.cc unavailable")
}
lookupLocalIP = func() string {
return "9.9.9.9"
}
defer func() {
lookupOutboundIP = previousOutboundLookup
lookupLocalIP = previousLocalLookup
}()
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.NodeIP != "9.9.9.9" {
t.Fatalf("expected local fallback IP, got %s", cfg.NodeIP)
}
}
func TestLoadEnvOverridesConfigFile(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
if err := os.WriteFile(configPath, []byte(`{"server_url":"http://old:3000","agent_token":"old","node_name":"edge-01","node_ip":"10.0.0.8","openresty_path":"/old/openresty"}`), 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
t.Setenv("OPENFLARE_SERVER_URL", "http://new:3000")
t.Setenv("OPENFLARE_AGENT_TOKEN", "new-token")
t.Setenv("OPENFLARE_OPENRESTY_PATH", "/new/openresty")
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.ServerURL != "http://new:3000" {
t.Fatalf("expected server url from env, got %s", cfg.ServerURL)
}
if cfg.AccessToken != "new-token" {
t.Fatalf("expected token from env, got %s", cfg.AccessToken)
}
if cfg.OpenrestyPath != "/new/openresty" {
t.Fatalf("expected openresty path from env, got %s", cfg.OpenrestyPath)
}
}
func TestLoadUsesMillisecondsForIntervals(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
payload := map[string]any{
"server_url": "http://127.0.0.1:3000",
"agent_token": "token",
"node_name": "edge-01",
"node_ip": "10.0.0.8",
"heartbeat_interval": 30000,
"request_timeout": 1500,
}
data, err := json.Marshal(payload)
if err != nil {
t.Fatalf("failed to marshal config: %v", err)
}
if err = os.WriteFile(configPath, data, 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if cfg.HeartbeatInterval.Duration() != 30*time.Second {
t.Fatalf("unexpected heartbeat interval: %s", cfg.HeartbeatInterval)
}
if cfg.RequestTimeout.Duration() != 1500*time.Millisecond {
t.Fatalf("unexpected request timeout: %s", cfg.RequestTimeout)
}
}
func TestSavePersistsMillisecondsAndOmitsRuntimeVersions(t *testing.T) {
dir := t.TempDir()
configPath := filepath.Join(dir, "agent.json")
if err := os.WriteFile(configPath, []byte(`{"server_url":"http://127.0.0.1:3000","agent_token":"token","node_name":"edge-01","node_ip":"10.0.0.8"}`), 0o644); err != nil {
t.Fatalf("failed to write config: %v", err)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
cfg.ExtVersion = "1.27.1.2"
cfg.HeartbeatInterval = MillisecondDuration(5 * time.Second)
cfg.RequestTimeout = MillisecondDuration(7 * time.Second)
cfg.OpenrestyResolvers = []string{"10.0.0.2", "1.1.1.1"}
if err = cfg.Save(); err != nil {
t.Fatalf("Save failed: %v", err)
}
data, err := os.ReadFile(configPath)
if err != nil {
t.Fatalf("failed to read saved config: %v", err)
}
var decoded map[string]any
if err = json.Unmarshal(data, &decoded); err != nil {
t.Fatalf("failed to decode saved config: %v", err)
}
if _, ok := decoded["agent_version"]; ok {
t.Fatal("agent_version should not be persisted")
}
if _, ok := decoded["nginx_version"]; ok {
t.Fatal("nginx_version should not be persisted")
}
if decoded["heartbeat_interval"] != float64(5000) {
t.Fatalf("unexpected heartbeat interval: %#v", decoded["heartbeat_interval"])
}
if decoded["request_timeout"] != float64(7000) {
t.Fatalf("unexpected request timeout: %#v", decoded["request_timeout"])
}
resolvers, ok := decoded["openresty_resolvers"].([]any)
if !ok || len(resolvers) != 2 || resolvers[0] != "10.0.0.2" || resolvers[1] != "1.1.1.1" {
t.Fatalf("unexpected resolvers: %#v", decoded["openresty_resolvers"])
}
if decoded["openresty_observability_port"] != float64(defaultOpenRestyObservabilityPort) {
t.Fatalf("unexpected observability port: %#v", decoded["openresty_observability_port"])
}
if decoded["observability_replay_minutes"] != float64(defaultObservabilityReplayMinutes) {
t.Fatalf("unexpected observability replay minutes: %#v", decoded["observability_replay_minutes"])
}
if _, ok := decoded["nginx_path"]; ok {
t.Fatal("legacy nginx_path should not be persisted")
}
}
func TestInitialAuthToken(t *testing.T) {
tests := []struct {
name string
agentToken string
discoveryToken string
expected string
}{
{
name: "prefer agent token",
agentToken: "agent-token",
discoveryToken: "discovery-token",
expected: "agent-token",
},
{
name: "fallback to discovery token",
agentToken: " ",
discoveryToken: "discovery-token",
expected: "discovery-token",
},
{
name: "nil config returns empty string",
agentToken: "",
discoveryToken: "",
expected: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var cfg *Config
if tt.name != "nil config returns empty string" {
cfg = &Config{
AccessToken: tt.agentToken,
DiscoveryToken: tt.discoveryToken,
}
}
if token := cfg.InitialAuthToken(); token != tt.expected {
t.Fatalf("unexpected initial auth token: %q", token)
}
})
}
}
func TestNodeIPPriority(t *testing.T) {
tests := []struct {
name string
ip string
expected int
}{
{
name: "public ipv4 preferred",
ip: "8.8.8.8",
expected: 2,
},
{
name: "private ipv4 fallback",
ip: "10.0.0.8",
expected: 1,
},
{
name: "link local ignored",
ip: "169.254.1.10",
expected: -1,
},
{
name: "loopback ignored",
ip: "127.0.0.1",
expected: -1,
},
{
name: "nil ignored",
ip: "",
expected: -1,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var parsed net.IP
if tt.ip != "" {
parsed = net.ParseIP(tt.ip)
}
if got := nodeIPPriority(parsed); got != tt.expected {
t.Fatalf("unexpected priority for %q: got %d want %d", tt.ip, got, tt.expected)
}
})
}
}
+54
View File
@@ -0,0 +1,54 @@
package config
import (
"encoding/json"
"fmt"
"strconv"
"strings"
"time"
)
type MillisecondDuration time.Duration
func (d MillisecondDuration) Duration() time.Duration {
return time.Duration(d)
}
func (d MillisecondDuration) String() string {
return time.Duration(d).String()
}
func (d *MillisecondDuration) UnmarshalJSON(data []byte) error {
raw := strings.TrimSpace(string(data))
if raw == "" || raw == "null" {
*d = 0
return nil
}
if strings.HasPrefix(raw, "\"") {
var text string
if err := json.Unmarshal(data, &text); err != nil {
return err
}
text = strings.TrimSpace(text)
if text == "" {
*d = 0
return nil
}
parsed, err := time.ParseDuration(text)
if err != nil {
return fmt.Errorf("invalid duration string %q: %w", text, err)
}
*d = MillisecondDuration(parsed)
return nil
}
ms, err := strconv.ParseInt(raw, 10, 64)
if err != nil {
return fmt.Errorf("invalid duration milliseconds %q: %w", raw, err)
}
*d = MillisecondDuration(time.Duration(ms) * time.Millisecond)
return nil
}
func (d MillisecondDuration) MarshalJSON() ([]byte, error) {
return json.Marshal(time.Duration(d).Milliseconds())
}
+3
View File
@@ -0,0 +1,3 @@
package config
var Version = "dev"
+8
View File
@@ -0,0 +1,8 @@
package geoipdata
import "embed"
//go:embed GeoLite2-Country.mmdb
var FS embed.FS
const DefaultMMDBName = "GeoLite2-Country.mmdb"
@@ -0,0 +1,67 @@
package geoipupdate
import (
"context"
"fmt"
"io/fs"
"log/slog"
"os"
"path/filepath"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/geoipdata"
"github.com/Rain-kl/Wavelet/pkg/geoip"
)
type Updater struct {
MMDBPath string
DownloadURL string
UpdateInterval time.Duration
}
func (u *Updater) EnsureInitialDatabase() error {
path := filepath.Clean(u.MMDBPath)
if path == "" || path == "." {
return nil
}
if _, err := os.Stat(path); err == nil {
return nil
} else if !os.IsNotExist(err) {
return fmt.Errorf("stat mmdb file failed: %w", err)
}
data, err := fs.ReadFile(geoipdata.FS, geoipdata.DefaultMMDBName)
if err != nil {
return fmt.Errorf("read embedded mmdb failed: %w", err)
}
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return fmt.Errorf("create mmdb directory failed: %w", err)
}
if err := os.WriteFile(path, data, 0o644); err != nil {
return fmt.Errorf("write initial mmdb failed: %w", err)
}
slog.Info("initialized GeoIP mmdb from embedded database", "path", path, "size", len(data))
return nil
}
func (u *Updater) Run(ctx context.Context) {
if u == nil || u.MMDBPath == "" || u.UpdateInterval <= 0 {
return
}
if err := u.EnsureInitialDatabase(); err != nil {
slog.Warn("initialize GeoIP mmdb failed", "path", u.MMDBPath, "error", err)
}
ticker := time.NewTicker(u.UpdateInterval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
if err := geoip.DownloadMaxMindDatabase(u.MMDBPath, u.DownloadURL); err != nil {
slog.Warn("update GeoIP mmdb failed", "path", u.MMDBPath, "error", err)
continue
}
slog.Info("GeoIP mmdb updated", "path", u.MMDBPath)
}
}
}
@@ -0,0 +1,24 @@
package geoipupdate
import (
"os"
"path/filepath"
"testing"
)
func TestEnsureInitialDatabaseCopiesEmbeddedMMDB(t *testing.T) {
tempDir := t.TempDir()
path := filepath.Join(tempDir, "GeoLite2-Country.mmdb")
updater := &Updater{MMDBPath: path}
if err := updater.EnsureInitialDatabase(); err != nil {
t.Fatalf("EnsureInitialDatabase failed: %v", err)
}
info, err := os.Stat(path)
if err != nil {
t.Fatalf("expected mmdb to exist: %v", err)
}
if info.Size() == 0 {
t.Fatal("expected copied mmdb to be non-empty")
}
}
+33
View File
@@ -0,0 +1,33 @@
package heartbeat
import (
"context"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
type Client interface {
RegisterNode(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error)
Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
SetToken(token string)
}
type Service struct {
client Client
}
func New(client Client) *Service {
return &Service{client: client}
}
func (s *Service) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) {
return s.client.RegisterNode(ctx, payload)
}
func (s *Service) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) {
return s.client.Heartbeat(ctx, payload)
}
func (s *Service) SetToken(token string) {
s.client.SetToken(token)
}
+194
View File
@@ -0,0 +1,194 @@
package httpclient
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
type Client struct {
baseURL string
token string
httpClient *http.Client
}
func New(baseURL string, token string, timeout time.Duration) *Client {
return &Client{
baseURL: strings.TrimRight(baseURL, "/"),
token: token,
httpClient: &http.Client{
Timeout: timeout,
},
}
}
func (c *Client) RegisterNode(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) {
slog.Debug("http register node request", "node_id", payload.NodeID, "current_version", payload.CurrentVersion)
resp := protocol.APIResponse[protocol.RegisterNodeResponse]{}
if err := c.postJSON(ctx, "/api/v1/agent/nodes/register", payload, &resp); err != nil {
return nil, err
}
if err := apiError(resp.ErrorMsg); err != nil {
return nil, err
}
slog.Debug("http register node response", "node_id", resp.Data.NodeID)
return &resp.Data, nil
}
func (c *Client) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) {
resp := protocol.APIResponse[protocol.HeartbeatData]{}
if err := c.postJSON(ctx, "/api/v1/agent/nodes/heartbeat", payload, &resp); err != nil {
return nil, err
}
if err := apiError(resp.ErrorMsg); err != nil {
return nil, err
}
return &protocol.HeartbeatResult{
AgentSettings: resp.Data.AgentSettings,
ActiveConfig: resp.Data.ActiveConfig,
WAFIPGroups: resp.Data.WAFIPGroups,
}, nil
}
func (c *Client) GetActiveConfig(ctx context.Context) (*protocol.ActiveConfigResponse, error) {
resp := protocol.APIResponse[protocol.ActiveConfigResponse]{}
if err := c.getJSON(ctx, "/api/v1/agent/config-versions/active", &resp); err != nil {
return nil, err
}
if err := apiError(resp.ErrorMsg); err != nil {
return nil, err
}
slog.Debug("http get active config response", "version", resp.Data.Version, "checksum", resp.Data.Checksum, "support_files", len(resp.Data.SupportFiles))
return &resp.Data, nil
}
func (c *Client) ReportApplyLog(ctx context.Context, payload protocol.ApplyLogPayload) error {
slog.Debug("http report apply log request", "node_id", payload.NodeID, "version", payload.Version, "result", payload.Result)
resp := protocol.APIResponse[json.RawMessage]{}
if err := c.postJSON(ctx, "/api/v1/agent/apply-logs", payload, &resp); err != nil {
return err
}
return apiError(resp.ErrorMsg)
}
func (c *Client) SyncWAFIPGroups(ctx context.Context, payload protocol.WAFIPGroupSyncRequest) (*protocol.WAFIPGroupSyncResponse, error) {
resp := protocol.APIResponse[protocol.WAFIPGroupSyncResponse]{}
if err := c.postJSON(ctx, "/api/v1/agent/waf/ip-groups/sync", payload, &resp); err != nil {
return nil, err
}
if err := apiError(resp.ErrorMsg); err != nil {
return nil, err
}
return &resp.Data, nil
}
func (c *Client) DownloadPagesDeploymentPackage(ctx context.Context, deploymentID uint) ([]byte, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+fmt.Sprintf("/api/v1/agent/pages/deployments/%d/package", deploymentID), nil)
if err != nil {
return nil, err
}
req.Header.Set("X-Agent-Token", c.token)
res, err := c.httpClient.Do(req)
if err != nil {
return nil, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return nil, readHTTPError(res)
}
return io.ReadAll(res.Body)
}
func (c *Client) SetToken(token string) {
c.token = strings.TrimSpace(token)
slog.Debug("http client token updated")
}
func (c *Client) getJSON(ctx context.Context, path string, target any) error {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+path, nil)
if err != nil {
return err
}
req.Header.Set("X-Agent-Token", c.token)
return c.do(req, target)
}
func (c *Client) postJSON(ctx context.Context, path string, body any, target any) error {
data, err := json.Marshal(body)
if err != nil {
return err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, bytes.NewReader(data))
if err != nil {
return err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("X-Agent-Token", c.token)
return c.do(req, target)
}
func (c *Client) do(req *http.Request, target any) error {
res, err := c.httpClient.Do(req)
if err != nil {
slog.Error("http request failed", "method", req.Method, "path", req.URL.Path, "error", err)
return err
}
defer func(Body io.ReadCloser) {
err := Body.Close()
if err != nil {
slog.Error("failed to close response body", "error", err)
}
}(res.Body)
body, err := io.ReadAll(res.Body)
if err != nil {
slog.Error("http response read failed", "method", req.Method, "path", req.URL.Path, "error", err)
return err
}
if res.StatusCode != http.StatusOK {
slog.Warn("http request returned non-200", "method", req.Method, "path", req.URL.Path, "status", res.Status)
return readBodyError(body, res.Status)
}
if target == nil {
return nil
}
if err = json.Unmarshal(body, target); err != nil {
slog.Error("http response decode failed", "method", req.Method, "path", req.URL.Path, "error", err)
return err
}
return nil
}
func apiError(msg string) error {
if strings.TrimSpace(msg) == "" {
return nil
}
return errors.New(msg)
}
func readHTTPError(res *http.Response) error {
body, err := io.ReadAll(res.Body)
if err != nil {
return errors.New(res.Status)
}
return readBodyError(body, res.Status)
}
func readBodyError(body []byte, fallback string) error {
var errBody struct {
ErrorMsg string `json:"error_msg"`
}
if err := json.Unmarshal(body, &errBody); err == nil && strings.TrimSpace(errBody.ErrorMsg) != "" {
return errors.New(errBody.ErrorMsg)
}
return errors.New(fallback)
}
+29
View File
@@ -0,0 +1,29 @@
package logging
import (
"log/slog"
"os"
"strings"
)
func Setup() {
opts := &slog.HandlerOptions{
AddSource: true,
Level: parseLevel(os.Getenv("LOG_LEVEL")),
}
handler := slog.NewTextHandler(os.Stdout, opts)
slog.SetDefault(slog.New(handler))
}
func parseLevel(value string) slog.Level {
switch strings.ToLower(strings.TrimSpace(value)) {
case "debug":
return slog.LevelDebug
case "warn", "warning":
return slog.LevelWarn
case "error":
return slog.LevelError
default:
return slog.LevelInfo
}
}
File diff suppressed because it is too large Load Diff
+968
View File
@@ -0,0 +1,968 @@
package nginx
import (
"context"
"errors"
"net"
"net/http"
"os"
"path/filepath"
"reflect"
"runtime"
"strings"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
type runCall struct {
name string
args []string
}
type fakeRunner struct {
calls []runCall
runFn func(name string, args ...string) ([]byte, error)
}
type fakeExecutor struct {
testErr error
reloadErr error
}
type scriptedExecutor struct {
testErrors []error
testCalls int
reloadErrors []error
reloadCalls int
}
func (r *fakeRunner) Run(ctx context.Context, name string, args ...string) ([]byte, error) {
r.calls = append(r.calls, runCall{name: name, args: append([]string{}, args...)})
if r.runFn != nil {
return r.runFn(name, args...)
}
return nil, nil
}
func (e *fakeExecutor) Test(ctx context.Context) error {
return e.testErr
}
func (e *fakeExecutor) Reload(ctx context.Context) error {
return e.reloadErr
}
func (e *fakeExecutor) EnsureRuntime(ctx context.Context, recreate bool) error {
return nil
}
func (e *fakeExecutor) CheckHealth(ctx context.Context) error {
return e.testErr
}
func (e *fakeExecutor) Restart(ctx context.Context) error {
return e.reloadErr
}
func (e *scriptedExecutor) Test(ctx context.Context) error {
index := e.testCalls
e.testCalls++
if index >= len(e.testErrors) {
return nil
}
return e.testErrors[index]
}
func (e *scriptedExecutor) Reload(ctx context.Context) error {
index := e.reloadCalls
e.reloadCalls++
if index >= len(e.reloadErrors) {
return nil
}
return e.reloadErrors[index]
}
func (e *scriptedExecutor) EnsureRuntime(ctx context.Context, recreate bool) error {
return nil
}
func (e *scriptedExecutor) CheckHealth(ctx context.Context) error {
return nil
}
func (e *scriptedExecutor) Restart(ctx context.Context) error {
return nil
}
func TestPathExecutorCommands(t *testing.T) {
runner := &fakeRunner{}
executor := &PathExecutor{
Path: "/usr/local/openresty/nginx/sbin/openresty",
ConfigPath: "/data/etc/nginx/nginx.conf",
Runner: runner,
}
if err := executor.Test(context.Background()); err != nil {
t.Fatalf("Test failed: %v", err)
}
if err := executor.Reload(context.Background()); err != nil {
t.Fatalf("Reload failed: %v", err)
}
expected := []runCall{
{name: "/usr/local/openresty/nginx/sbin/openresty", args: []string{"-t", "-c", "/data/etc/nginx/nginx.conf"}},
{name: "/usr/local/openresty/nginx/sbin/openresty", args: []string{"-s", "reload", "-c", "/data/etc/nginx/nginx.conf"}},
}
if !reflect.DeepEqual(runner.calls, expected) {
t.Fatalf("unexpected calls: %#v", runner.calls)
}
}
func TestPathExecutorEnsureRuntimeNoop(t *testing.T) {
runner := &fakeRunner{}
executor := &PathExecutor{
Path: "/usr/local/openresty/nginx/sbin/openresty",
ConfigPath: "/data/etc/nginx/nginx.conf",
Runner: runner,
}
if err := executor.EnsureRuntime(context.Background(), true); err != nil {
t.Fatalf("EnsureRuntime failed: %v", err)
}
if len(runner.calls) != 2 {
t.Fatalf("expected test and reload calls, got %d", len(runner.calls))
}
}
func TestPathExecutorRestartIgnoresMissingPID(t *testing.T) {
runner := &fakeRunner{
runFn: func(name string, args ...string) ([]byte, error) {
if len(args) == 2 && args[0] == "-s" && args[1] == "quit" {
return []byte("openresty: [error] invalid PID number \"\" in \"/usr/local/openresty/nginx/logs/nginx.pid\""), errors.New("exit status 1")
}
return []byte(""), nil
},
}
executor := &PathExecutor{
Path: "/usr/local/openresty/nginx/sbin/openresty",
ConfigPath: "/data/etc/nginx/nginx.conf",
Runner: runner,
}
if err := executor.Restart(context.Background()); err != nil {
t.Fatalf("Restart failed: %v", err)
}
if len(runner.calls) != 2 {
t.Fatalf("expected 2 restart calls, got %d", len(runner.calls))
}
}
func TestPathExecutorReloadStartsWhenRuntimeIsNotRunning(t *testing.T) {
runner := &fakeRunner{
runFn: func(name string, args ...string) ([]byte, error) {
if len(args) >= 2 && args[0] == "-s" && args[1] == "reload" {
return []byte("openresty: [error] invalid PID number \"\" in \"/usr/local/openresty/nginx/logs/nginx.pid\""), errors.New("exit status 1")
}
return []byte(""), nil
},
}
executor := &PathExecutor{
Path: "/usr/local/openresty/nginx/sbin/openresty",
ConfigPath: "/data/etc/nginx/nginx.conf",
Runner: runner,
}
if err := executor.Reload(context.Background()); err != nil {
t.Fatalf("Reload failed: %v", err)
}
expected := []runCall{
{name: "/usr/local/openresty/nginx/sbin/openresty", args: []string{"-s", "reload", "-c", "/data/etc/nginx/nginx.conf"}},
{name: "/usr/local/openresty/nginx/sbin/openresty", args: []string{"-c", "/data/etc/nginx/nginx.conf"}},
}
if !reflect.DeepEqual(runner.calls, expected) {
t.Fatalf("unexpected calls: %#v", runner.calls)
}
}
func TestDetectVersionFromBinary(t *testing.T) {
version, err := detectVersion(context.Background(), ExecutorOptions{
NginxPath: "/usr/local/openresty/nginx/sbin/openresty",
}, &fakeRunner{
runFn: func(name string, args ...string) ([]byte, error) {
return []byte("nginx version: openresty/1.27.1.2\n"), nil
},
})
if err != nil {
t.Fatalf("detectVersion failed: %v", err)
}
if version != "1.27.1.2" {
t.Fatalf("unexpected version: %s", version)
}
}
func TestManagerApplyAndChecksumIncludeMainConfig(t *testing.T) {
tempDir := t.TempDir()
mainPath := filepath.Join(tempDir, "nginx.conf")
routePath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
certDir := filepath.Join(tempDir, "certs")
accessLogPath := filepath.Join(tempDir, "var", "log", "openflare", "access.log")
manager := &Manager{
MainConfigPath: mainPath,
RouteConfigPath: routePath,
AccessLogPath: accessLogPath,
CertDir: certDir,
NginxCertDir: "/etc/nginx/openflare-certs",
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
Executor: &fakeExecutor{},
}
outcome := manager.Apply(
context.Background(),
"include __OPENFLARE_ROUTE_CONFIG__;\naccess_log __OPENFLARE_ACCESS_LOG__ openflare_json;\n",
"ssl_certificate __OPENFLARE_CERT_DIR__/1.crt;\n",
[]protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
)
if outcome.Status != ApplyStatusSuccess {
t.Fatalf("Apply failed: %#v", outcome)
}
mainData, err := os.ReadFile(mainPath)
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
expectedMain := "include " + routePath + ";\naccess_log " + filepath.ToSlash(accessLogPath) + " openflare_json;\n"
if string(mainData) != expectedMain {
t.Fatalf("unexpected main config: %s", string(mainData))
}
routeData, err := os.ReadFile(routePath)
if err != nil {
t.Fatalf("failed to read route config: %v", err)
}
if string(routeData) != "ssl_certificate /etc/nginx/openflare-certs/1.crt;\n" {
t.Fatalf("unexpected route config: %s", string(routeData))
}
value, err := manager.CurrentChecksum()
if err != nil {
t.Fatalf("CurrentChecksum failed: %v", err)
}
expected := bundleChecksum(
"include __OPENFLARE_ROUTE_CONFIG__;\naccess_log __OPENFLARE_ACCESS_LOG__ openflare_json;\n",
"ssl_certificate __OPENFLARE_CERT_DIR__/1.crt;\n",
[]protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
)
if value != expected {
t.Fatalf("unexpected checksum: got %s want %s", value, expected)
}
}
func TestParseExtVersionIgnoresDockerEntrypointPaths(t *testing.T) {
output := strings.Join([]string{
"/docker-entrypoint.sh: /docker-entrypoint.d/10-listen-on-ipv6-by-default.sh: info: can not modify /etc/nginx/conf.d/default.conf (read-only file system?)",
"nginx version: openresty/1.27.1.2",
}, "\n")
version := parseExtVersion(output)
if version != "1.27.1.2" {
t.Fatalf("unexpected version: %s", version)
}
}
func TestManagerApplyWritesSupportFilesAndReplacesPlaceholder(t *testing.T) {
tempDir := t.TempDir()
manager := &Manager{
MainConfigPath: filepath.Join(tempDir, "nginx.conf"),
RouteConfigPath: filepath.Join(tempDir, "routes.conf"),
CertDir: filepath.Join(tempDir, "certs"),
NginxCertDir: "/etc/nginx/openflare-certs",
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
OpenrestyObservabilityListen: "18081",
OpenrestyResolverDirective: " resolver 127.0.0.11 valid=30s ipv6=off;\n resolver_timeout 5s;\n",
Executor: &fakeExecutor{},
}
outcome := manager.Apply(context.Background(), "include __OPENFLARE_ROUTE_CONFIG__;\n__OPENFLARE_RESOLVER_DIRECTIVE__server { listen __OPENFLARE_OBSERVABILITY_LISTEN__; }", "ssl_certificate __OPENFLARE_CERT_DIR__/1.crt;", []protocol.SupportFile{
{Path: "1.crt", Content: "cert-data"},
{Path: "1.key", Content: "key-data"},
})
if outcome.Status != ApplyStatusSuccess {
t.Fatalf("Apply failed: %#v", outcome)
}
routeData, err := os.ReadFile(manager.RouteConfigPath)
if err != nil {
t.Fatalf("failed to read route config: %v", err)
}
if !strings.Contains(string(routeData), "/etc/nginx/openflare-certs/1.crt") {
t.Fatalf("expected placeholder replacement in route config, got %s", string(routeData))
}
renderedRoute := manager.renderRouteConfig("access_by_lua_file __OPENFLARE_LUA_DIR__/pow/check.lua;\nlocation /.within.website/x/cmd/anubis/static/ { alias __OPENFLARE_POW_STATIC_DIR__/; }\n")
if !strings.Contains(renderedRoute, "access_by_lua_file /etc/nginx/openflare-lua/pow/check.lua;") {
t.Fatalf("expected lua dir placeholder replacement in route config, got %s", renderedRoute)
}
if !strings.Contains(renderedRoute, "alias /etc/nginx/openflare-lua/pow/static/;") {
t.Fatalf("expected pow static dir placeholder replacement in route config, got %s", renderedRoute)
}
mainData, err := os.ReadFile(manager.MainConfigPath)
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
if !strings.Contains(string(mainData), "listen 18081;") {
t.Fatalf("expected observability listen placeholder replacement in main config, got %s", string(mainData))
}
if !strings.Contains(string(mainData), "resolver 127.0.0.11 valid=30s ipv6=off;") {
t.Fatalf("expected resolver directive placeholder replacement in main config, got %s", string(mainData))
}
certData, err := os.ReadFile(filepath.Join(manager.CertDir, "1.crt"))
if err != nil {
t.Fatalf("failed to read cert file: %v", err)
}
if string(certData) != "cert-data" {
t.Fatalf("unexpected cert file content: %s", string(certData))
}
luaInfo, err := os.Stat(filepath.Join(manager.LuaDir, "log.lua"))
if err != nil {
t.Fatalf("expected managed lua file to exist, stat err = %v", err)
}
if runtime.GOOS != "windows" && luaInfo.Mode().Perm() != 0o644 {
t.Fatalf("unexpected lua mode: %o", luaInfo.Mode().Perm())
}
}
func TestManagerCheckHealthUsesStubStatusInsteadOfConfigTest(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("Listen failed: %v", err)
}
port := listener.Addr().(*net.TCPAddr).Port
server := &http.Server{
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/openflare/stub_status" {
http.NotFound(w, r)
return
}
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("Active connections: 1\n"))
}),
}
go func() {
_ = server.Serve(listener)
}()
defer server.Shutdown(context.Background())
mainPath := filepath.Join(t.TempDir(), "nginx.conf")
if err := os.WriteFile(mainPath, []byte("main"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{
MainConfigPath: mainPath,
OpenrestyObservabilityPort: port,
Executor: &fakeExecutor{
testErr: errors.New("openresty -t should not be called"),
},
}
if err := manager.CheckHealth(context.Background()); err != nil {
t.Fatalf("CheckHealth failed: %v", err)
}
}
func TestManagerCheckHealthFailsWhenStubStatusUnavailable(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("Listen failed: %v", err)
}
port := listener.Addr().(*net.TCPAddr).Port
if err := listener.Close(); err != nil {
t.Fatalf("listener close failed: %v", err)
}
mainPath := filepath.Join(t.TempDir(), "nginx.conf")
if err := os.WriteFile(mainPath, []byte("main"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{
MainConfigPath: mainPath,
OpenrestyObservabilityPort: port,
Executor: &fakeExecutor{},
}
if err := manager.CheckHealth(context.Background()); err == nil {
t.Fatal("expected CheckHealth to fail when stub_status is unavailable")
}
}
func TestResolverDirectiveUsesExplicitResolvers(t *testing.T) {
got := ResolverDirective([]string{"10.0.0.2", "1.1.1.1"})
if !strings.Contains(got, "resolver 10.0.0.2 1.1.1.1") {
t.Fatalf("expected explicit resolver directive, got %q", got)
}
}
func TestParseResolverAddressesFiltersLoopbackForDocker(t *testing.T) {
content := strings.Join([]string{
"nameserver 127.0.0.53",
"nameserver 10.0.0.2",
"nameserver ::1",
"nameserver 1.1.1.1",
}, "\n")
got := parseResolverAddresses(content, true)
expected := []string{"10.0.0.2", "1.1.1.1"}
if !reflect.DeepEqual(got, expected) {
t.Fatalf("unexpected docker resolvers: got %#v want %#v", got, expected)
}
}
func TestParseResolverAddressesKeepsLoopbackForLocalBinary(t *testing.T) {
content := strings.Join([]string{
"nameserver 127.0.0.53",
"nameserver 10.0.0.2",
}, "\n")
got := parseResolverAddresses(content, false)
expected := []string{"127.0.0.53", "10.0.0.2"}
if !reflect.DeepEqual(got, expected) {
t.Fatalf("unexpected local resolvers: got %#v want %#v", got, expected)
}
}
func TestRequiresRuntimeResolver(t *testing.T) {
testCases := []struct {
name string
originURL string
want bool
}{
{name: "hostname", originURL: "https://origin.internal", want: true},
{name: "ipv4", originURL: "https://10.0.0.8", want: false},
{name: "ipv6", originURL: "https://[2001:db8::1]", want: false},
{name: "invalid", originURL: "://bad", want: false},
}
for _, testCase := range testCases {
if got := RequiresRuntimeResolver(testCase.originURL); got != testCase.want {
t.Fatalf("%s: got %v want %v", testCase.name, got, testCase.want)
}
}
}
func TestWriteCertFilesKeepsBaseDirAndRemovesStaleFiles(t *testing.T) {
tempDir := t.TempDir()
certDir := filepath.Join(tempDir, "certs")
if err := os.MkdirAll(filepath.Join(certDir, "stale"), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
if err := os.WriteFile(filepath.Join(certDir, "stale", "old.crt"), []byte("old"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{CertDir: certDir}
if err := manager.writeCertFiles([]protocol.SupportFile{
{Path: "1.crt", Content: "cert"},
{Path: "1.key", Content: "key"},
}); err != nil {
t.Fatalf("writeCertFiles failed: %v", err)
}
if _, err := os.Stat(certDir); err != nil {
t.Fatalf("expected cert dir to persist, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(certDir, "stale", "old.crt")); !os.IsNotExist(err) {
t.Fatalf("expected stale cert file to be removed, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(certDir, "1.crt")); err != nil {
t.Fatalf("expected new cert file to exist, stat err = %v", err)
}
}
func TestEnsureLuaAssetsKeepsBaseDirAndRemovesStaleFiles(t *testing.T) {
tempDir := t.TempDir()
luaDir := filepath.Join(tempDir, "lua")
if err := os.MkdirAll(filepath.Join(luaDir, "stale"), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
if err := os.WriteFile(filepath.Join(luaDir, "stale", "old.lua"), []byte("old"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{LuaDir: luaDir}
if err := manager.EnsureLuaAssets(); err != nil {
t.Fatalf("EnsureLuaAssets failed: %v", err)
}
if _, err := os.Stat(luaDir); err != nil {
t.Fatalf("expected lua dir to persist, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(luaDir, "stale", "old.lua")); !os.IsNotExist(err) {
t.Fatalf("expected stale lua file to be removed, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(luaDir, "log.lua")); err != nil {
t.Fatalf("expected managed lua file to exist, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(luaDir, "pow", "check.lua")); err != nil {
t.Fatalf("expected managed pow lua file to exist, stat err = %v", err)
}
if _, err := os.Stat(filepath.Join(luaDir, "pow", "static", "js", "main.mjs")); err != nil {
t.Fatalf("expected managed pow static asset to exist, stat err = %v", err)
}
}
func TestCertFileMode(t *testing.T) {
testCases := []struct {
path string
want os.FileMode
}{
{path: "1.crt", want: 0o644},
{path: "1.pem", want: 0o644},
{path: "1.key", want: 0o600},
{path: "misc.txt", want: 0o644},
}
for _, testCase := range testCases {
if got := certFileMode(testCase.path); got != testCase.want {
t.Fatalf("unexpected mode for %s: got %o want %o", testCase.path, got, testCase.want)
}
}
}
func TestManagerEnsureLuaAssetsWritesReadableFiles(t *testing.T) {
tempDir := t.TempDir()
manager := &Manager{
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
RuntimeConfigDir: filepath.Join(tempDir, "runtime"),
}
err := manager.EnsureLuaAssets()
if err != nil {
t.Fatalf("EnsureLuaAssets failed: %v", err)
}
luaInfo, err := os.Stat(filepath.Join(manager.LuaDir, "log.lua"))
if err != nil {
t.Fatalf("failed to stat lua file: %v", err)
}
if luaInfo.Mode().Perm() != 0o644 {
t.Fatalf("unexpected lua mode: %o", luaInfo.Mode().Perm())
}
if _, err := os.Stat(filepath.Join(manager.LuaDir, "pow", "check.lua")); err != nil {
t.Fatalf("failed to stat pow lua file: %v", err)
}
data, err := os.ReadFile(filepath.Join(manager.LuaDir, "pow", "runtime.lua"))
if err != nil {
t.Fatalf("failed to read pow lua file: %v", err)
}
if !strings.Contains(string(data), filepath.ToSlash(manager.RuntimeConfigDir)+"/waf_config.json") {
t.Fatalf("expected pow lua to read runtime config dir, got %s", string(data))
}
}
func TestEnsureLuaAssetsLeavesRuntimePowConfigOutsideLuaDir(t *testing.T) {
tempDir := t.TempDir()
luaDir := filepath.Join(tempDir, "lua")
runtimeConfigDir := filepath.Join(tempDir, "runtime")
if err := os.MkdirAll(runtimeConfigDir, 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
powConfigPath := filepath.Join(runtimeConfigDir, "pow_config.json")
want := `[{"domains":["pow.example.com"],"enabled":true}]`
if err := os.WriteFile(powConfigPath, []byte(want), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{LuaDir: luaDir, RuntimeConfigDir: runtimeConfigDir}
if err := manager.EnsureLuaAssets(); err != nil {
t.Fatalf("EnsureLuaAssets failed: %v", err)
}
got, err := os.ReadFile(powConfigPath)
if err != nil {
t.Fatalf("expected pow_config.json to remain after EnsureLuaAssets: %v", err)
}
if string(got) != want {
t.Fatalf("unexpected pow_config.json content: got %s want %s", string(got), want)
}
if _, err := os.Stat(filepath.Join(luaDir, "pow_config.json")); !os.IsNotExist(err) {
t.Fatalf("expected lua pow_config.json to stay absent, stat err = %v", err)
}
}
func TestManagerApplyWritesPowConfigToRuntimeDirAndCleansLegacyCopies(t *testing.T) {
tempDir := t.TempDir()
certDir := filepath.Join(tempDir, "certs")
luaDir := filepath.Join(tempDir, "lua")
runtimeConfigDir := filepath.Join(tempDir, "runtime")
for _, dir := range []string{certDir, luaDir, runtimeConfigDir} {
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
}
for _, path := range []string{filepath.Join(certDir, "pow_config.json"), filepath.Join(luaDir, "pow_config.json")} {
if err := os.WriteFile(path, []byte("stale"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
}
manager := &Manager{
MainConfigPath: filepath.Join(tempDir, "nginx.conf"),
RouteConfigPath: filepath.Join(tempDir, "routes.conf"),
CertDir: certDir,
LuaDir: luaDir,
RuntimeConfigDir: runtimeConfigDir,
Executor: &fakeExecutor{},
}
outcome := manager.Apply(context.Background(), "main", "route", []protocol.SupportFile{
{Path: "pow_config.json", Content: "runtime"},
})
if outcome.Status != ApplyStatusSuccess {
t.Fatalf("Apply failed: %#v", outcome)
}
data, err := os.ReadFile(filepath.Join(runtimeConfigDir, "pow_config.json"))
if err != nil {
t.Fatalf("failed to read runtime pow config: %v", err)
}
if string(data) != "runtime" {
t.Fatalf("unexpected runtime pow config: %s", string(data))
}
for _, path := range []string{filepath.Join(certDir, "pow_config.json"), filepath.Join(luaDir, "pow_config.json")} {
if _, err := os.Stat(path); !os.IsNotExist(err) {
t.Fatalf("expected legacy pow config to be removed from %s, stat err = %v", path, err)
}
}
}
func TestManagerCurrentChecksumIncludesPowConfig(t *testing.T) {
tempDir := t.TempDir()
mainPath := filepath.Join(tempDir, "nginx.conf")
routePath := filepath.Join(tempDir, "routes.conf")
luaDir := filepath.Join(tempDir, "lua")
runtimeConfigDir := filepath.Join(tempDir, "runtime")
manager := &Manager{
MainConfigPath: mainPath,
RouteConfigPath: routePath,
LuaDir: luaDir,
NginxLuaDir: "/etc/nginx/openflare-lua",
RuntimeConfigDir: runtimeConfigDir,
Executor: &fakeExecutor{},
}
outcome := manager.Apply(
context.Background(),
"access_log __OPENFLARE_ACCESS_LOG__ openflare_json;\n",
"location /.within.website/x/cmd/anubis/static/ { alias __OPENFLARE_POW_STATIC_DIR__/; }\n",
[]protocol.SupportFile{{Path: "pow_config.json", Content: `[{"domains":["pow.example.com"],"enabled":true}]`}},
)
if outcome.Status != ApplyStatusSuccess {
t.Fatalf("Apply failed: %#v", outcome)
}
value, err := manager.CurrentChecksum()
if err != nil {
t.Fatalf("CurrentChecksum failed: %v", err)
}
expected := bundleChecksum(
"access_log __OPENFLARE_ACCESS_LOG__ openflare_json;\n",
"location /.within.website/x/cmd/anubis/static/ { alias __OPENFLARE_POW_STATIC_DIR__/; }\n",
[]protocol.SupportFile{{Path: "pow_config.json", Content: `[{"domains":["pow.example.com"],"enabled":true}]`}},
)
if value != expected {
t.Fatalf("unexpected checksum with pow config: got %s want %s", value, expected)
}
}
func TestManagedPowLuaFilesUseInternalChallengeFlow(t *testing.T) {
if !strings.Contains(openRestyPowRuntimeLua, `return ngx.exec("/.within.website/x/cmd/anubis/api/make-challenge")`) {
t.Fatal("expected pow runtime lua to internally execute make-challenge instead of issuing a 302 redirect")
}
if strings.Contains(openRestyPowRuntimeLua, "ngx.redirect(") {
t.Fatal("expected pow runtime lua to avoid external redirects for challenge rendering")
}
if !strings.Contains(openRestyPowChallengeLua, `<h1 id="title" class="centered-div">`) {
t.Fatal("expected challenge html to include Anubis-compatible title node")
}
if !strings.Contains(openRestyPowChallengeLua, `<div id="progress" role="progressbar" aria-labelledby="status"><div class="bar-inner"></div></div>`) {
t.Fatal("expected challenge html to include Anubis-compatible progress markup")
}
if !strings.Contains(openRestyPowChallengeLua, `<script id="anubis_public_url" type="application/json">"__openflare_internal__"</script>`) {
t.Fatal("expected challenge html to force Anubis frontend to reuse the current URL as redir target")
}
if !strings.Contains(openRestyPowRuntimeLua, `pow_sessions:set(session_key, "1", session_ttl)`) {
t.Fatal("expected pow runtime lua to refresh the PoW session TTL on each valid request")
}
if !strings.Contains(openRestyPowRuntimeLua, `ngx.header["Set-Cookie"] = session_cookie(cookie_val, session_ttl)`) {
t.Fatal("expected pow runtime lua to refresh the browser session cookie on each valid request")
}
if !strings.Contains(openRestyPowChallengeLua, `local session_ttl = config.session_ttl or 600`) {
t.Fatal("expected challenge.lua to default session TTL to 10 minutes")
}
if !strings.Contains(openRestyPowVerifyLua, `local session_ttl = challenge_info.session_ttl or 600`) {
t.Fatal("expected verify.lua to default session TTL to 10 minutes")
}
if !strings.Contains(openRestyPowVerifyLua, `if ngx.var.scheme == "https" then`) {
t.Fatal("expected verify.lua to only mark the session cookie as Secure for HTTPS requests")
}
}
func TestManagedWAFLuaTreatsWhitelistAsAllowlist(t *testing.T) {
if !strings.Contains(openRestyWAFRuntimeLua, "local function first_allowlist_group(groups)") {
t.Fatal("expected waf runtime to detect allowlist rule groups")
}
if !strings.Contains(openRestyWAFRuntimeLua, "local allowlist_group = first_allowlist_group(groups)") {
t.Fatal("expected waf runtime to enter allowlist mode when whitelist rules exist")
}
if !strings.Contains(openRestyWAFRuntimeLua, "return exit_with_group(allowlist_group)") {
t.Fatal("expected waf runtime to block requests that miss configured whitelists")
}
}
func TestManagerRollbackRestoresCertFiles(t *testing.T) {
tempDir := t.TempDir()
routePath := filepath.Join(tempDir, "routes.conf")
mainPath := filepath.Join(tempDir, "nginx.conf")
certDir := filepath.Join(tempDir, "certs")
if err := os.MkdirAll(certDir, 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
if err := os.WriteFile(mainPath, []byte("old-main"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
if err := os.WriteFile(routePath, []byte("old-route"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
if err := os.WriteFile(filepath.Join(certDir, "1.crt"), []byte("old-cert"), 0o600); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{
MainConfigPath: mainPath,
RouteConfigPath: routePath,
CertDir: certDir,
NginxCertDir: "/etc/nginx/openflare-certs",
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
Executor: &fakeExecutor{
reloadErr: errors.New("openresty reload failed"),
},
}
outcome := manager.Apply(context.Background(), "new-main", "new-route", []protocol.SupportFile{
{Path: "1.crt", Content: "new-cert"},
})
if outcome.Status != ApplyStatusFatal {
t.Fatalf("expected fatal apply outcome, got %#v", outcome)
}
mainData, err := os.ReadFile(mainPath)
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
if string(mainData) != "old-main" {
t.Fatalf("expected main rollback, got %s", string(mainData))
}
routeData, err := os.ReadFile(routePath)
if err != nil {
t.Fatalf("failed to read route config: %v", err)
}
if string(routeData) != "old-route" {
t.Fatalf("expected route rollback, got %s", string(routeData))
}
certData, err := os.ReadFile(filepath.Join(certDir, "1.crt"))
if err != nil {
t.Fatalf("failed to read cert file: %v", err)
}
if string(certData) != "old-cert" {
t.Fatalf("expected cert rollback, got %s", string(certData))
}
}
func TestManagerApplyReturnsWarningWhenRollbackRecoversRuntime(t *testing.T) {
tempDir := t.TempDir()
routePath := filepath.Join(tempDir, "routes.conf")
mainPath := filepath.Join(tempDir, "nginx.conf")
certDir := filepath.Join(tempDir, "certs")
if err := os.MkdirAll(certDir, 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
if err := os.WriteFile(mainPath, []byte("old-main"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
if err := os.WriteFile(routePath, []byte("old-route"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
if err := os.WriteFile(filepath.Join(certDir, "1.crt"), []byte("old-cert"), 0o600); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
manager := &Manager{
MainConfigPath: mainPath,
RouteConfigPath: routePath,
CertDir: certDir,
NginxCertDir: "/etc/nginx/openflare-certs",
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
Executor: &scriptedExecutor{
reloadErrors: []error{errors.New("target config failed"), nil},
},
}
outcome := manager.Apply(context.Background(), "new-main", "new-route", []protocol.SupportFile{
{Path: "1.crt", Content: "new-cert"},
})
if outcome.Status != ApplyStatusWarning {
t.Fatalf("expected warning apply outcome, got %#v", outcome)
}
mainData, err := os.ReadFile(mainPath)
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
if string(mainData) != "old-main" {
t.Fatalf("expected main rollback, got %s", string(mainData))
}
}
func TestManagerApplyStartsSafeFallbackWhenNoRollbackConfigExists(t *testing.T) {
tempDir := t.TempDir()
routePath := filepath.Join(tempDir, "routes.conf")
mainPath := filepath.Join(tempDir, "nginx.conf")
executor := &scriptedExecutor{
testErrors: []error{errors.New("target config failed"), errors.New("rollback config missing"), nil},
}
manager := &Manager{
MainConfigPath: mainPath,
RouteConfigPath: routePath,
OpenrestyObservabilityListen: "127.0.0.1:18081",
Executor: executor,
}
outcome := manager.Apply(context.Background(), "bad-main", "bad-route", nil)
if outcome.Status != ApplyStatusWarning {
t.Fatalf("expected warning apply outcome, got %#v", outcome)
}
if !strings.Contains(outcome.Message, "fallback runtime started") {
t.Fatalf("expected fallback message, got %q", outcome.Message)
}
if executor.testCalls != 3 {
t.Fatalf("expected target, rollback, and fallback tests, got %d", executor.testCalls)
}
mainData, err := os.ReadFile(mainPath)
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
if !strings.Contains(string(mainData), "OpenFlare: No Valid Configuration") {
t.Fatalf("expected safe fallback main config, got %s", string(mainData))
}
if !strings.Contains(string(mainData), "listen 80 default_server") {
t.Fatalf("expected fallback to listen on port 80, got %s", string(mainData))
}
if !strings.Contains(string(mainData), "listen 127.0.0.1:18081") {
t.Fatalf("expected fallback to expose local stub_status port, got %s", string(mainData))
}
if !strings.Contains(string(mainData), "stub_status;") {
t.Fatalf("expected fallback to expose stub_status, got %s", string(mainData))
}
routeData, err := os.ReadFile(routePath)
if err != nil {
t.Fatalf("failed to read route config: %v", err)
}
if len(routeData) != 0 {
t.Fatalf("expected fallback route config to be empty, got %q", string(routeData))
}
}
func TestManagerCertFileTargetPathRejectsEscapes(t *testing.T) {
manager := &Manager{CertDir: filepath.Join(t.TempDir(), "certs")}
if err := os.MkdirAll(manager.CertDir, 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
absolutePath := "/tmp/evil.crt"
if runtime.GOOS == "windows" {
absolutePath = `C:/tmp/evil.crt`
}
testCases := []struct {
path string
shouldErr bool
}{
{path: "nested/1.crt", shouldErr: false},
{path: "../escape.crt", shouldErr: true},
{path: "..\\escape.crt", shouldErr: true},
{path: absolutePath, shouldErr: true},
{path: "", shouldErr: true},
}
for _, testCase := range testCases {
targetPath, err := manager.certFileTargetPath(testCase.path)
if testCase.shouldErr {
if err == nil {
t.Fatalf("expected path %q to be rejected, got target %q", testCase.path, targetPath)
}
continue
}
if err != nil {
t.Fatalf("expected path %q to be accepted: %v", testCase.path, err)
}
if !strings.HasPrefix(targetPath, manager.CertDir) {
t.Fatalf("expected target path %q to stay under %q", targetPath, manager.CertDir)
}
}
}
func TestManagerApplyRejectsCertFilePathTraversal(t *testing.T) {
tempDir := t.TempDir()
manager := &Manager{
MainConfigPath: filepath.Join(tempDir, "nginx.conf"),
RouteConfigPath: filepath.Join(tempDir, "routes.conf"),
CertDir: filepath.Join(tempDir, "certs"),
NginxCertDir: "/etc/nginx/openflare-certs",
LuaDir: filepath.Join(tempDir, "lua"),
NginxLuaDir: "/etc/nginx/openflare-lua",
Executor: &fakeExecutor{},
}
outcome := manager.Apply(context.Background(), "main", "route", []protocol.SupportFile{
{Path: "../escape.crt", Content: "bad"},
})
if outcome.Status != ApplyStatusWarning {
t.Fatalf("expected warning apply outcome, got %#v", outcome)
}
if _, statErr := os.Stat(filepath.Join(tempDir, "escape.crt")); !os.IsNotExist(statErr) {
t.Fatalf("expected escaped file to not exist, stat err = %v", statErr)
}
}
func TestManagerSyncWAFIPGroupsWritesDeltaRuntimeFile(t *testing.T) {
manager := &Manager{RuntimeConfigDir: t.TempDir()}
if err := manager.SyncWAFIPGroups([]protocol.WAFIPGroup{
{ID: 1, Enabled: true, IPList: []string{"203.0.113.10"}, Checksum: "sum-1"},
}); err != nil {
t.Fatalf("SyncWAFIPGroups failed: %v", err)
}
if err := manager.SyncWAFIPGroups([]protocol.WAFIPGroup{
{ID: 2, Enabled: true, IPList: []string{"198.51.100.10"}, Checksum: "sum-2"},
}); err != nil {
t.Fatalf("SyncWAFIPGroups second delta failed: %v", err)
}
checksums, err := manager.WAFIPGroupChecksums()
if err != nil {
t.Fatalf("WAFIPGroupChecksums failed: %v", err)
}
if checksums["1"] != "sum-1" || checksums["2"] != "sum-2" {
t.Fatalf("expected merged checksums, got %#v", checksums)
}
data, err := os.ReadFile(filepath.Join(manager.RuntimeConfigDir, WAFIPGroupsConfigFileName))
if err != nil {
t.Fatalf("failed to read runtime ip group file: %v", err)
}
text := string(data)
if !strings.Contains(text, "203.0.113.10") || !strings.Contains(text, "198.51.100.10") {
t.Fatalf("expected runtime file to keep both groups, got %s", text)
}
}
func TestObservabilityListenAddress(t *testing.T) {
if got := ObservabilityListenAddress(18081); got != "127.0.0.1:18081" {
t.Fatalf("unexpected default observability listen address: %s", got)
}
if got := ObservabilityListenAddress(18081); got != "127.0.0.1:18081" {
t.Fatalf("unexpected path observability listen address: %s", got)
}
}
+98
View File
@@ -0,0 +1,98 @@
package nginx
const DefaultMimeTypes = `
types {
text/html html htm shtml;
text/css css;
text/xml xml;
image/gif gif;
image/jpeg jpeg jpg;
application/javascript js;
application/atom+xml atom;
application/rss+xml rss;
text/mathml mml;
text/plain txt;
text/vnd.sun.j2me.app-descriptor jad;
text/vnd.wap.wml wml;
text/x-component htc;
image/png png;
image/svg+xml svg svgz;
image/tiff tif tiff;
image/vnd.wap.wbmp wbmp;
image/webp webp;
image/x-icon ico;
image/x-jng jng;
image/x-ms-bmp bmp;
application/font-woff woff;
application/java-archive jar war ear;
application/json json;
application/mac-binhex40 hqx;
application/msword doc;
application/pdf pdf;
application/postscript ps eps ai;
application/rtf rtf;
application/vnd.apple.mpegurl m3u8;
application/vnd.google-earth.kml+xml kml;
application/vnd.google-earth.kmz kmz;
application/vnd.ms-excel xls;
application/vnd.ms-fontobject eot;
application/vnd.ms-powerpoint ppt;
application/vnd.oasis.opendocument.graphics odg;
application/vnd.oasis.opendocument.presentation odp;
application/vnd.oasis.opendocument.spreadsheet ods;
application/vnd.oasis.opendocument.text odt;
application/vnd.openxmlformats-officedocument.presentationml.presentation
pptx;
application/vnd.openxmlformats-officedocument.spreadsheetml.sheet
xlsx;
application/vnd.openxmlformats-officedocument.wordprocessingml.document
docx;
application/vnd.wap.wmlc wmlc;
application/x-7z-compressed 7z;
application/x-cocoa cco;
application/x-java-archive-diff jardiff;
application/x-java-jnlp-file jnlp;
application/x-makeself run;
application/x-perl pl pm;
application/x-pilot prc pdb;
application/x-rar-compressed rar;
application/x-redhat-package-manager rpm;
application/x-sea sea;
application/x-shockwave-flash swf;
application/x-stuffit sit;
application/x-tcl tcl tk;
application/x-x509-ca-cert der pem crt;
application/x-xpinstall xpi;
application/xhtml+xml xhtml;
application/xspf+xml xspf;
application/zip zip;
application/octet-stream bin exe dll;
application/octet-stream deb;
application/octet-stream dmg;
application/octet-stream iso img;
application/octet-stream msi msp msm;
audio/midi mid midi kar;
audio/mpeg mp3;
audio/ogg ogg;
audio/x-m4a m4a;
audio/x-realaudio ra;
video/3gpp 3gpp 3gp;
video/mp2t ts;
video/mp4 mp4;
video/mpeg mpeg mpg;
video/quicktime mov;
video/webm webm;
video/x-flv flv;
video/x-m4v m4v;
video/x-mng mng;
video/x-ms-asf asx asf;
video/x-ms-wmv wmv;
video/x-msvideo avi;
}
`
@@ -0,0 +1,158 @@
package nginx
import "github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
const (
openRestyObservabilityWindowTTL = "7200"
openRestyObservabilityWindowSize = "60"
)
const openRestyObservabilityInitLua = `local dict = ngx.shared.openflare_observability
if not dict then
return
end
return
`
const openRestyObservabilityLogLua = `local dict = ngx.shared.openflare_observability
if not dict then
return
end
local request_uri = tostring(ngx.var.uri or "")
if request_uri == "/openflare/observability" or request_uri == "/openflare/stub_status" then
return
end
local ttl = ` + openRestyObservabilityWindowTTL + `
local now = ngx.time()
local window_size = ` + openRestyObservabilityWindowSize + `
local window_start = now - (now % window_size)
local function ensure_counter(key)
dict:add(key, 0, ttl)
end
local function incr(key, delta)
ensure_counter(key)
local value, err = dict:incr(key, delta)
if not value and err == "not found" then
dict:set(key, delta, ttl)
end
end
local function remember_value(list_key, marker_key, value)
if value == "" then
return
end
if not dict:add(marker_key, 1, ttl) then
return
end
local existing = dict:get(list_key)
if not existing or existing == "" then
dict:set(list_key, value, ttl)
return
end
dict:set(list_key, existing .. "\n" .. value, ttl)
end
local window_prefix = tostring(window_start)
incr("request_count:" .. window_prefix, 1)
local status = tostring(ngx.status or 0)
if status ~= "0" then
incr("status:" .. window_prefix .. ":" .. status, 1)
remember_value(
"status_keys:" .. window_prefix,
"status_marker:" .. window_prefix .. ":" .. status,
status
)
if tonumber(status) and tonumber(status) >= 500 then
incr("error_count:" .. window_prefix, 1)
end
end
local host = tostring(ngx.var.host or "")
if host ~= "" then
incr("domain:" .. window_prefix .. ":" .. host, 1)
remember_value(
"domain_keys:" .. window_prefix,
"domain_marker:" .. window_prefix .. ":" .. host,
host
)
end
local remote_addr = tostring(ngx.var.binary_remote_addr or ngx.var.remote_addr or "")
if remote_addr ~= "" and dict:add("visitor:" .. window_prefix .. ":" .. remote_addr, 1, ttl) then
incr("unique_visitor_count:" .. window_prefix, 1)
end
local request_length = tonumber(ngx.var.request_length) or 0
if request_length > 0 then
incr("openresty_rx_bytes:" .. window_prefix, request_length)
end
local bytes_sent = tonumber(ngx.var.bytes_sent) or tonumber(ngx.var.body_bytes_sent) or 0
if bytes_sent > 0 then
incr("openresty_tx_bytes:" .. window_prefix, bytes_sent)
end
`
const openRestyObservabilityReadLua = `local cjson = require "cjson.safe"
local dict = ngx.shared.openflare_observability
if not dict then
ngx.status = ngx.HTTP_SERVICE_UNAVAILABLE
ngx.say(cjson.encode({ message = "shared dict unavailable" }))
return
end
local now = ngx.time()
local window_size = ` + openRestyObservabilityWindowSize + `
local window_start = now - (now % window_size)
local current_window = tostring(window_start)
local function read_counter(key)
return tonumber(dict:get(key) or 0) or 0
end
local function read_map(window_id, prefix, list_key)
local result = {}
local raw = dict:get(list_key .. ":" .. window_id)
if not raw or raw == "" then
return result
end
for value in string.gmatch(raw, "[^\n]+") do
result[value] = read_counter(prefix .. ":" .. window_id .. ":" .. value)
end
return result
end
local payload = {
window_started_at_unix = window_start,
window_ended_at_unix = now,
request_count = read_counter("request_count:" .. current_window),
error_count = read_counter("error_count:" .. current_window),
unique_visitor_count = read_counter("unique_visitor_count:" .. current_window),
status_codes = read_map(current_window, "status", "status_keys"),
top_domains = read_map(current_window, "domain", "domain_keys"),
source_countries = {},
openresty_rx_bytes = read_counter("openresty_rx_bytes:" .. current_window),
openresty_tx_bytes = read_counter("openresty_tx_bytes:" .. current_window)
}
ngx.header.content_type = "application/json"
ngx.say(cjson.encode(payload))
`
func ManagedObservabilityLuaFiles() []protocol.SupportFile {
return []protocol.SupportFile{
{Path: "init.lua", Content: openRestyObservabilityInitLua},
{Path: "log.lua", Content: openRestyObservabilityLogLua},
{Path: "read.lua", Content: openRestyObservabilityReadLua},
{Path: "observability/init.lua", Content: openRestyObservabilityInitLua},
{Path: "observability/log.lua", Content: openRestyObservabilityLogLua},
{Path: "observability/read.lua", Content: openRestyObservabilityReadLua},
}
}
+607
View File
@@ -0,0 +1,607 @@
package nginx
import (
"embed"
"path/filepath"
"strings"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
//go:embed pow_static
var powStaticFS embed.FS
const openRestyPowRuntimeLua = `local _M = {}
function _M.check()
local source = debug.getinfo(1, "S").source or ""
if string.sub(source, 1, 1) == "@" then
local script_path = string.sub(source, 2)
local base_dir = string.match(script_path, "^(.*)/pow/[^/]+%.lua$")
if base_dir and base_dir ~= "" and not string.find(package.path, base_dir, 1, true) then
package.path = base_dir .. "/?.lua;" .. base_dir .. "/?/init.lua;" .. package.path
end
end
local cjson = require "cjson.safe"
local policy = require "pow.policy"
local pow_config_dict = ngx.shared.openflare_pow_config
local pow_sessions = ngx.shared.openflare_pow_sessions
local function session_cookie(value, ttl)
local cookie = "__openflare_pow=" .. value .. "; Path=/; HttpOnly; SameSite=Lax; Max-Age=" .. tostring(ttl)
if ngx.var.scheme == "https" then
cookie = cookie .. "; Secure"
end
return cookie
end
-- Lazy-load pow_config from file; reload when content changes
local function load_pow_config()
local config_paths = {
"__OPENFLARE_RUNTIME_CONFIG_DIR__/waf_config.json",
"/etc/nginx/openflare-lua/waf_config.json",
"/usr/local/openresty/nginx/conf/waf_config.json"
}
for _, config_path in ipairs(config_paths) do
local f = io.open(config_path, "r")
if f then
local content = f:read("*a")
f:close()
local current_hash = ngx.md5(content or "")
if current_hash == pow_config_dict:get("_config_hash") then
return
end
-- Clear old domain/site entries
local old_keys = pow_config_dict:get("_domain_keys")
if old_keys then
for domain in string.gmatch(old_keys, "[^\n]+") do
pow_config_dict:delete(domain)
end
end
local domain_keys = {}
if content and content ~= "" and content ~= "{}" then
local ok, decoded = pcall(cjson.decode, content)
if ok and decoded and decoded.rule_groups and decoded.site_rule_groups then
-- Build rule groups map (group ID -> PoWConfig)
local groups = {}
for _, group in ipairs(decoded.rule_groups) do
if group.pow_enabled then
groups[tostring(group.id)] = group.pow_config
end
end
-- Build site name to pow_config map
for site, group_ids in pairs(decoded.site_rule_groups) do
local pow_config = nil
-- Check custom group IDs first
for _, id in ipairs(group_ids) do
pow_config = groups[tostring(id)]
if pow_config then
break
end
end
-- If not found, check global group IDs
if not pow_config then
for _, group in ipairs(decoded.rule_groups) do
if group.is_global and group.pow_enabled then
pow_config = group.pow_config
break
end
end
end
if pow_config then
pow_config_dict:set(site, cjson.encode({enabled = true, config = pow_config}), 0)
domain_keys[#domain_keys+1] = site
end
end
end
end
pow_config_dict:set("_domain_keys", table.concat(domain_keys, "\n"), 0)
pow_config_dict:set("_config_hash", current_hash, 0)
return
end
end
end
load_pow_config()
local host = ngx.var.host
if not host or host == "" then
return
end
local site = ngx.var.openflare_waf_site or ""
if site == "" then
site = host
end
local config_raw = pow_config_dict:get(site)
if not config_raw then
return
end
local ok, route_config = pcall(cjson.decode, config_raw)
if not ok or not route_config then
return
end
if not route_config.enabled then
return
end
local config = route_config.config or {}
local session_ttl = config.session_ttl or 600
local uri = ngx.var.uri or ""
local ua = ngx.var.http_user_agent or ""
local remote_ip = ngx.var.remote_addr or ""
-- Check whitelist: if matched, skip PoW
local whitelist = config.whitelist or {}
if policy.match_any(remote_ip, ua, uri, whitelist) then
return
end
-- Check blacklist: if matched, require PoW
local blacklist = config.blacklist or {}
local has_blacklist = policy.has_entries(blacklist)
local need_pow = false
if has_blacklist then
need_pow = policy.match_any(remote_ip, ua, uri, blacklist)
else
-- No blacklist means all non-whitelisted need PoW
need_pow = true
end
if not need_pow then
return
end
-- Check valid session cookie
local cookie_val = ngx.var["cookie___openflare_pow"]
if cookie_val and cookie_val ~= "" then
local session_key = host .. ":" .. cookie_val
local session_data = pow_sessions:get(session_key)
if session_data then
pow_sessions:set(session_key, "1", session_ttl)
ngx.header["Set-Cookie"] = session_cookie(cookie_val, session_ttl)
return
end
end
-- If requesting the challenge API endpoints, let them through (handled by content_by_lua)
local anubis_api_prefix = "/.within.website/x/cmd/anubis/api/"
local anubis_static_prefix = "/.within.website/x/cmd/anubis/static/"
if string.sub(uri, 1, #anubis_api_prefix) == anubis_api_prefix then
return
end
if string.sub(uri, 1, #anubis_static_prefix) == anubis_static_prefix then
return
end
-- Render the challenge page through an internal redirect so the browser stays
-- on the originally requested URL instead of seeing a 302 hop.
ngx.req.set_uri_args({
redir = ngx.var.scheme .. "://" .. host .. uri .. (ngx.var.args and ("?" .. ngx.var.args) or ""),
host = host
})
return ngx.exec("/.within.website/x/cmd/anubis/api/make-challenge")
end
return _M
`
const openRestyPowCheckLua = `local source = debug.getinfo(1, "S").source or ""
if string.sub(source, 1, 1) == "@" then
local script_path = string.sub(source, 2)
local base_dir = string.match(script_path, "^(.*)/pow/[^/]+%.lua$")
if base_dir and base_dir ~= "" and not string.find(package.path, base_dir, 1, true) then
package.path = base_dir .. "/?.lua;" .. base_dir .. "/?/init.lua;" .. package.path
end
end
return require("pow.runtime").check()
`
const openRestyPowChallengeLua = `local cjson = require "cjson.safe"
local pow_config_dict = ngx.shared.openflare_pow_config
local pow_challenges = ngx.shared.openflare_pow_challenges
local function generate_entropy()
local pieces = {
tostring(ngx.now()),
tostring(ngx.worker.pid()),
tostring(math.random()),
ngx.var.remote_addr or "",
ngx.var.http_user_agent or "",
ngx.var.request_id or "",
}
return table.concat(pieces, ":")
end
local args = ngx.req.get_uri_args()
local host = args["host"] or ngx.var.host or ""
local redir = args["redir"] or ""
local site = ngx.var.openflare_waf_site or ""
if site == "" then
site = host
end
local config_raw = pow_config_dict:get(site)
if not config_raw then
ngx.status = 403
ngx.say("PoW not configured for this site")
return
end
local ok, route_config = pcall(cjson.decode, config_raw)
if not ok or not route_config or not route_config.enabled then
ngx.status = 403
ngx.say("PoW not enabled for this site")
return
end
local config = route_config.config or {}
local difficulty = config.difficulty or 4
local algorithm = config.algorithm or "fast"
local challenge_ttl = config.challenge_ttl or 300
local session_ttl = config.session_ttl or 600
-- Generate challenge data without depending on ngx.random_bytes, which is not
-- available in every OpenResty runtime build.
local entropy = generate_entropy()
local challenge_id = ngx.md5(entropy .. ":id")
local challenge_data = ngx.md5(entropy .. ":data-a") .. ngx.md5(entropy .. ":data-b")
-- Store challenge
local challenge_info = cjson.encode({
data = challenge_data,
difficulty = difficulty,
host = host,
redir = redir,
session_ttl = session_ttl
})
pow_challenges:set(challenge_id, challenge_info, challenge_ttl)
local static_prefix = "/.within.website/x/cmd/anubis/static/"
local accept_lang = ngx.var.http_accept_language or ""
local lang = "en"
if string.find(accept_lang, "zh") then
lang = "zh-CN"
end
local t_title = "Making sure you're not a bot!"
local t_status = "Loading..."
local t_protected = "This site is protected by a Proof-of-Work challenge. Your browser will solve a small puzzle before the upstream response is shown."
local t_why = "Why am I seeing this?"
local t_why_desc = "OpenFlare is asking your browser to complete a lightweight computation to distinguish normal browser traffic from automated abuse. This should finish automatically."
local t_noscript = "JavaScript is required to pass this verification. Please enable JavaScript and reload."
if lang == "zh-CN" then
t_title = "正在确认你是不是机器人!"
t_status = "加载中..."
t_protected = "本网站受工作量证明(Proof-of-Work)挑战保护。在显示源站响应之前,您的浏览器将解决一个微型谜题。"
t_why = "为什么我会看到这个?"
t_why_desc = "OpenFlare 正在要求您的浏览器完成一项轻量级计算,以区分正常的浏览器流量和自动化的恶意请求。这应该会自动完成。"
t_noscript = "很遗憾,您必须启用 JavaScript 才能通过这项验证。请开启 JavaScript 并刷新页面。"
end
ngx.header.content_type = "text/html; charset=utf-8"
ngx.say([[<!DOCTYPE html>
<html lang="]] .. lang .. [[">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="robots" content="noindex,nofollow">
<title>]] .. t_title .. [[</title>
<link rel="stylesheet" href="]] .. static_prefix .. [[css/xess.css">
<style>
body,html{height:100%;display:flex;justify-content:center;align-items:center;margin-left:auto;margin-right:auto}
.centered-div{text-align:center}
#status{font-variant-numeric:tabular-nums}
#progress{display:none;width:min(20rem,90%);height:2rem;border-radius:1rem;overflow:hidden;margin:1rem 0 2rem;outline-offset:2px;outline:#b16286 solid 4px}
.bar-inner{background-color:#b16286;height:100%;width:0;transition:width .25s ease-in}
</style>
<script id="anubis_version" type="application/json">"openflare-pow"</script>
<script id="anubis_challenge" type="application/json">]] .. cjson.encode({
challenge = {
id = challenge_id,
randomData = challenge_data,
method = algorithm
},
rules = {
difficulty = difficulty,
algorithm = algorithm
}
}) .. [[</script>
<script id="anubis_base_prefix" type="application/json">""</script>
<script id="anubis_public_url" type="application/json">"__openflare_internal__"</script>
</head>
<body id="top">
<main>
<h1 id="title" class="centered-div">]] .. t_title .. [[</h1>
<div class="centered-div">
<img id="image" style="width:100%;max-width:256px;" src="]] .. static_prefix .. [[img/pensive.webp?cacheBuster=openflare-pow">
<p id="status">]] .. t_status .. [[</p>
<p>]] .. t_protected .. [[</p>
<div id="progress" role="progressbar" aria-labelledby="status"><div class="bar-inner"></div></div>
<details>
<summary>]] .. t_why .. [[</summary>
<p>]] .. t_why_desc .. [[</p>
</details>
<noscript><p>]] .. t_noscript .. [[</p></noscript>
</div>
</main>
<script type="module" src="]] .. static_prefix .. [[js/main.mjs"></script>
</body>
</html>]])
`
const openRestyPowVerifyLua = `local cjson = require "cjson.safe"
local pow_challenges = ngx.shared.openflare_pow_challenges
local pow_sessions = ngx.shared.openflare_pow_sessions
local args = ngx.req.get_uri_args()
local challenge_id = args["id"] or ""
local response = args["response"] or ""
local nonce_str = args["nonce"] or ""
local redir = args["redir"] or ""
local elapsed = args["elapsedTime"] or ""
if challenge_id == "" or response == "" or nonce_str == "" then
ngx.status = 400
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "missing parameters"}))
return
end
local nonce = tonumber(nonce_str)
if not nonce then
ngx.status = 400
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "invalid nonce"}))
return
end
-- Get stored challenge
local challenge_raw = pow_challenges:get(challenge_id)
if not challenge_raw then
ngx.status = 410
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "challenge expired or not found"}))
return
end
local ok, challenge_info = pcall(cjson.decode, challenge_raw)
if not ok or not challenge_info then
ngx.status = 500
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "invalid challenge data"}))
return
end
local challenge_data = challenge_info.data or ""
local difficulty = challenge_info.difficulty or 4
local host = challenge_info.host or ngx.var.host or ""
local session_ttl = challenge_info.session_ttl or 600
-- Compute SHA-256(challenge_data + nonce)
local calc_string = challenge_data .. tostring(math.floor(nonce))
local calculated = ngx.sha1_bin ~= nil and "" or ""
-- Use resty.sha256 for proper SHA-256
local sha256 = require "resty.sha256"
local str = require "resty.string"
local hasher = sha256:new()
hasher:update(calc_string)
local hash_bytes = hasher:final()
local hash_hex = str.to_hex(hash_bytes)
-- Verify hash matches response
if hash_hex ~= string.lower(response) then
ngx.status = 403
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "hash mismatch"}))
return
end
-- Verify difficulty (leading zeros in hex)
local prefix = string.rep("0", difficulty)
if string.sub(hash_hex, 1, difficulty) ~= prefix then
ngx.status = 403
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({error = "insufficient difficulty"}))
return
end
-- Invalidate challenge (prevent replay)
pow_challenges:delete(challenge_id)
-- Generate session token
local session_token = str.to_hex(ngx.sha1_bin(challenge_id .. ngx.now() .. tostring(ngx.worker.pid())))
-- Store session
pow_sessions:set(host .. ":" .. session_token, "1", session_ttl)
-- Set cookie. Secure cookies are not sent over HTTP, so only add Secure when
-- the current request itself is HTTPS.
local cookie = "__openflare_pow=" .. session_token .. "; Path=/; HttpOnly; SameSite=Lax; Max-Age=" .. tostring(session_ttl)
if ngx.var.scheme == "https" then
cookie = cookie .. "; Secure"
end
ngx.header["Set-Cookie"] = cookie
if redir ~= "" then
return ngx.redirect(redir)
end
ngx.header.content_type = "application/json"
ngx.say(cjson.encode({ok = true}))
`
const openRestyPowPolicyLua = `local M = {}
local function match_ip(remote_ip, ips)
if not ips or #ips == 0 then return false end
for _, ip in ipairs(ips) do
if ip == remote_ip then
return true
end
end
return false
end
local function match_cidr(remote_ip, cidrs)
if not cidrs or #cidrs == 0 then return false end
for _, cidr in ipairs(cidrs) do
local m, err = ngx.re.match(cidr, "^(\\\\d{1,3}\\\\.\\\\d{1,3}\\\\.\\\\d{1,3}\\\\.\\\\d{1,3})/(\\\\d{1,2})$")
if m then
local mask_bits = tonumber(m[2])
if mask_bits and mask_bits >= 0 and mask_bits <= 32 then
local function ip_to_num(ip_str)
local parts = {}
for part in string.gmatch(ip_str, "%d+") do
parts[#parts+1] = tonumber(part) or 0
end
if #parts ~= 4 then return 0 end
return parts[1]*16777216 + parts[2]*65536 + parts[3]*256 + parts[4]
end
local remote_num = ip_to_num(remote_ip)
local net_num = ip_to_num(m[1])
if mask_bits == 0 then
return true
end
local mask = math.floor(2^(32 - mask_bits))
mask = 4294967296 - mask
if bit.band(remote_num, mask) == bit.band(net_num, mask) then
return true
end
end
end
end
return false
end
local function match_path(uri, patterns)
if not patterns or #patterns == 0 then return false end
for _, pattern in ipairs(patterns) do
local ok, match = pcall(ngx.re.match, uri, "^" .. ngx.re.gsub(pattern, "([%^%$%(%)%%%.%[%]%+%-%?])", function(c)
if c == "*" then return ".*" end
return "%" .. c
end) .. "$", "i")
if ok and match then
return true
end
end
return false
end
local function match_path_regex(uri, patterns)
if not patterns or #patterns == 0 then return false end
for _, pattern in ipairs(patterns) do
local ok, match = pcall(ngx.re.match, uri, pattern)
if ok and match then
return true
end
end
return false
end
local function match_ua(ua, patterns)
if not patterns or #patterns == 0 then return false end
for _, pattern in ipairs(patterns) do
if ua and string.find(ua, pattern, 1, true) then
return true
end
end
return false
end
function M.match_any(remote_ip, ua, uri, list)
if not list then return false end
if match_ip(remote_ip, list.ips) then return true end
if match_cidr(remote_ip, list.ip_cidrs) then return true end
if match_path(uri, list.paths) then return true end
if match_path_regex(uri, list.path_regexes) then return true end
if match_ua(ua, list.user_agents) then return true end
return false
end
function M.has_entries(list)
if not list then return false end
return (#(list.ips or {}) + #(list.ip_cidrs or {}) + #(list.paths or {}) + #(list.path_regexes or {}) + #(list.user_agents or {})) > 0
end
return M
`
func ManagedPowLuaFiles() []protocol.SupportFile {
return []protocol.SupportFile{
{Path: "pow/runtime.lua", Content: openRestyPowRuntimeLua},
{Path: "pow/check.lua", Content: openRestyPowCheckLua},
{Path: "pow/challenge.lua", Content: openRestyPowChallengeLua},
{Path: "pow/verify.lua", Content: openRestyPowVerifyLua},
{Path: "pow/policy.lua", Content: openRestyPowPolicyLua},
}
}
func ManagedPowStaticFiles() ([]protocol.SupportFile, error) {
var files []protocol.SupportFile
entries, err := powStaticFS.ReadDir("pow_static")
if err != nil {
return nil, err
}
var walk func(dir string) error
walk = func(dir string) error {
entries, err := powStaticFS.ReadDir(dir)
if err != nil {
return err
}
for _, entry := range entries {
fullPath := filepath.Join(dir, entry.Name())
if entry.IsDir() {
if err := walk(fullPath); err != nil {
return err
}
continue
}
data, err := powStaticFS.ReadFile(fullPath)
if err != nil {
return err
}
// Convert pow_static/css/xess.css -> pow/static/css/xess.css
relPath := strings.TrimPrefix(fullPath, "pow_static/")
files = append(files, protocol.SupportFile{
Path: "pow/static/" + relPath,
Content: string(data),
})
}
return nil
}
for _, entry := range entries {
fullPath := filepath.Join("pow_static", entry.Name())
if entry.IsDir() {
if err := walk(fullPath); err != nil {
return nil, err
}
} else {
data, err := powStaticFS.ReadFile(fullPath)
if err != nil {
return nil, err
}
relPath := strings.TrimPrefix(fullPath, "pow_static/")
files = append(files, protocol.SupportFile{
Path: "pow/static/" + relPath,
Content: string(data),
})
}
}
return files, nil
}
@@ -0,0 +1,7 @@
@font-face {
font-family: "Podkova";
font-style: normal;
font-weight: 400 800;
font-display: swap;
src: url("podkova.woff2") format("woff2");
}
@@ -0,0 +1,149 @@
:root {
--body-sans-font: Geist, sans-serif;
--body-preformatted-font: Iosevka Curly Iaso, monospace;
--body-title-font: Podkova, serif;
--background: #1d2021;
--text: #f9f5d7;
--text-selection: #d3869b;
--preformatted-background: #3c3836;
--link-foreground: #b16286;
--link-background: #282828;
--blockquote-border-left: 1px solid #bdae93;
--progress-bar-outline: #b16286 solid 4px;
--progress-bar-fill: #b16286;
}
@media (prefers-color-scheme: light) {
:root {
--background: #f9f5d7;
--text: #1d2021;
--text-selection: #d3869b;
--preformatted-background: #ebdbb2;
--link-foreground: #b16286;
--link-background: #fbf1c7;
--blockquote-border-left: 1px solid #655c54;
}
}
@font-face {
font-family: "Geist";
font-style: normal;
font-weight: 100 900;
font-display: swap;
src: url("./static/geist.woff2") format("woff2");
}
@font-face {
font-family: "Podkova";
font-style: normal;
font-weight: 400 800;
font-display: swap;
src: url("./static/podkova.woff2") format("woff2");
}
@font-face {
font-family: "Iosevka Curly";
font-style: monospace;
font-display: swap;
src: url("./static/iosevka-curly.woff2") format("woff2");
}
main {
font-family: var(--body-sans-font);
max-width: 50rem;
padding: 2rem;
margin: auto;
}
::selection {
background: var(--text-selection);
}
body {
background: var(--background);
color: var(--text);
}
body,
html {
height: 100%;
display: flex;
justify-content: center;
align-items: center;
margin-left: auto;
margin-right: auto;
}
.centered-div {
text-align: center;
}
#status {
font-variant-numeric: tabular-nums;
}
.centered-div {
text-align: center;
}
#status {
font-variant-numeric: tabular-nums;
}
#progress {
display: none;
width: min(20rem, 90%);
height: 2rem;
border-radius: 1rem;
overflow: hidden;
margin: 1rem 0 2rem;
outline-offset: 2px;
outline: var(--progress-bar-outline);
}
.bar-inner {
background-color: var(--progress-bar-fill);
height: 100%;
width: 0;
transition: width 0.25s ease-in;
}
@media (prefers-reduced-motion: no-preference) {
.bar-inner {
transition: width 0.25s ease-in;
}
}
pre {
background-color: var(--preformatted-background);
padding: 1em;
border: 0;
font-family: var(--body-preformatted-font);
}
a,
a:active,
a:visited {
color: var(--link-foreground);
background-color: var(--link-background);
}
h1,
h2,
h3,
h4,
h5 {
margin-bottom: 0.1rem;
font-family: var(--body-title-font);
}
blockquote {
border-left: var(--blockquote-border-left);
margin: 0.5em 10px;
padding: 0.5em 10px;
}
footer {
text-align: center;
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 30 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 28 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 26 KiB

@@ -0,0 +1,32 @@
/*
@licstart The following is the entire license notice for the
JavaScript code in this page.
Copyright (c) 2025 Xe Iaso <xe.iaso@techaro.lol>
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
Includes code from https://github.com/aws/aws-sdk-js-crypto-helpers which is
used under the terms of the Apache 2 license.
@licend The above is the entire license notice
for the JavaScript code in this page.
*/
(()=>{var k=()=>navigator.hardwareConcurrency!==void 0?navigator.hardwareConcurrency:1;function n(c,b,w=5,e=null,g,u=Math.trunc(Math.max(k()/2,1))){console.debug("fast algo");let s="purejs";return window.isSecureContext&&(s="webcrypto"),(navigator.userAgent.includes("Firefox")||navigator.userAgent.includes("Goanna"))&&(console.log("Firefox detected, using pure-JS fallback"),s="purejs"),new Promise((p,l)=>{let m=`${c.basePrefix}/.within.website/x/cmd/anubis/static/js/worker/sha256-${s}.mjs?cacheBuster=${c.version}`,f=[],d=!1,a=()=>{console.log("PoW aborted"),i(),l(new DOMException("Aborted","AbortError"))},i=()=>{d||(d=!0,f.forEach(r=>r.terminate()),e?.removeEventListener("abort",a))};if(e!=null){if(e.aborted)return a();e.addEventListener("abort",a,{once:!0})}for(let r=0;r<u;r++){let t=new Worker(m);t.onmessage=o=>{typeof o.data=="number"?g?.(o.data):(i(),p(o.data))},t.onerror=o=>{i(),l(o)},t.postMessage({data:b,difficulty:w,nonce:r,threads:u}),f.push(t)}})}var P={fast:n,slow:n};})();
//# sourceMappingURL=index.mjs.map
@@ -0,0 +1,32 @@
/*
@licstart The following is the entire license notice for the
JavaScript code in this page.
Copyright (c) 2025 Xe Iaso <xe.iaso@techaro.lol>
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
Includes code from https://github.com/aws/aws-sdk-js-crypto-helpers which is
used under the terms of the Apache 2 license.
@licend The above is the entire license notice
for the JavaScript code in this page.
*/
(()=>{var I=()=>navigator.hardwareConcurrency!==void 0?navigator.hardwareConcurrency:1;function _(e,n,s=5,o=null,i,u=Math.trunc(Math.max(I()/2,1))){console.debug("fast algo");let a="purejs";return window.isSecureContext&&(a="webcrypto"),(navigator.userAgent.includes("Firefox")||navigator.userAgent.includes("Goanna"))&&(console.log("Firefox detected, using pure-JS fallback"),a="purejs"),new Promise((E,x)=>{let M=`${e.basePrefix}/.within.website/x/cmd/anubis/static/js/worker/sha256-${a}.mjs?cacheBuster=${e.version}`,p=[],d=!1,b=()=>{console.log("PoW aborted"),h(),x(new DOMException("Aborted","AbortError"))},h=()=>{d||(d=!0,p.forEach(c=>c.terminate()),o?.removeEventListener("abort",b))};if(o!=null){if(o.aborted)return b();o.addEventListener("abort",b,{once:!0})}for(let c=0;c<u;c++){let g=new Worker(M);g.onmessage=m=>{typeof m.data=="number"?i?.(m.data):(h(),E(m.data))},g.onerror=m=>{h(),x(m)},g.postMessage({data:n,difficulty:s,nonce:c,threads:u}),p.push(g)}})}var j={fast:_,slow:_};var v=(e="",n={})=>{let s=new URL(e,window.location.href);return Object.entries(n).forEach(([o,i])=>s.searchParams.set(o,i)),s.toString()},L=e=>{let n=document.getElementById(e);return n===null?null:JSON.parse(n.textContent)},k=(e,n,s)=>v(`${s}/.within.website/x/cmd/anubis/static/img/${e}.webp`,{cacheBuster:n});var W=async()=>document.documentElement.lang,S=async e=>{let n=L("anubis_base_prefix");if(n!==null)try{return await(await fetch(`${n}/.within.website/x/cmd/anubis/static/locales/${e}.json`)).json()}catch(s){if(console.warn(`Failed to load translations for ${e}, falling back to English`),e!=="en")return await S("en");throw s}},C=()=>{let e=L("anubis_public_url");if(e!==null)return e&&window.location.href.startsWith(e)?new URLSearchParams(window.location.search).get("redir"):window.location.href},$={},D,A=async()=>{D=await W(),$=await S(D)},r=e=>$[`js_${e}`]||$[e]||e;(async()=>{await A();let e=[{name:"Web Workers",msg:r("web_workers_error"),value:window.Worker},{name:"Cookies",msg:r("cookies_error"),value:navigator.cookieEnabled}],n=document.getElementById("status"),s=document.getElementById("image"),o=document.getElementById("title"),i=document.getElementById("progress"),u=L("anubis_version"),a=L("anubis_base_prefix"),E=document.querySelector("details"),x=!1;E&&E.addEventListener("toggle",()=>{E.open&&(x=!0)});let M=({titleMsg:l,statusMsg:f,imageSrc:w})=>{o.innerHTML=l,n.innerHTML=f,s.src=w,i.style.display="none"};n.innerHTML=r("calculating");for(let{value:l,name:f,msg:w}of e)if(!l){M({titleMsg:`${r("missing_feature")} ${f}`,statusMsg:w,imageSrc:k("reject",u,a)});return}let{challenge:p,rules:d}=L("anubis_challenge"),b=j[d.algorithm];if(!b){M({titleMsg:r("challenge_error"),statusMsg:r("challenge_error_msg"),imageSrc:k("reject",u,a)});return}n.innerHTML=`${r("calculating_difficulty")} ${d.difficulty}, `,i.style.display="inline-block";let h=document.createTextNode(`${r("speed")} 0kH/s`);n.appendChild(h);let c=0,g=!1,m=Math.pow(16,-d.difficulty);try{let l=Date.now(),{hash:f,nonce:w}=await b({basePrefix:a,version:u},p.randomData,d.difficulty,null,t=>{let y=Date.now()-l;y-c>1e3&&(c=y,h.data=`${r("speed")} ${(t/y).toFixed(3)}kH/s`);let T=Math.pow(1-m,t),P=(1-Math.pow(T,2))*100;i["aria-valuenow"]=P,i.firstElementChild!==null&&(i.firstElementChild.style.width=`${P}%`),T<.1&&!g&&(n.append(document.createElement("br"),document.createTextNode(r("verification_longer"))),g=!0)}),H=Date.now();if(console.log({hash:f,nonce:w}),x){let y=function(){let T=C();window.location.replace(v(`${a}/.within.website/x/cmd/anubis/api/pass-challenge`,{id:p.id,response:f,nonce:w,redir:T,elapsedTime:H-l}))},t=document.getElementById("progress");t.style.display="flex",t.style.alignItems="center",t.style.justifyContent="center",t.style.height="2rem",t.style.borderRadius="1rem",t.style.cursor="pointer",t.style.background="#b16286",t.style.color="white",t.style.fontWeight="bold",t.style.outline="4px solid #b16286",t.style.outlineOffset="2px",t.style.width="min(20rem, 90%)",t.style.margin="1rem auto 2rem",t.innerHTML=r("finished_reading"),t.onclick=y,setTimeout(y,3e4)}else{let t=C();window.location.replace(v(`${a}/.within.website/x/cmd/anubis/api/pass-challenge`,{id:p.id,response:f,nonce:w,redir:t,elapsedTime:H-l}))}}catch(l){M({titleMsg:r("calculation_error"),statusMsg:`${r("calculation_error_msg")} ${l.message}`,imageSrc:k("reject",u,a)})}})();})();
//# sourceMappingURL=main.mjs.map
File diff suppressed because one or more lines are too long
@@ -0,0 +1,32 @@
/*
@licstart The following is the entire license notice for the
JavaScript code in this page.
Copyright (c) 2025 Xe Iaso <xe.iaso@techaro.lol>
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
Includes code from https://github.com/aws/aws-sdk-js-crypto-helpers which is
used under the terms of the Apache 2 license.
@licend The above is the entire license notice
for the JavaScript code in this page.
*/
(()=>{var h=new TextEncoder,y=async e=>{let s=h.encode(e);return await crypto.subtle.digest("SHA-256",s)},g=e=>e.reduce((s,a)=>s+a.toString(16).padStart(2,"0"),"");addEventListener("message",async({data:e})=>{let{data:s,difficulty:a,threads:d}=e,t=e.nonce,f=t===0,o=0,c=Math.floor(a/2),l=a%2!==0;for(;;){let u=await y(s+t),i=new Uint8Array(u),r=!0;for(let n=0;n<c;n++)if(i[n]!==0){r=!1;break}if(r&&l&&i[c]>>4!==0&&(r=!1),r){let n=g(i);postMessage({hash:n,data:s,difficulty:a,nonce:t});return}t+=d,o++,t%1!==0&&(t=Math.trunc(t)),f&&(o&1023)===0&&postMessage(t)}});})();
//# sourceMappingURL=sha256-webcrypto.mjs.map
@@ -0,0 +1,66 @@
{
"loading": "Loading...",
"why_am_i_seeing": "Why am I seeing this?",
"protected_by": "Protected by",
"protected_from": "From",
"made_with": "Made with ❤️ in 🇨🇦",
"mascot_design": "Mascot design by",
"ai_companies_explanation": "You are seeing this because the administrator of this website has set up Anubis to protect the server against the scourge of AI companies aggressively scraping websites. This can and does cause downtime for the websites, which makes their resources inaccessible for everyone.",
"anubis_compromise": "Anubis is a compromise. Anubis uses a Proof-of-Work scheme in the vein of Hashcash, a proposed proof-of-work scheme for reducing email spam. The idea is that at individual scales the additional load is ignorable, but at mass scraper levels it adds up and makes scraping much more expensive.",
"hack_purpose": "Ultimately, this is a placeholder solution so that more time can be spent on fingerprinting and identifying headless browsers (EG: via how they do font rendering) so that the challenge proof of work page doesn't need to be presented to users that are much more likely to be legitimate.",
"simplified_explanation": "This is a measure against bots and malicious requests similar to a CAPTCHA. However, instead of having to do work yourself, your browser is given a calculation task that it has to solve to ensure that it is a valid client. This concept is called <a href=\"https://en.wikipedia.org/wiki/Proof_of_work\">Proof of Work</a>. The task is calculated in a few seconds and you are granted access to the website. Thank you for your understanding and patience.",
"jshelter_note": "Please note that Anubis requires the use of modern JavaScript features that plugins like JShelter will disable. Please disable JShelter or other such plugins for this domain.",
"version_info": "This website is running Anubis version",
"try_again": "Try again",
"go_home": "Go home",
"contact_webmaster": "or if you believe you should not be blocked, please contact the webmaster at",
"connection_security": "Please wait a moment while we ensure the security of your connection.",
"javascript_required": "Sadly, you must enable JavaScript to get past this challenge. This is required because AI companies have changed the social contract around how website hosting works. A no-JS solution is a work-in-progress.",
"benchmark_requires_js": "Running the benchmark tool requires JavaScript to be enabled.",
"difficulty": "Difficulty:",
"algorithm": "Algorithm:",
"compare": "Compare:",
"time": "Time",
"iters": "Iters",
"time_a": "Time A",
"iters_a": "Iters A",
"time_b": "Time B",
"iters_b": "Iters B",
"static_check_endpoint": "This is just a check endpoint for your reverse proxy to use.",
"authorization_required": "Authorization required",
"cookies_disabled": "Your browser is configured to disable cookies. Anubis requires cookies for the legitimate interest of making sure you are a valid client. Please enable cookies for this domain",
"access_denied": "Access Denied: error code",
"dronebl_entry": "DroneBL reported an entry",
"see_dronebl_lookup": "see",
"internal_server_error": "Internal Server Error: administrator has misconfigured Anubis. Please contact the administrator and ask them to look for the logs around",
"invalid_redirect": "Invalid redirect",
"redirect_not_parseable": "Redirect URL not parseable",
"redirect_domain_not_allowed": "Redirect domain not allowed",
"missing_required_forwarded_headers": "Missing required X-Forwarded-* headers",
"failed_to_sign_jwt": "failed to sign JWT",
"invalid_invocation": "Invalid invocation of MakeChallenge",
"client_error_browser": "Client Error: Please ensure your browser is up to date and try again later.",
"oh_noes": "Oh noes!",
"benchmarking_anubis": "Benchmarking Anubis!",
"you_are_not_a_bot": "You are not a bot!",
"making_sure_not_bot": "Making sure you're not a bot!",
"celphase": "CELPHASE",
"js_web_crypto_error": "Your browser doesn't have a functioning web.crypto element. Are you viewing this over a secure context?",
"js_web_workers_error": "Your browser doesn't support web workers (Anubis uses this to avoid freezing your browser). Do you have a plugin like JShelter installed?",
"js_cookies_error": "Your browser doesn't store cookies. Anubis uses cookies to determine which clients have passed challenges by storing a signed token in a cookie. Please enable storing cookies for this domain. The names of the cookies Anubis stores may vary without notice. Cookie names and values are not part of the public API.",
"js_context_not_secure": "Your context is not secure!",
"js_context_not_secure_msg": "Try connecting over HTTPS or let the admin know to set up HTTPS. For more information, see <a href=\"https://developer.mozilla.org/en-US/docs/Web/Security/Secure_Contexts#when_is_a_context_considered_secure\">MDN</a>.",
"js_calculating": "Calculating...",
"js_missing_feature": "Missing feature",
"js_challenge_error": "Challenge error!",
"js_challenge_error_msg": "Failed to resolve check algorithm. You may want to reload the page.",
"js_calculating_difficulty": "Calculating...<br/>Difficulty:",
"js_speed": "Speed:",
"js_verification_longer": "Verification is taking longer than expected. Please do not refresh the page.",
"js_success": "Success!",
"js_done_took": "Done! Took",
"js_iterations": "iterations",
"js_finished_reading": "I've finished reading, continue →",
"js_calculation_error": "Calculation error!",
"js_calculation_error_msg": "Failed to calculate challenge:"
}
@@ -0,0 +1,66 @@
{
"loading": "加载中...",
"why_am_i_seeing": "为什么我会看到这个?",
"protected_by": "本网站由",
"protected_from": "保护,来自",
"made_with": "在 🇨🇦 用 ❤️ 制作",
"mascot_design": "吉祥物由",
"ai_companies_explanation": "您会看到这个画面,是因为网站管理员启用了 Anubis 来保护服务器,避免 AI 公司大量爬取网站内容。这类行为会导致网站崩溃,让所有用户都无法正常访问资源。",
"anubis_compromise": "Anubis 是一种折中做法。它采用了类似 Hashcash 的工作量证明机制(Proof-of-Work),该机制最初是为了减少垃圾邮件而提出。其核心概念是:对个别用户而言,额外的计算负担可以忽略,但对大规模爬虫来说,累积起来的成本将大幅增加,从而让爬取行为变得更困难。",
"hack_purpose": "最终,这是一个占位符解决方案,以便将更多时间用于指纹识别和识别无头浏览器(例如:通过它们如何进行字体渲染),从而无需向更可能是合法用户的用户呈现挑战工作量证明页面。",
"jshelter_note": "请注意,Anubis 需要使用现代 JavaScript 功能,而像 JShelter 这类插件可能会阻挡这些功能。请为此域名停用 JShelter 或类似的插件。",
"version_info": "这个网站正在运行的 Anubis 版本为",
"try_again": "再试一次",
"go_home": "返回首页",
"contact_webmaster": "或者您觉得您不应该被封锁,请联系网站管理员于",
"connection_security": "请稍等,我们需要在继续之前检查您的连接安全性。",
"javascript_required": "很遗憾,您必须启用 JavaScript 才能通过这项验证。这是因为 AI 公司已经改变了网站托管的社会契约,因此我们必须采取这样的保护机制。无需 JavaScript 的解决方案仍在开发中。",
"benchmark_requires_js": "运行基准测试工具需要启用 JavaScript。",
"difficulty": "难度:",
"algorithm": "算法:",
"compare": "比较:",
"time": "时间",
"iters": "迭代",
"time_a": "时间 A",
"iters_a": "迭代 A",
"time_b": "时间 B",
"iters_b": "迭代 B",
"static_check_endpoint": "这是提供给您的反向代理服务器使用的检查端点。",
"authorization_required": "需要认证",
"cookies_disabled": "您的浏览器目前已禁用 Cookie,为了确认您是合法用户,Anubis 需要启用 Cookie。 请您为此域名启用 Cookie",
"access_denied": "拒绝访问:错误代码",
"dronebl_entry": "DroneBL 报告了一条记录",
"see_dronebl_lookup": "见",
"internal_server_error": "内部服务器错误:管理员错误地配置了 Anubis。 请联系管理员要求他们检查日志",
"invalid_redirect": "无效的重定向",
"redirect_not_parseable": "重定向 URL 无法解析",
"redirect_domain_not_allowed": "重定向的域名并不允许",
"failed_to_sign_jwt": "签署 JWT 失败",
"invalid_invocation": "无效的 MakeChallenge 调用",
"client_error_browser": "客户端错误:请确保您的浏览器是最新版本并稍候再试。",
"oh_noes": "哎呀糟糕了!",
"benchmarking_anubis": "正在进行 Anubis 性能测试!",
"you_are_not_a_bot": "你不是机器人!",
"making_sure_not_bot": "正在确认你是不是机器人!",
"celphase": "CELPHASE 设计",
"js_web_crypto_error": "您的浏览器无法正常使用 web.crypto 组件。您是否通过安全连接(HTTPS)查看此网站?",
"js_web_workers_error": "您的浏览器并不支持 Web workers (Anubis 使用这个来避免冻结您的浏览器 )您有安装像是 JShelter 之类的插件吗?",
"js_cookies_error": "您的浏览器无法存储 Cookie。 Anubis 会使用 Cookie 存储签署的凭证,以判断用户是否已通过验证。请为此域名启用 Cookie 存储功能。 请注意,Anubis 存储的 Cookie 名称可能会变动,且其名称与内容不属于公开 API 的一部分。",
"js_context_not_secure": "您的内容并不安全",
"js_context_not_secure_msg": "请尝试使用 HTTPS 连接,或联系网站管理员设置 HTTPS。更多信息请参见 <a href=\"https://developer.mozilla.org/en-US/docs/Web/Security/Secure_Contexts#when_is_a_context_considered_secure\">MDN</a>。",
"js_calculating": "计算中...",
"js_missing_feature": "缺少功能",
"js_challenge_error": "挑战错误!",
"js_challenge_error_msg": "解决检查算法失败。 您可能会想要刷新页面。",
"js_calculating_difficulty": "计算中...<br/>难度:",
"js_speed": "速度:",
"js_verification_longer": "验证所花的时间高于预期。 请不要刷新页面。",
"js_success": "成功!",
"js_done_took": "完成! 花费",
"js_iterations": "迭代",
"js_finished_reading": "我读完了,继续 →",
"js_calculation_error": "计算错误!",
"js_calculation_error_msg": "计算挑战失败:",
"missing_required_forwarded_headers": "缺少必要的 X-Forwarded-* 头",
"simplified_explanation": "这是一种类似于验证码的措施,用于防止机器人和恶意请求。但是,您无需自己动手,您的浏览器会收到一个计算任务,必须解决该任务以确保它是有效的客户端。这个概念称为<a href=\"https://en.wikipedia.org/wiki/Proof_of_work\">工作量证明</a>。该任务在几秒钟内计算完毕,您将被授予访问网站的权限。感谢您的理解和耐心。"
}
+310
View File
@@ -0,0 +1,310 @@
package nginx
import "github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
const openRestyWAFRuntimeLua = `local _M = {}
function _M.check()
local cjson = require "cjson.safe"
local config_dict = ngx.shared.openflare_waf_config
local function read_file(path)
local f = io.open(path, "r")
if not f then
return nil
end
local content = f:read("*a")
f:close()
return content
end
local function load_config()
local paths = {
"__OPENFLARE_RUNTIME_CONFIG_DIR__/waf_config.json",
"/etc/nginx/openflare-lua/waf_config.json",
"/usr/local/openresty/nginx/conf/waf_config.json"
}
for _, path in ipairs(paths) do
local content = read_file(path)
if content and content ~= "" then
local hash = ngx.md5(content)
if config_dict:get("_config_hash") == hash then
local cached = config_dict:get("_config_json")
if cached then
local decoded = cjson.decode(cached)
if decoded then
return decoded
end
end
end
local decoded = cjson.decode(content)
if decoded then
config_dict:set("_config_hash", hash, 0)
config_dict:set("_config_json", content, 0)
return decoded
end
end
end
return nil
end
local function load_ip_groups()
local paths = {
"__OPENFLARE_RUNTIME_CONFIG_DIR__/waf_ip_groups.json",
"/etc/nginx/openflare-lua/waf_ip_groups.json",
"/usr/local/openresty/nginx/conf/waf_ip_groups.json"
}
for _, path in ipairs(paths) do
local content = read_file(path)
if content and content ~= "" then
local hash = ngx.md5(content)
if config_dict:get("_ip_groups_hash") == hash then
local cached = config_dict:get("_ip_groups_json")
if cached then
local decoded = cjson.decode(cached)
if decoded then
return decoded
end
end
end
local decoded = cjson.decode(content)
if decoded then
config_dict:set("_ip_groups_hash", hash, 0)
config_dict:set("_ip_groups_json", content, 0)
return decoded
end
end
end
return { groups = {} }
end
local function list_contains(items, value)
if not items or type(items) ~= "table" or not value or value == "" then
return false
end
for _, item in ipairs(items) do
if item == value then
return true
end
end
return false
end
local function table_has_items(items)
return type(items) == "table" and #items > 0
end
local function parse_ipv4(value)
local a, b, c, d = string.match(value or "", "^(%d+)%.(%d+)%.(%d+)%.(%d+)$")
if not a then
return nil
end
a, b, c, d = tonumber(a), tonumber(b), tonumber(c), tonumber(d)
if a > 255 or b > 255 or c > 255 or d > 255 then
return nil
end
return ((a * 256 + b) * 256 + c) * 256 + d
end
local function ipv4_in_cidr(ip, cidr)
local base, bits = string.match(cidr or "", "^([^/]+)/(%d+)$")
if not base then
return false
end
bits = tonumber(bits)
if not bits or bits < 0 or bits > 32 then
return false
end
local ip_num = parse_ipv4(ip)
local base_num = parse_ipv4(base)
if not ip_num or not base_num then
return false
end
if bits == 0 then
return true
end
local mask = 4294967295 - (2 ^ (32 - bits) - 1)
return (ip_num - (ip_num % (2 ^ (32 - bits)))) == (base_num - (base_num % (2 ^ (32 - bits))))
end
local function ip_matches(items, ip)
if not items or type(items) ~= "table" or not ip or ip == "" then
return false
end
for _, item in ipairs(items) do
if item == ip then
return true
end
if string.find(item, "/", 1, true) and ipv4_in_cidr(ip, item) then
return true
end
end
return false
end
local function ip_matches_group_ids(group_ids, ip, ip_groups_config)
if not group_ids or type(group_ids) ~= "table" or not ip or ip == "" then
return false
end
local groups = (ip_groups_config or {}).groups or {}
for _, id in ipairs(group_ids) do
local group = groups[tostring(id)]
if group and group.enabled and ip_matches(group.ip_list, ip) then
return true
end
end
return false
end
local function lookup_country(ip)
local ok, maxminddb = pcall(require, "resty.maxminddb")
if not ok or not maxminddb then
return nil
end
local paths = {
"__OPENFLARE_RUNTIME_CONFIG_DIR__/GeoLite2-Country.mmdb",
"/etc/openflare/GeoLite2-Country.mmdb",
"/usr/local/share/openflare/GeoLite2-Country.mmdb"
}
for _, path in ipairs(paths) do
local opened = pcall(maxminddb.init, path)
if opened then
local res, err = maxminddb.lookup(ip)
if res and res.country and res.country.iso_code then
return string.upper(res.country.iso_code)
end
end
end
return nil
end
local function group_by_id(config)
local result = {}
for _, group in ipairs(config.rule_groups or {}) do
result[tostring(group.id)] = group
end
return result
end
local function active_groups(config, groups)
local site = ngx.var.openflare_waf_site or ""
local ids = (config.site_rule_groups or {})[site]
local result = {}
for _, group in ipairs(config.rule_groups or {}) do
if group.is_global then
result[#result + 1] = group
end
end
if ids then
local by_id = group_by_id(config)
for _, id in ipairs(ids) do
local group = by_id[tostring(id)]
if group and not group.is_global then
result[#result + 1] = group
end
end
end
return result
end
local function exit_with_group(group)
ngx.ctx.openflare_waf_blocked = true
ngx.status = tonumber(group.block_status_code) or 418
local body = group.block_response_body or ""
if body ~= "" then
ngx.header["Content-Type"] = "text/html; charset=utf-8"
ngx.say(body)
end
return ngx.exit(ngx.status)
end
local function first_allowlist_group(groups)
for _, group in ipairs(groups) do
if table_has_items(group.ip_whitelist)
or table_has_items(group.ip_whitelist_group_ids)
or table_has_items(group.country_whitelist) then
return group
end
end
return nil
end
local config = load_config()
if not config then
if config_dict:add("_missing_config_logged", true, 60) then
ngx.log(ngx.WARN, "openflare waf config is missing or invalid; requests will be allowed")
end
return
end
local ip = ngx.var.remote_addr or ""
local groups = active_groups(config)
local ip_groups_config = load_ip_groups()
if #groups == 0 then
if config_dict:add("_empty_groups_logged", true, 60) then
ngx.log(ngx.WARN, "openflare waf has no active rule group for site: ", ngx.var.openflare_waf_site or "")
end
return
end
for _, group in ipairs(groups) do
if ip_matches(group.ip_whitelist, ip) or ip_matches_group_ids(group.ip_whitelist_group_ids, ip, ip_groups_config) then
return
end
end
local country = nil
for _, group in ipairs(groups) do
if type(group.country_whitelist) == "table" and #group.country_whitelist > 0 then
country = country or lookup_country(ip)
if list_contains(group.country_whitelist, country) then
return
end
end
end
local allowlist_group = first_allowlist_group(groups)
if allowlist_group then
return exit_with_group(allowlist_group)
end
for _, group in ipairs(groups) do
if ip_matches(group.ip_blacklist, ip) or ip_matches_group_ids(group.ip_blacklist_group_ids, ip, ip_groups_config) then
return exit_with_group(group)
end
end
for _, group in ipairs(groups) do
if type(group.country_blacklist) == "table" and #group.country_blacklist > 0 then
country = country or lookup_country(ip)
if list_contains(group.country_blacklist, country) then
return exit_with_group(group)
end
end
end
return "ok"
end
return _M
`
const openRestyWAFCheckLua = `local source = debug.getinfo(1, "S").source or ""
if string.sub(source, 1, 1) == "@" then
local script_path = string.sub(source, 2)
local base_dir = string.match(script_path, "^(.*)/waf/[^/]+%.lua$")
if base_dir and base_dir ~= "" and not string.find(package.path, base_dir, 1, true) then
package.path = base_dir .. "/?.lua;" .. base_dir .. "/?/init.lua;" .. package.path
end
end
return require("waf.runtime").check()
`
func ManagedWAFLuaFiles() []protocol.SupportFile {
return []protocol.SupportFile{
{Path: "waf/runtime.lua", Content: openRestyWAFRuntimeLua},
{Path: "waf/check.lua", Content: openRestyWAFCheckLua},
}
}
@@ -0,0 +1,404 @@
package observability
import (
"bufio"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"os"
"path/filepath"
"runtime"
"strconv"
"strings"
"syscall"
"time"
"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 BuildProfile(cfg *config.Config, stateStore *state.Store) *protocol.NodeSystemProfile {
profile := collectProfile(cfg)
if profile == nil {
return nil
}
fingerprint := fingerprintProfile(profile)
if stateStore == nil {
return profile
}
snapshot, err := stateStore.Load()
if err != nil {
return profile
}
if snapshot.LastProfileFingerprint == fingerprint {
return nil
}
snapshot.LastProfileFingerprint = fingerprint
if err = stateStore.Save(snapshot); err != nil {
return profile
}
return profile
}
func BuildSnapshot(cfg *config.Config, stateStore *state.Store) *protocol.NodeMetricSnapshot {
now := time.Now().UTC()
metric := &protocol.NodeMetricSnapshot{
CapturedAtUnix: now.Unix(),
}
memTotal, memUsed := readMemInfo()
metric.MemoryTotalBytes = memTotal
metric.MemoryUsedBytes = memUsed
storageTotal, storageUsed := statFilesystem(cfg.DataDir)
metric.StorageTotalBytes = storageTotal
metric.StorageUsedBytes = storageUsed
metric.NetworkRxBytes, metric.NetworkTxBytes = readLinuxNetworkTotals()
metric.DiskReadBytes, metric.DiskWriteBytes = readLinuxDiskTotals()
if stateStore == nil {
return metric
}
totalCPU, idleCPU := readLinuxCPUStat()
snapshot, err := stateStore.Load()
if err != nil {
return metric
}
if snapshot.LastCPUStatTotal > 0 && totalCPU > snapshot.LastCPUStatTotal && idleCPU >= snapshot.LastCPUStatIdle {
deltaTotal := totalCPU - snapshot.LastCPUStatTotal
deltaIdle := idleCPU - snapshot.LastCPUStatIdle
if deltaTotal > 0 && deltaIdle <= deltaTotal {
metric.CPUUsagePercent = (float64(deltaTotal-deltaIdle) / float64(deltaTotal)) * 100
}
}
snapshot.LastCPUStatTotal = totalCPU
snapshot.LastCPUStatIdle = idleCPU
snapshot.LastMetricAtUnix = now.Unix()
_ = stateStore.Save(snapshot)
return metric
}
func BuildOpenrestyObservation(managed *ManagedOpenRestyMetrics) *protocol.NodeOpenrestyObservation {
if managed == nil {
return nil
}
return &protocol.NodeOpenrestyObservation{
CapturedAtUnix: time.Now().UTC().Unix(),
OpenrestyRxBytes: managed.OpenrestyRxBytes,
OpenrestyTxBytes: managed.OpenrestyTxBytes,
OpenrestyConnections: managed.OpenrestyConnections,
}
}
func BuildHealthEvents(snapshot *state.Snapshot) []protocol.NodeHealthEvent {
if snapshot == nil {
return []protocol.NodeHealthEvent{}
}
events := make([]protocol.NodeHealthEvent, 0, 2)
nowUnix := time.Now().UTC().Unix()
if strings.TrimSpace(snapshot.OpenrestyStatus) == protocol.OpenrestyStatusUnhealthy {
events = append(events, protocol.NodeHealthEvent{
EventType: "openresty_unhealthy",
Severity: "critical",
Message: strings.TrimSpace(snapshot.OpenrestyMessage),
TriggeredAtUnix: nowUnix,
})
}
if strings.TrimSpace(snapshot.LastError) != "" {
events = append(events, protocol.NodeHealthEvent{
EventType: "sync_error",
Severity: "warning",
Message: strings.TrimSpace(snapshot.LastError),
TriggeredAtUnix: nowUnix,
})
}
return events
}
func collectProfile(cfg *config.Config) *protocol.NodeSystemProfile {
hostname, _ := os.Hostname()
osName, osVersion := readLinuxOSRelease()
kernelVersion := readFirstLine("/proc/sys/kernel/osrelease")
cpuModel := readLinuxCPUModel()
totalMemory, _ := readMemInfo()
totalDisk, _ := statFilesystem(cfg.DataDir)
uptimeSeconds := readLinuxUptimeSeconds()
return &protocol.NodeSystemProfile{
Hostname: strings.TrimSpace(hostname),
OSName: osName,
OSVersion: osVersion,
KernelVersion: kernelVersion,
Architecture: runtime.GOARCH,
CPUModel: cpuModel,
CPUCores: runtime.NumCPU(),
TotalMemoryBytes: totalMemory,
TotalDiskBytes: totalDisk,
UptimeSeconds: uptimeSeconds,
ReportedAtUnix: time.Now().UTC().Unix(),
}
}
func fingerprintProfile(profile *protocol.NodeSystemProfile) string {
raw, err := json.Marshal(profile)
if err != nil {
return ""
}
sum := sha256.Sum256(raw)
return hex.EncodeToString(sum[:])
}
func readLinuxOSRelease() (string, string) {
file, err := os.Open("/etc/os-release")
if err != nil {
return runtime.GOOS, ""
}
defer file.Close()
values := make(map[string]string)
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" || strings.HasPrefix(line, "#") {
continue
}
key, value, ok := strings.Cut(line, "=")
if !ok {
continue
}
values[key] = strings.Trim(value, `"`)
}
if pretty := strings.TrimSpace(values["PRETTY_NAME"]); pretty != "" {
return pretty, strings.TrimSpace(values["VERSION_ID"])
}
name := strings.TrimSpace(values["NAME"])
if name == "" {
name = runtime.GOOS
}
return name, strings.TrimSpace(values["VERSION_ID"])
}
func readLinuxCPUModel() string {
file, err := os.Open("/proc/cpuinfo")
if err != nil {
return ""
}
defer file.Close()
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(strings.ToLower(line), "model name") {
_, value, ok := strings.Cut(line, ":")
if ok {
return strings.TrimSpace(value)
}
}
}
return ""
}
func readMemInfo() (int64, int64) {
file, err := os.Open("/proc/meminfo")
if err != nil {
return 0, 0
}
defer file.Close()
var memTotalKB int64
var memAvailableKB int64
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := scanner.Text()
switch {
case strings.HasPrefix(line, "MemTotal:"):
memTotalKB = parseMemInfoValue(line)
case strings.HasPrefix(line, "MemAvailable:"):
memAvailableKB = parseMemInfoValue(line)
}
}
total := memTotalKB * 1024
if total == 0 {
return 0, 0
}
used := total - (memAvailableKB * 1024)
if used < 0 {
used = 0
}
return total, used
}
func parseMemInfoValue(line string) int64 {
fields := strings.Fields(line)
if len(fields) < 2 {
return 0
}
value, err := strconv.ParseInt(fields[1], 10, 64)
if err != nil {
return 0
}
return value
}
func readLinuxUptimeSeconds() int64 {
content, err := os.ReadFile("/proc/uptime")
if err != nil {
return 0
}
fields := strings.Fields(string(content))
if len(fields) == 0 {
return 0
}
value, err := strconv.ParseFloat(fields[0], 64)
if err != nil {
return 0
}
return int64(value)
}
func readLinuxCPUStat() (uint64, uint64) {
content, err := os.ReadFile("/proc/stat")
if err != nil {
return 0, 0
}
lines := strings.Split(string(content), "\n")
for _, line := range lines {
if !strings.HasPrefix(line, "cpu ") {
continue
}
fields := strings.Fields(line)
if len(fields) < 5 {
return 0, 0
}
var total uint64
for i := 1; i < len(fields); i++ {
value, err := strconv.ParseUint(fields[i], 10, 64)
if err != nil {
return 0, 0
}
total += value
if i == 4 {
// idle
}
}
idle, err := strconv.ParseUint(fields[4], 10, 64)
if err != nil {
return 0, 0
}
return total, idle
}
return 0, 0
}
func readLinuxNetworkTotals() (int64, int64) {
file, err := os.Open("/proc/net/dev")
if err != nil {
return 0, 0
}
defer file.Close()
var rx int64
var tx int64
scanner := bufio.NewScanner(file)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if !strings.Contains(line, ":") {
continue
}
name, data, ok := strings.Cut(line, ":")
if !ok {
continue
}
if strings.TrimSpace(name) == "lo" {
continue
}
fields := strings.Fields(data)
if len(fields) < 16 {
continue
}
rxValue, err := strconv.ParseInt(fields[0], 10, 64)
if err == nil {
rx += rxValue
}
txValue, err := strconv.ParseInt(fields[8], 10, 64)
if err == nil {
tx += txValue
}
}
return rx, tx
}
func readLinuxDiskTotals() (int64, int64) {
file, err := os.Open("/proc/diskstats")
if err != nil {
return 0, 0
}
defer file.Close()
var readBytes int64
var writeBytes int64
scanner := bufio.NewScanner(file)
for scanner.Scan() {
fields := strings.Fields(scanner.Text())
if len(fields) < 14 {
continue
}
device := fields[2]
if shouldSkipDiskDevice(device) {
continue
}
readSectors, err := strconv.ParseInt(fields[5], 10, 64)
if err == nil {
readBytes += readSectors * 512
}
writeSectors, err := strconv.ParseInt(fields[9], 10, 64)
if err == nil {
writeBytes += writeSectors * 512
}
}
return readBytes, writeBytes
}
func shouldSkipDiskDevice(device string) bool {
switch {
case device == "":
return true
case strings.HasPrefix(device, "loop"),
strings.HasPrefix(device, "ram"),
strings.HasPrefix(device, "dm-"):
return true
default:
return false
}
}
func statFilesystem(path string) (int64, int64) {
if strings.TrimSpace(path) == "" {
path = string(os.PathSeparator)
}
absPath := filepath.Clean(path)
var stat syscall.Statfs_t
if err := syscall.Statfs(absPath, &stat); err != nil {
return 0, 0
}
total := int64(stat.Blocks) * int64(stat.Bsize)
free := int64(stat.Bavail) * int64(stat.Bsize)
used := total - free
if used < 0 {
used = 0
}
return total, used
}
func readFirstLine(path string) string {
content, err := os.ReadFile(path)
if err != nil {
return ""
}
return strings.TrimSpace(string(content))
}
@@ -0,0 +1,130 @@
package observability
import (
"encoding/json"
"fmt"
"io"
"net/http"
"regexp"
"strconv"
"strings"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
const openRestyObservabilityPath = "/openflare/observability"
const openRestyStubStatusPath = "/openflare/stub_status"
var stubStatusActivePattern = regexp.MustCompile(`Active connections:\s+(\d+)`)
type ManagedOpenRestyMetrics struct {
TrafficReport *protocol.NodeTrafficReport
OpenrestyRxBytes int64
OpenrestyTxBytes int64
OpenrestyConnections int64
}
type openRestyObservabilityResponse struct {
WindowStartedAtUnix int64 `json:"window_started_at_unix"`
WindowEndedAtUnix int64 `json:"window_ended_at_unix"`
RequestCount int64 `json:"request_count"`
ErrorCount int64 `json:"error_count"`
UniqueVisitorCount int64 `json:"unique_visitor_count"`
StatusCodes map[string]int64 `json:"status_codes"`
TopDomains map[string]int64 `json:"top_domains"`
SourceCountries map[string]int64 `json:"source_countries"`
OpenrestyRxBytes int64 `json:"openresty_rx_bytes"`
OpenrestyTxBytes int64 `json:"openresty_tx_bytes"`
}
func CollectManagedOpenRestyMetrics(cfg *config.Config) *ManagedOpenRestyMetrics {
if cfg == nil || cfg.OpenrestyObservabilityPort <= 0 {
return nil
}
baseURL := fmt.Sprintf("http://127.0.0.1:%d", cfg.OpenrestyObservabilityPort)
client := &http.Client{Timeout: 1500 * time.Millisecond}
observabilityResp := openRestyObservabilityResponse{}
if err := fetchLocalJSON(client, baseURL+openRestyObservabilityPath, &observabilityResp); err != nil {
return nil
}
result := &ManagedOpenRestyMetrics{
TrafficReport: &protocol.NodeTrafficReport{
WindowStartedAtUnix: observabilityResp.WindowStartedAtUnix,
WindowEndedAtUnix: observabilityResp.WindowEndedAtUnix,
RequestCount: observabilityResp.RequestCount,
ErrorCount: observabilityResp.ErrorCount,
UniqueVisitorCount: observabilityResp.UniqueVisitorCount,
StatusCodes: normalizeCountMap(observabilityResp.StatusCodes),
TopDomains: normalizeCountMap(observabilityResp.TopDomains),
SourceCountries: normalizeCountMap(observabilityResp.SourceCountries),
},
OpenrestyRxBytes: observabilityResp.OpenrestyRxBytes,
OpenrestyTxBytes: observabilityResp.OpenrestyTxBytes,
}
if text, err := fetchLocalText(client, baseURL+openRestyStubStatusPath); err == nil {
result.OpenrestyConnections = parseStubStatusActiveConnections(text)
}
return result
}
func fetchLocalJSON(client *http.Client, url string, target any) error {
resp, err := client.Get(url)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("unexpected local observability status: %s", resp.Status)
}
return json.NewDecoder(resp.Body).Decode(target)
}
func fetchLocalText(client *http.Client, url string) (string, error) {
resp, err := client.Get(url)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("unexpected local stub status: %s", resp.Status)
}
data, err := io.ReadAll(resp.Body)
if err != nil {
return "", err
}
return string(data), nil
}
func parseStubStatusActiveConnections(raw string) int64 {
matches := stubStatusActivePattern.FindStringSubmatch(raw)
if len(matches) != 2 {
return 0
}
value, err := strconv.ParseInt(matches[1], 10, 64)
if err != nil {
return 0
}
return value
}
func normalizeCountMap(values map[string]int64) map[string]int64 {
if len(values) == 0 {
return map[string]int64{}
}
result := make(map[string]int64, len(values))
for key, value := range values {
key = strings.TrimSpace(key)
if key == "" || value <= 0 {
continue
}
result[key] = value
}
return result
}
@@ -0,0 +1,82 @@
package observability
import (
"net"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
)
func TestCollectManagedOpenRestyMetrics(t *testing.T) {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("Listen failed: %v", err)
}
port := listener.Addr().(*net.TCPAddr).Port
mux := http.NewServeMux()
mux.HandleFunc(openRestyObservabilityPath, func(writer http.ResponseWriter, request *http.Request) {
writer.Header().Set("Content-Type", "application/json")
_, _ = writer.Write([]byte(`{"window_started_at_unix":1710403200,"window_ended_at_unix":1710403210,"request_count":12,"error_count":2,"unique_visitor_count":5,"status_codes":{"200":10,"502":2},"top_domains":{"app.example.com":9,"api.example.com":3},"source_countries":{},"openresty_rx_bytes":4096,"openresty_tx_bytes":8192}`))
})
mux.HandleFunc(openRestyStubStatusPath, func(writer http.ResponseWriter, request *http.Request) {
_, _ = writer.Write([]byte("Active connections: 7 \nserver accepts handled requests\n 10 10 12 \nReading: 1 Writing: 2 Waiting: 4 \n"))
})
server := httptest.NewUnstartedServer(mux)
server.Listener = listener
server.Start()
defer server.Close()
metrics := CollectManagedOpenRestyMetrics(&config.Config{
OpenrestyObservabilityPort: port,
})
if metrics == nil || metrics.TrafficReport == nil {
t.Fatalf("expected managed openresty metrics, got %+v", metrics)
}
if metrics.TrafficReport.RequestCount != 12 || metrics.TrafficReport.ErrorCount != 2 {
t.Fatalf("unexpected traffic report: %+v", metrics.TrafficReport)
}
if metrics.OpenrestyRxBytes != 4096 || metrics.OpenrestyTxBytes != 8192 {
t.Fatalf("unexpected openresty byte counters: %+v", metrics)
}
if metrics.OpenrestyConnections != 7 {
t.Fatalf("unexpected openresty connections: %+v", metrics)
}
}
func TestParseStubStatusActiveConnections(t *testing.T) {
if value := parseStubStatusActiveConnections("Active connections: 19\n"); value != 19 {
t.Fatalf("unexpected active connections: %d", value)
}
}
func TestNormalizeCountMapDropsEmptyKeys(t *testing.T) {
normalized := normalizeCountMap(map[string]int64{
"": 4,
" 200 ": 3,
"app.example.com": 0,
})
if len(normalized) != 1 || normalized["200"] != 3 {
t.Fatalf("unexpected normalized map: %+v", normalized)
}
}
func TestCollectManagedOpenRestyMetricsHandlesUnavailableEndpoint(t *testing.T) {
cfg := &config.Config{OpenrestyObservabilityPort: 1}
if metrics := CollectManagedOpenRestyMetrics(cfg); metrics != nil {
t.Fatalf("expected nil metrics for unavailable endpoint, got %+v", metrics)
}
}
func TestOpenRestyObservabilityPathsAreStable(t *testing.T) {
if !strings.HasPrefix(openRestyObservabilityPath, "/openflare/") {
t.Fatalf("unexpected observability path: %s", openRestyObservabilityPath)
}
if !strings.HasPrefix(openRestyStubStatusPath, "/openflare/") {
t.Fatalf("unexpected stub status path: %s", openRestyStubStatusPath)
}
}
@@ -0,0 +1,362 @@
package observability
import (
"bufio"
"encoding/json"
"errors"
"io"
"log/slog"
"os"
"regexp"
"sort"
"strconv"
"strings"
"time"
"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"
)
type accessLogRecord struct {
Timestamp string `json:"ts"`
Host string `json:"host"`
RemoteAddr string `json:"remote_addr"`
Path string `json:"path"`
Status int `json:"status"`
BytesSent int64 `json:"bytes_sent"`
RequestLength int64 `json:"request_length"`
}
var combinedAccessLogPattern = regexp.MustCompile(`^(\S+)\s+\S+\s+\S+\s+\[([^]]+)]\s+"\S+\s+(\S+)(?:\s+[^"]*)?"\s+(\d{3})\s+\S+`)
type trafficAggregate struct {
windowStartedAt time.Time
windowEndedAt time.Time
requestCount int64
errorCount int64
openrestyRxBytes int64
openrestyTxBytes int64
statusCodes map[string]int64
topDomains map[string]int64
visitors map[string]struct{}
logs []protocol.NodeAccessLog
}
func BuildTrafficReport(cfg *config.Config, stateStore *state.Store, managed *ManagedOpenRestyMetrics) *protocol.NodeTrafficReport {
report, _, _ := BuildTrafficObservability(cfg, stateStore, managed)
return report
}
func BuildTrafficObservability(cfg *config.Config, stateStore *state.Store, managed *ManagedOpenRestyMetrics) (*protocol.NodeTrafficReport, []protocol.NodeAccessLog, *ManagedOpenRestyMetrics) {
if cfg == nil || stateStore == nil {
if managed != nil && managed.TrafficReport != nil {
return managed.TrafficReport, nil, managed
}
return nil, nil, managed
}
aggregate := readAccessLogDelta(cfg, stateStore)
var accessLogs []protocol.NodeAccessLog
if aggregate != nil {
accessLogs = aggregate.accessLogs()
}
if managed != nil && managed.TrafficReport != nil {
return managed.TrafficReport, accessLogs, managed
}
if aggregate == nil {
return nil, accessLogs, managed
}
fallbackManaged := aggregate.managedMetrics()
return aggregate.report(), accessLogs, fallbackManaged
}
func readAccessLogDelta(cfg *config.Config, stateStore *state.Store) *trafficAggregate {
snapshot, err := stateStore.Load()
if err != nil {
return nil
}
logPath := managedAccessLogPath(cfg)
file, err := os.Open(logPath)
if err != nil {
if os.IsNotExist(err) {
if snapshot.AccessLogOffset != 0 {
snapshot.AccessLogOffset = 0
_ = stateStore.Save(snapshot)
}
return nil
}
return nil
}
defer func(file *os.File) {
err := file.Close()
if err != nil {
slog.Error("failed to close access log file", "error", err)
}
}(file)
info, err := file.Stat()
if err != nil {
return nil
}
offset := snapshot.AccessLogOffset
if offset < 0 || offset > info.Size() {
offset = 0
}
if _, err = file.Seek(offset, io.SeekStart); err != nil {
return nil
}
reader := bufio.NewReader(file)
currentOffset := offset
aggregate := newTrafficAggregate()
for {
line, readErr := reader.ReadBytes('\n')
if len(line) > 0 {
currentOffset += int64(len(line))
aggregate.consume(line)
}
if errors.Is(readErr, io.EOF) {
break
}
if readErr != nil {
return nil
}
}
snapshot.AccessLogOffset = currentOffset
_ = stateStore.Save(snapshot)
return aggregate
}
func managedAccessLogPath(cfg *config.Config) string {
if cfg == nil || strings.TrimSpace(cfg.AccessLogPath) == "" {
return ""
}
return cfg.AccessLogPath
}
func newTrafficAggregate() *trafficAggregate {
return &trafficAggregate{
statusCodes: make(map[string]int64),
topDomains: make(map[string]int64),
visitors: make(map[string]struct{}),
}
}
func (aggregate *trafficAggregate) consume(line []byte) {
trimmed := strings.TrimSpace(string(line))
if trimmed == "" {
return
}
record, ok := parseAccessLogRecord(trimmed)
if !ok {
return
}
if aggregate.windowStartedAt.IsZero() || record.Timestamp.Before(aggregate.windowStartedAt) {
aggregate.windowStartedAt = record.Timestamp
}
if aggregate.windowEndedAt.IsZero() || record.Timestamp.After(aggregate.windowEndedAt) {
aggregate.windowEndedAt = record.Timestamp
}
aggregate.requestCount++
if record.Status >= 500 {
aggregate.errorCount++
}
if record.Status > 0 {
aggregate.statusCodes[strconv.Itoa(record.Status)]++
}
if record.RequestLength > 0 {
aggregate.openrestyRxBytes += record.RequestLength
}
if record.BytesSent > 0 {
aggregate.openrestyTxBytes += record.BytesSent
}
if host := strings.TrimSpace(record.Host); host != "" {
aggregate.topDomains[host]++
}
if remoteAddr := strings.TrimSpace(record.RemoteAddr); remoteAddr != "" {
aggregate.visitors[remoteAddr] = struct{}{}
}
aggregate.logs = append(aggregate.logs, protocol.NodeAccessLog{
LoggedAtUnix: record.Timestamp.Unix(),
RemoteAddr: strings.TrimSpace(record.RemoteAddr),
Host: strings.TrimSpace(record.Host),
Path: normalizeAccessLogPath(record.Path),
StatusCode: record.Status,
})
}
type parsedAccessLogRecord struct {
Timestamp time.Time
Host string
RemoteAddr string
Path string
Status int
BytesSent int64
RequestLength int64
}
func parseAccessLogRecord(raw string) (parsedAccessLogRecord, bool) {
record, ok := parseJSONAccessLogRecord(raw)
if ok {
return record, true
}
return parseCombinedAccessLogRecord(raw)
}
func parseJSONAccessLogRecord(raw string) (parsedAccessLogRecord, bool) {
var record accessLogRecord
if err := json.Unmarshal([]byte(raw), &record); err != nil {
return parsedAccessLogRecord{}, false
}
timestamp, err := parseAccessLogTime(record.Timestamp)
if err != nil {
return parsedAccessLogRecord{}, false
}
return parsedAccessLogRecord{
Timestamp: timestamp,
Host: strings.TrimSpace(record.Host),
RemoteAddr: strings.TrimSpace(record.RemoteAddr),
Path: normalizeAccessLogPath(record.Path),
Status: record.Status,
BytesSent: record.BytesSent,
RequestLength: record.RequestLength,
}, true
}
func parseCombinedAccessLogRecord(raw string) (parsedAccessLogRecord, bool) {
matches := combinedAccessLogPattern.FindStringSubmatch(raw)
if len(matches) != 5 {
return parsedAccessLogRecord{}, false
}
timestamp, err := parseAccessLogTime(matches[2])
if err != nil {
return parsedAccessLogRecord{}, false
}
status, err := strconv.Atoi(matches[4])
if err != nil {
return parsedAccessLogRecord{}, false
}
return parsedAccessLogRecord{
Timestamp: timestamp,
RemoteAddr: strings.TrimSpace(matches[1]),
Path: normalizeAccessLogPath(matches[3]),
Status: status,
}, true
}
func (aggregate *trafficAggregate) report() *protocol.NodeTrafficReport {
if aggregate.requestCount == 0 || aggregate.windowStartedAt.IsZero() || aggregate.windowEndedAt.IsZero() {
return nil
}
return &protocol.NodeTrafficReport{
WindowStartedAtUnix: aggregate.windowStartedAt.Unix(),
WindowEndedAtUnix: aggregate.windowEndedAt.Unix(),
RequestCount: aggregate.requestCount,
ErrorCount: aggregate.errorCount,
UniqueVisitorCount: int64(len(aggregate.visitors)),
StatusCodes: cloneTrafficCounts(aggregate.statusCodes, 0),
TopDomains: topCounts(aggregate.topDomains, 8),
SourceCountries: map[string]int64{},
}
}
func (aggregate *trafficAggregate) accessLogs() []protocol.NodeAccessLog {
if aggregate == nil || len(aggregate.logs) == 0 {
return []protocol.NodeAccessLog{}
}
return append([]protocol.NodeAccessLog(nil), aggregate.logs...)
}
func (aggregate *trafficAggregate) managedMetrics() *ManagedOpenRestyMetrics {
if aggregate == nil {
return nil
}
report := aggregate.report()
if report == nil && aggregate.openrestyRxBytes <= 0 && aggregate.openrestyTxBytes <= 0 {
return nil
}
return &ManagedOpenRestyMetrics{
TrafficReport: report,
OpenrestyRxBytes: aggregate.openrestyRxBytes,
OpenrestyTxBytes: aggregate.openrestyTxBytes,
}
}
func parseAccessLogTime(value string) (time.Time, error) {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
return time.Time{}, errors.New("empty access log time")
}
timestamp, err := time.Parse(time.RFC3339, trimmed)
if err == nil {
return timestamp, nil
}
return time.Parse("02/Jan/2006:15:04:05 -0700", trimmed)
}
func cloneTrafficCounts(values map[string]int64, limit int) map[string]int64 {
if len(values) == 0 {
return map[string]int64{}
}
items := make([]trafficCountItem, 0, len(values))
for key, value := range values {
items = append(items, trafficCountItem{key: key, value: value})
}
sort.Slice(items, func(i int, j int) bool {
if items[i].value == items[j].value {
return items[i].key < items[j].key
}
return items[i].value > items[j].value
})
if limit > 0 && len(items) > limit {
items = items[:limit]
}
result := make(map[string]int64, len(items))
for _, item := range items {
result[item.key] = item.value
}
return result
}
type trafficCountItem struct {
key string
value int64
}
const accessLogPathMaxRunes = 100
func normalizeAccessLogPath(value string) string {
trimmed := strings.TrimSpace(value)
if trimmed == "" {
return ""
}
if strings.HasPrefix(trimmed, "http://") || strings.HasPrefix(trimmed, "https://") {
return truncateAccessLogPath(trimmed)
}
if strings.HasPrefix(trimmed, "/") {
return truncateAccessLogPath(trimmed)
}
return truncateAccessLogPath("/" + trimmed)
}
func truncateAccessLogPath(value string) string {
runes := []rune(value)
if len(runes) <= accessLogPathMaxRunes {
return value
}
return string(runes[:accessLogPathMaxRunes])
}
func topCounts(values map[string]int64, limit int) map[string]int64 {
return cloneTrafficCounts(values, limit)
}
@@ -0,0 +1,188 @@
package observability
import (
"os"
"path/filepath"
"strings"
"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 TestBuildTrafficReportAggregatesManagedAccessLog(t *testing.T) {
tempDir := t.TempDir()
routeConfigPath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
if err := os.MkdirAll(filepath.Dir(routeConfigPath), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
logPath := filepath.Join(filepath.Dir(routeConfigPath), "openflare_access.log")
content := []byte(
"{\"ts\":\"2026-03-14T08:00:00Z\",\"host\":\"app.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.1\",\"status\":200}\n" +
"{\"ts\":\"2026-03-14T08:00:05Z\",\"host\":\"app.example.com\",\"path\":\"/healthz\",\"remote_addr\":\"10.0.0.2\",\"status\":503}\n" +
"{\"ts\":\"2026-03-14T08:00:08Z\",\"host\":\"api.example.com\",\"path\":\"/api\",\"remote_addr\":\"10.0.0.1\",\"status\":200}\n",
)
if err := os.WriteFile(logPath, content, 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
report := BuildTrafficReport(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if report == nil {
t.Fatal("expected traffic report")
}
if report.RequestCount != 3 || report.ErrorCount != 1 || report.UniqueVisitorCount != 2 {
t.Fatalf("unexpected traffic report counters: %+v", report)
}
if report.StatusCodes["200"] != 2 || report.StatusCodes["503"] != 1 {
t.Fatalf("unexpected status codes: %+v", report.StatusCodes)
}
if report.TopDomains["app.example.com"] != 2 || report.TopDomains["api.example.com"] != 1 {
t.Fatalf("unexpected top domains: %+v", report.TopDomains)
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if snapshot.AccessLogOffset != int64(len(content)) {
t.Fatalf("unexpected access log offset: %d", snapshot.AccessLogOffset)
}
secondReport := BuildTrafficReport(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if secondReport != nil {
t.Fatalf("expected no report without appended lines, got %+v", secondReport)
}
}
func TestBuildTrafficReportResetsOffsetAfterTruncate(t *testing.T) {
tempDir := t.TempDir()
routeConfigPath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
if err := os.MkdirAll(filepath.Dir(routeConfigPath), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
logPath := filepath.Join(filepath.Dir(routeConfigPath), "openflare_access.log")
if err := os.WriteFile(logPath, []byte("{\"ts\":\"2026-03-14T09:00:00Z\",\"host\":\"app.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.3\",\"status\":200}\n"), 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
if err := stateStore.Save(&state.Snapshot{AccessLogOffset: 4096}); err != nil {
t.Fatalf("Save failed: %v", err)
}
report := BuildTrafficReport(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if report == nil || report.RequestCount != 1 {
t.Fatalf("expected one request after truncate reset, got %+v", report)
}
}
func TestBuildTrafficObservabilityReturnsAccessLogs(t *testing.T) {
tempDir := t.TempDir()
routeConfigPath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
if err := os.MkdirAll(filepath.Dir(routeConfigPath), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
logPath := filepath.Join(filepath.Dir(routeConfigPath), "openflare_access.log")
content := []byte(
"{\"ts\":\"2026-03-14T08:00:00Z\",\"host\":\"app.example.com\",\"path\":\"/login\",\"remote_addr\":\"10.0.0.1\",\"status\":200,\"request_length\":128,\"bytes_sent\":512}\n" +
"{\"ts\":\"2026-03-14T08:00:05Z\",\"host\":\"api.example.com\",\"path\":\"/v1/ping\",\"remote_addr\":\"10.0.0.2\",\"status\":502,\"request_length\":64,\"bytes_sent\":256}\n",
)
if err := os.WriteFile(logPath, content, 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
report, accessLogs, fallbackMetrics := BuildTrafficObservability(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if report == nil || report.RequestCount != 2 {
t.Fatalf("expected traffic report, got %+v", report)
}
if len(accessLogs) != 2 {
t.Fatalf("expected access logs, got %+v", accessLogs)
}
if fallbackMetrics == nil || fallbackMetrics.OpenrestyRxBytes != 192 || fallbackMetrics.OpenrestyTxBytes != 768 {
t.Fatalf("expected fallback throughput metrics, got %+v", fallbackMetrics)
}
if accessLogs[0].Path != "/login" || accessLogs[1].Path != "/v1/ping" {
t.Fatalf("unexpected access log paths: %+v", accessLogs)
}
}
func TestBuildTrafficObservabilityTruncatesLongAccessLogPath(t *testing.T) {
tempDir := t.TempDir()
routeConfigPath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
if err := os.MkdirAll(filepath.Dir(routeConfigPath), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
logPath := filepath.Join(filepath.Dir(routeConfigPath), "openflare_access.log")
longPath := "/" + strings.Repeat("a", 140)
content := []byte(
"{\"ts\":\"2026-03-14T08:00:00Z\",\"host\":\"app.example.com\",\"path\":\"" + longPath + "\",\"remote_addr\":\"10.0.0.1\",\"status\":200}\n",
)
if err := os.WriteFile(logPath, content, 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
_, accessLogs, _ := BuildTrafficObservability(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if len(accessLogs) != 1 {
t.Fatalf("expected one access log, got %+v", accessLogs)
}
if got := len([]rune(accessLogs[0].Path)); got != accessLogPathMaxRunes {
t.Fatalf("expected truncated path length %d, got %d (%q)", accessLogPathMaxRunes, got, accessLogs[0].Path)
}
}
func TestBuildTrafficReportParsesCombinedAccessLog(t *testing.T) {
tempDir := t.TempDir()
routeConfigPath := filepath.Join(tempDir, "conf.d", "openflare_routes.conf")
if err := os.MkdirAll(filepath.Dir(routeConfigPath), 0o755); err != nil {
t.Fatalf("MkdirAll failed: %v", err)
}
logPath := filepath.Join(filepath.Dir(routeConfigPath), "openflare_access.log")
content := []byte(
"10.0.0.1 - - [14/Mar/2026:08:00:00 +0000] \"GET / HTTP/1.1\" 200 123 \"-\" \"curl/8.0\"\n" +
"10.0.0.2 - - [14/Mar/2026:08:00:05 +0000] \"GET /healthz HTTP/1.1\" 502 64 \"-\" \"curl/8.0\"\n" +
"10.0.0.1 - - [14/Mar/2026:08:00:10 +0000] \"GET /api HTTP/1.1\" 200 256 \"-\" \"curl/8.0\"\n",
)
if err := os.WriteFile(logPath, content, 0o644); err != nil {
t.Fatalf("WriteFile failed: %v", err)
}
stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
report := BuildTrafficReport(&config.Config{AccessLogPath: logPath}, stateStore, nil)
if report == nil {
t.Fatal("expected traffic report from combined access log")
}
if report.RequestCount != 3 || report.ErrorCount != 1 || report.UniqueVisitorCount != 2 {
t.Fatalf("unexpected combined log counters: %+v", report)
}
if report.StatusCodes["200"] != 2 || report.StatusCodes["502"] != 1 {
t.Fatalf("unexpected combined log status codes: %+v", report.StatusCodes)
}
if len(report.TopDomains) != 0 {
t.Fatalf("expected combined access log to omit top domains when host is unavailable, got %+v", report.TopDomains)
}
}
func TestBuildTrafficReportReturnsManagedWindowEvenWhenRequestCountZero(t *testing.T) {
report := BuildTrafficReport(nil, nil, &ManagedOpenRestyMetrics{
TrafficReport: &protocol.NodeTrafficReport{
WindowStartedAtUnix: 1710403200,
WindowEndedAtUnix: 1710403260,
RequestCount: 0,
ErrorCount: 0,
UniqueVisitorCount: 0,
StatusCodes: map[string]int64{},
TopDomains: map[string]int64{},
SourceCountries: map[string]int64{},
},
})
if report == nil {
t.Fatal("expected managed traffic report to be returned even when request count is zero")
}
if report.RequestCount != 0 || report.WindowStartedAtUnix != 1710403200 || report.WindowEndedAtUnix != 1710403260 {
t.Fatalf("unexpected managed traffic report: %+v", report)
}
}
+207
View File
@@ -0,0 +1,207 @@
package protocol
import "encoding/json"
type APIResponse[T any] struct {
ErrorMsg string `json:"error_msg"`
Data T `json:"data"`
}
type HeartbeatData struct {
AgentSettings *AgentSettings `json:"agent_settings"`
ActiveConfig *ActiveConfigMeta `json:"active_config"`
WAFIPGroups []WAFIPGroup `json:"waf_ip_groups,omitempty"`
}
type HeartbeatResult struct {
AgentSettings *AgentSettings
ActiveConfig *ActiveConfigMeta
WAFIPGroups []WAFIPGroup
}
type AgentSettings struct {
HeartbeatInterval int `json:"heartbeat_interval"`
WebsocketUpgradeEnabled bool `json:"websocket_upgrade_enabled"`
AutoUpdate bool `json:"auto_update"`
UpdateRepo string `json:"update_repo"`
UpdateNow bool `json:"update_now"`
UpdateChannel string `json:"update_channel"`
UpdateTag string `json:"update_tag"`
RestartOpenrestyNow bool `json:"restart_openresty_now"`
}
const (
WSMessageTypeStatus = "status"
WSMessageTypeSettings = "settings"
WSMessageTypeActiveConfig = "active_config"
WSMessageTypeForceSyncConfig = "force_sync_config"
WSMessageTypeWAFIPGroups = "waf_ip_groups"
WSMessageTypePing = "ping"
WSMessageTypePong = "pong"
)
type WSMessage struct {
Type string `json:"type"`
Payload json.RawMessage `json:"payload,omitempty"`
}
type WSOutboundMessage struct {
Type string `json:"type"`
Payload any `json:"payload,omitempty"`
}
type WebSocketConnection interface {
URL() string
SendStatus(payload NodePayload) error
SendPong() error
Receive() (WSMessage, error)
Close() error
}
const (
OpenrestyStatusHealthy = "healthy"
OpenrestyStatusUnhealthy = "unhealthy"
OpenrestyStatusUnknown = "unknown"
)
type NodePayload struct {
NodeID string `json:"node_id"`
Name string `json:"name"`
IP string `json:"ip"`
Version string `json:"version"`
ExtVersion string `json:"ext_version"`
CurrentVersion string `json:"current_version"`
LastError string `json:"last_error"`
OpenrestyStatus string `json:"openresty_status"`
OpenrestyMessage string `json:"openresty_message"`
Profile *NodeSystemProfile `json:"profile,omitempty"`
Snapshot *NodeMetricSnapshot `json:"snapshot,omitempty"`
OpenrestyObservation *NodeOpenrestyObservation `json:"openresty_observation,omitempty"`
TrafficReport *NodeTrafficReport `json:"traffic_report,omitempty"`
AccessLogs []NodeAccessLog `json:"access_logs,omitempty"`
BufferedObservability []BufferedObservabilityRecord `json:"buffered_observability,omitempty"`
HealthEvents []NodeHealthEvent `json:"health_events"`
WAFIPGroupChecksums map[string]string `json:"waf_ip_group_checksums,omitempty"`
}
type NodeSystemProfile struct {
Hostname string `json:"hostname"`
OSName string `json:"os_name"`
OSVersion string `json:"os_version"`
KernelVersion string `json:"kernel_version"`
Architecture string `json:"architecture"`
CPUModel string `json:"cpu_model"`
CPUCores int `json:"cpu_cores"`
TotalMemoryBytes int64 `json:"total_memory_bytes"`
TotalDiskBytes int64 `json:"total_disk_bytes"`
UptimeSeconds int64 `json:"uptime_seconds"`
ReportedAtUnix int64 `json:"reported_at_unix"`
}
type NodeMetricSnapshot struct {
CapturedAtUnix int64 `json:"captured_at_unix"`
CPUUsagePercent float64 `json:"cpu_usage_percent"`
MemoryUsedBytes int64 `json:"memory_used_bytes"`
MemoryTotalBytes int64 `json:"memory_total_bytes"`
StorageUsedBytes int64 `json:"storage_used_bytes"`
StorageTotalBytes int64 `json:"storage_total_bytes"`
DiskReadBytes int64 `json:"disk_read_bytes"`
DiskWriteBytes int64 `json:"disk_write_bytes"`
NetworkRxBytes int64 `json:"network_rx_bytes"`
NetworkTxBytes int64 `json:"network_tx_bytes"`
}
type NodeOpenrestyObservation struct {
CapturedAtUnix int64 `json:"captured_at_unix"`
OpenrestyRxBytes int64 `json:"openresty_rx_bytes"`
OpenrestyTxBytes int64 `json:"openresty_tx_bytes"`
OpenrestyConnections int64 `json:"openresty_connections"`
}
type NodeTrafficReport struct {
WindowStartedAtUnix int64 `json:"window_started_at_unix"`
WindowEndedAtUnix int64 `json:"window_ended_at_unix"`
RequestCount int64 `json:"request_count"`
ErrorCount int64 `json:"error_count"`
UniqueVisitorCount int64 `json:"unique_visitor_count"`
StatusCodes map[string]int64 `json:"status_codes"`
TopDomains map[string]int64 `json:"top_domains"`
SourceCountries map[string]int64 `json:"source_countries"`
}
type NodeAccessLog struct {
LoggedAtUnix int64 `json:"logged_at_unix"`
RemoteAddr string `json:"remote_addr"`
Host string `json:"host"`
Path string `json:"path"`
StatusCode int `json:"status_code"`
}
type BufferedObservabilityRecord struct {
WindowStartedAtUnix int64 `json:"window_started_at_unix"`
Snapshot *NodeMetricSnapshot `json:"snapshot,omitempty"`
OpenrestyObservation *NodeOpenrestyObservation `json:"openresty_observation,omitempty"`
TrafficReport *NodeTrafficReport `json:"traffic_report,omitempty"`
AccessLogs []NodeAccessLog `json:"access_logs,omitempty"`
}
type NodeHealthEvent struct {
EventType string `json:"event_type"`
Severity string `json:"severity"`
Message string `json:"message"`
TriggeredAtUnix int64 `json:"triggered_at_unix"`
Metadata map[string]string `json:"metadata,omitempty"`
}
type RegisterNodeResponse struct {
NodeID string `json:"node_id"`
AccessToken string `json:"agent_token"`
Name string `json:"name"`
}
type ApplyLogPayload struct {
NodeID string `json:"node_id"`
Version string `json:"version"`
Result string `json:"result"`
Message string `json:"message"`
Checksum string `json:"checksum"`
MainConfigChecksum string `json:"main_config_checksum"`
RouteConfigChecksum string `json:"route_config_checksum"`
SupportFileCount int `json:"support_file_count"`
}
type ActiveConfigResponse struct {
Version string `json:"version"`
Checksum string `json:"checksum"`
SourceConfigJSON string `json:"source_config_json"`
SupportFiles []SupportFile `json:"support_files"`
CreatedAt string `json:"created_at"`
}
type ActiveConfigMeta struct {
Version string `json:"version"`
Checksum string `json:"checksum"`
}
type WAFIPGroup struct {
ID uint `json:"id"`
Name string `json:"name"`
Type string `json:"type"`
Enabled bool `json:"enabled"`
IPList []string `json:"ip_list"`
Checksum string `json:"checksum"`
}
type WAFIPGroupSyncRequest struct {
IDs []uint `json:"ids,omitempty"`
Checksums map[string]string `json:"checksums,omitempty"`
}
type WAFIPGroupSyncResponse struct {
Groups []WAFIPGroup `json:"groups"`
}
type SupportFile struct {
Path string `json:"path"`
Content string `json:"content"`
}
@@ -0,0 +1,229 @@
package state
import (
"encoding/json"
"os"
"path/filepath"
"sort"
"strconv"
"sync"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
const observabilityBufferWindowSeconds = 60
type ObservabilityBufferRecord struct {
WindowStartedAtUnix int64 `json:"window_started_at_unix"`
Snapshot *protocol.NodeMetricSnapshot `json:"snapshot,omitempty"`
OpenrestyObservation *protocol.NodeOpenrestyObservation `json:"openresty_observation,omitempty"`
TrafficReport *protocol.NodeTrafficReport `json:"traffic_report,omitempty"`
AccessLogs []protocol.NodeAccessLog `json:"access_logs,omitempty"`
QueuedAtUnix int64 `json:"queued_at_unix"`
}
type ObservabilityBufferStore struct {
path string
mu sync.Mutex
}
func NewObservabilityBufferStore(path string) *ObservabilityBufferStore {
return &ObservabilityBufferStore{path: filepath.Clean(path)}
}
func (s *ObservabilityBufferStore) Upsert(record ObservabilityBufferRecord, retainAfterUnix int64) error {
if s == nil || record.WindowStartedAtUnix <= 0 || (record.Snapshot == nil && record.OpenrestyObservation == nil && record.TrafficReport == nil && len(record.AccessLogs) == 0) {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
records, err := s.loadUnlocked()
if err != nil {
return err
}
records = pruneObservabilityBufferRecords(records, retainAfterUnix)
replaced := false
for index := range records {
if records[index].WindowStartedAtUnix != record.WindowStartedAtUnix {
continue
}
records[index] = mergeObservabilityBufferRecord(records[index], record)
replaced = true
break
}
if !replaced {
records = append(records, record)
}
sort.Slice(records, func(i int, j int) bool {
return records[i].WindowStartedAtUnix < records[j].WindowStartedAtUnix
})
return s.saveUnlocked(records)
}
func mergeObservabilityBufferRecord(existing ObservabilityBufferRecord, incoming ObservabilityBufferRecord) ObservabilityBufferRecord {
merged := existing
if incoming.Snapshot != nil {
merged.Snapshot = incoming.Snapshot
}
if incoming.OpenrestyObservation != nil {
merged.OpenrestyObservation = incoming.OpenrestyObservation
}
if incoming.TrafficReport != nil {
merged.TrafficReport = incoming.TrafficReport
}
merged.AccessLogs = mergeAccessLogs(existing.AccessLogs, incoming.AccessLogs)
if incoming.QueuedAtUnix > 0 {
merged.QueuedAtUnix = incoming.QueuedAtUnix
}
return merged
}
func mergeAccessLogs(existing []protocol.NodeAccessLog, incoming []protocol.NodeAccessLog) []protocol.NodeAccessLog {
if len(existing) == 0 && len(incoming) == 0 {
return nil
}
merged := make([]protocol.NodeAccessLog, 0, len(existing)+len(incoming))
seen := make(map[string]struct{}, len(existing)+len(incoming))
appendIfNeeded := func(items []protocol.NodeAccessLog) {
for _, item := range items {
key := accessLogKey(item)
if key == "" {
continue
}
if _, ok := seen[key]; ok {
continue
}
seen[key] = struct{}{}
merged = append(merged, item)
}
}
appendIfNeeded(existing)
appendIfNeeded(incoming)
sort.Slice(merged, func(i int, j int) bool {
if merged[i].LoggedAtUnix == merged[j].LoggedAtUnix {
return accessLogKey(merged[i]) < accessLogKey(merged[j])
}
return merged[i].LoggedAtUnix < merged[j].LoggedAtUnix
})
return merged
}
func accessLogKey(item protocol.NodeAccessLog) string {
return strconv.FormatInt(item.LoggedAtUnix, 10) + "|" + item.RemoteAddr + "|" + item.Host + "|" + item.Path + "|" + strconv.Itoa(item.StatusCode)
}
func (s *ObservabilityBufferStore) Replayable(currentWindowStartedAtUnix int64, retainAfterUnix int64) ([]ObservabilityBufferRecord, error) {
if s == nil {
return nil, nil
}
s.mu.Lock()
defer s.mu.Unlock()
records, err := s.loadUnlocked()
if err != nil {
return nil, err
}
records = pruneObservabilityBufferRecords(records, retainAfterUnix)
if err = s.saveUnlocked(records); err != nil {
return nil, err
}
result := make([]ObservabilityBufferRecord, 0, len(records))
for _, record := range records {
if currentWindowStartedAtUnix > 0 && record.WindowStartedAtUnix >= currentWindowStartedAtUnix {
continue
}
result = append(result, record)
}
return result, nil
}
func (s *ObservabilityBufferStore) Ack(windowStartedAtUnix []int64, retainAfterUnix int64) error {
if s == nil || len(windowStartedAtUnix) == 0 {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
records, err := s.loadUnlocked()
if err != nil {
return err
}
acked := make(map[int64]struct{}, len(windowStartedAtUnix))
for _, value := range windowStartedAtUnix {
if value > 0 {
acked[value] = struct{}{}
}
}
filtered := make([]ObservabilityBufferRecord, 0, len(records))
for _, record := range records {
if _, ok := acked[record.WindowStartedAtUnix]; ok {
continue
}
filtered = append(filtered, record)
}
filtered = pruneObservabilityBufferRecords(filtered, retainAfterUnix)
return s.saveUnlocked(filtered)
}
func (s *ObservabilityBufferStore) loadUnlocked() ([]ObservabilityBufferRecord, error) {
data, err := os.ReadFile(s.path)
if err != nil {
if os.IsNotExist(err) {
return []ObservabilityBufferRecord{}, nil
}
return nil, err
}
if len(data) == 0 {
return []ObservabilityBufferRecord{}, nil
}
var records []ObservabilityBufferRecord
if err = json.Unmarshal(data, &records); err != nil {
return nil, err
}
return records, nil
}
func (s *ObservabilityBufferStore) saveUnlocked(records []ObservabilityBufferRecord) error {
if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil {
return err
}
data, err := json.MarshalIndent(records, "", " ")
if err != nil {
return err
}
return os.WriteFile(s.path, data, 0o644)
}
func ObservabilityWindowStartedAt(snapshot *protocol.NodeMetricSnapshot, openresty *protocol.NodeOpenrestyObservation, traffic *protocol.NodeTrafficReport) int64 {
if traffic != nil && traffic.WindowStartedAtUnix > 0 {
return traffic.WindowStartedAtUnix - (traffic.WindowStartedAtUnix % observabilityBufferWindowSeconds)
}
if openresty != nil && openresty.CapturedAtUnix > 0 {
return openresty.CapturedAtUnix - (openresty.CapturedAtUnix % observabilityBufferWindowSeconds)
}
if snapshot == nil || snapshot.CapturedAtUnix <= 0 {
return 0
}
return snapshot.CapturedAtUnix - (snapshot.CapturedAtUnix % observabilityBufferWindowSeconds)
}
func pruneObservabilityBufferRecords(records []ObservabilityBufferRecord, retainAfterUnix int64) []ObservabilityBufferRecord {
if len(records) == 0 {
return []ObservabilityBufferRecord{}
}
filtered := make([]ObservabilityBufferRecord, 0, len(records))
for _, record := range records {
if record.WindowStartedAtUnix <= 0 {
continue
}
if retainAfterUnix > 0 && record.WindowStartedAtUnix < retainAfterUnix {
continue
}
filtered = append(filtered, record)
}
sort.Slice(filtered, func(i int, j int) bool {
return filtered[i].WindowStartedAtUnix < filtered[j].WindowStartedAtUnix
})
return filtered
}
@@ -0,0 +1,98 @@
package state
import (
"path/filepath"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
func TestObservabilityBufferStoreUpsertReplayAndAck(t *testing.T) {
store := NewObservabilityBufferStore(filepath.Join(t.TempDir(), "observability-buffer.json"))
if err := store.Upsert(ObservabilityBufferRecord{
WindowStartedAtUnix: 1710403200,
Snapshot: &protocol.NodeMetricSnapshot{CapturedAtUnix: 1710403205},
TrafficReport: &protocol.NodeTrafficReport{WindowStartedAtUnix: 1710403200, WindowEndedAtUnix: 1710403260, RequestCount: 5},
QueuedAtUnix: 1710403205,
}, 1710403000); err != nil {
t.Fatalf("first upsert failed: %v", err)
}
if err := store.Upsert(ObservabilityBufferRecord{
WindowStartedAtUnix: 1710403200,
Snapshot: &protocol.NodeMetricSnapshot{CapturedAtUnix: 1710403255},
TrafficReport: &protocol.NodeTrafficReport{WindowStartedAtUnix: 1710403200, WindowEndedAtUnix: 1710403260, RequestCount: 12},
QueuedAtUnix: 1710403255,
}, 1710403000); err != nil {
t.Fatalf("second upsert failed: %v", err)
}
if err := store.Upsert(ObservabilityBufferRecord{
WindowStartedAtUnix: 1710403260,
Snapshot: &protocol.NodeMetricSnapshot{CapturedAtUnix: 1710403265},
TrafficReport: &protocol.NodeTrafficReport{WindowStartedAtUnix: 1710403260, WindowEndedAtUnix: 1710403320, RequestCount: 2},
QueuedAtUnix: 1710403265,
}, 1710403000); err != nil {
t.Fatalf("third upsert failed: %v", err)
}
records, err := store.Replayable(1710403260, 1710403000)
if err != nil {
t.Fatalf("Replayable failed: %v", err)
}
if len(records) != 1 {
t.Fatalf("expected one replayable record before current window, got %d", len(records))
}
if records[0].TrafficReport == nil || records[0].TrafficReport.RequestCount != 12 {
t.Fatalf("expected replayable record to keep latest upsert, got %+v", records[0])
}
if err = store.Ack([]int64{1710403200}, 1710403000); err != nil {
t.Fatalf("Ack failed: %v", err)
}
records, err = store.Replayable(0, 1710403000)
if err != nil {
t.Fatalf("Replayable after ack failed: %v", err)
}
if len(records) != 1 || records[0].WindowStartedAtUnix != 1710403260 {
t.Fatalf("unexpected records after ack: %+v", records)
}
}
func TestObservabilityBufferStoreMergesAccessLogsWithinWindow(t *testing.T) {
store := NewObservabilityBufferStore(filepath.Join(t.TempDir(), "observability-buffer.json"))
if err := store.Upsert(ObservabilityBufferRecord{
WindowStartedAtUnix: 1710403200,
AccessLogs: []protocol.NodeAccessLog{
{LoggedAtUnix: 1710403201, RemoteAddr: "10.0.0.1", Host: "app.example.com", Path: "/a", StatusCode: 200},
},
}, 1710403000); err != nil {
t.Fatalf("first upsert failed: %v", err)
}
if err := store.Upsert(ObservabilityBufferRecord{
WindowStartedAtUnix: 1710403200,
AccessLogs: []protocol.NodeAccessLog{
{LoggedAtUnix: 1710403201, RemoteAddr: "10.0.0.1", Host: "app.example.com", Path: "/a", StatusCode: 200},
{LoggedAtUnix: 1710403205, RemoteAddr: "10.0.0.2", Host: "app.example.com", Path: "/b", StatusCode: 502},
},
}, 1710403000); err != nil {
t.Fatalf("second upsert failed: %v", err)
}
records, err := store.Replayable(0, 1710403000)
if err != nil {
t.Fatalf("Replayable failed: %v", err)
}
if len(records) != 1 || len(records[0].AccessLogs) != 2 {
t.Fatalf("expected merged access logs, got %+v", records)
}
}
func TestObservabilityWindowStartedAt(t *testing.T) {
if value := ObservabilityWindowStartedAt(nil, nil, &protocol.NodeTrafficReport{WindowStartedAtUnix: 1710403200}); value != 1710403200 {
t.Fatalf("unexpected traffic window start: %d", value)
}
if value := ObservabilityWindowStartedAt(&protocol.NodeMetricSnapshot{CapturedAtUnix: 1710403259}, nil, nil); value != 1710403200 {
t.Fatalf("unexpected snapshot-derived window start: %d", value)
}
}
+106
View File
@@ -0,0 +1,106 @@
package state
import (
"crypto/rand"
"encoding/hex"
"encoding/json"
"os"
"path/filepath"
"sync"
)
type Snapshot struct {
NodeID string `json:"node_id"`
CurrentVersion string `json:"current_version"`
CurrentChecksum string `json:"current_checksum"`
BlockedVersion string `json:"blocked_version"`
BlockedChecksum string `json:"blocked_checksum"`
BlockedReason string `json:"blocked_reason"`
LastError string `json:"last_error"`
OpenrestyStatus string `json:"openresty_status"`
OpenrestyMessage string `json:"openresty_message"`
LastProfileFingerprint string `json:"last_profile_fingerprint"`
LastCPUStatTotal uint64 `json:"last_cpu_stat_total"`
LastCPUStatIdle uint64 `json:"last_cpu_stat_idle"`
LastMetricAtUnix int64 `json:"last_metric_at_unix"`
AccessLogOffset int64 `json:"access_log_offset"`
}
type Store struct {
path string
mu sync.Mutex
}
func NewStore(path string) *Store {
return &Store{path: filepath.Clean(path)}
}
func (s *Store) Load() (*Snapshot, error) {
s.mu.Lock()
defer s.mu.Unlock()
return s.loadUnlocked()
}
func (s *Store) EnsureNodeID() (string, error) {
s.mu.Lock()
defer s.mu.Unlock()
snapshot, err := s.loadUnlocked()
if err != nil {
return "", err
}
if snapshot.NodeID != "" {
return snapshot.NodeID, nil
}
snapshot.NodeID, err = newNodeID()
if err != nil {
return "", err
}
if err = s.saveUnlocked(snapshot); err != nil {
return "", err
}
return snapshot.NodeID, nil
}
func (s *Store) Save(snapshot *Snapshot) error {
s.mu.Lock()
defer s.mu.Unlock()
return s.saveUnlocked(snapshot)
}
func (s *Store) loadUnlocked() (*Snapshot, error) {
data, err := os.ReadFile(s.path)
if err != nil {
if os.IsNotExist(err) {
return &Snapshot{}, nil
}
return nil, err
}
snapshot := &Snapshot{}
if len(data) == 0 {
return snapshot, nil
}
if err = json.Unmarshal(data, snapshot); err != nil {
return nil, err
}
return snapshot, nil
}
func (s *Store) saveUnlocked(snapshot *Snapshot) error {
if err := os.MkdirAll(filepath.Dir(s.path), 0o755); err != nil {
return err
}
data, err := json.MarshalIndent(snapshot, "", " ")
if err != nil {
return err
}
return os.WriteFile(s.path, data, 0o644)
}
func newNodeID() (string, error) {
buf := make([]byte, 8)
if _, err := rand.Read(buf); err != nil {
return "", err
}
return "node-" + hex.EncodeToString(buf), nil
}
+150
View File
@@ -0,0 +1,150 @@
package state
import (
"fmt"
"os"
"path/filepath"
"sync"
"testing"
)
func TestEnsureNodeIDPersists(t *testing.T) {
store := NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID1, err := store.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
nodeID2, err := store.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID second call failed: %v", err)
}
if nodeID1 == "" || nodeID1 != nodeID2 {
t.Fatal("expected node id to persist across calls")
}
}
func TestStore_Load_NonExistentFile(t *testing.T) {
// Loading from a non-existent path should succeed and return an empty Snapshot
tempFile := filepath.Join(t.TempDir(), "nonexistent.json")
store := NewStore(tempFile)
snap, err := store.Load()
if err != nil {
t.Fatalf("expected Load to succeed for non-existent file, got err: %v", err)
}
if snap == nil {
t.Fatal("expected non-nil snapshot")
}
if snap.NodeID != "" || snap.CurrentVersion != "" {
t.Errorf("expected empty snapshot, got: %+v", snap)
}
}
func TestStore_Load_EmptyFile(t *testing.T) {
// Loading from an empty file should succeed and return an empty Snapshot
tempFile := filepath.Join(t.TempDir(), "empty.json")
if err := os.WriteFile(tempFile, []byte(""), 0644); err != nil {
t.Fatalf("failed to create empty file: %v", err)
}
store := NewStore(tempFile)
snap, err := store.Load()
if err != nil {
t.Fatalf("expected Load to succeed for empty file, got err: %v", err)
}
if snap == nil {
t.Fatal("expected non-nil snapshot")
}
if snap.NodeID != "" {
t.Errorf("expected empty snapshot, got: %+v", snap)
}
}
func TestStore_Load_InvalidJSON(t *testing.T) {
// Loading from a corrupted file with invalid JSON should fail with parsing error
tempFile := filepath.Join(t.TempDir(), "corrupted.json")
if err := os.WriteFile(tempFile, []byte("{invalid-json"), 0644); err != nil {
t.Fatalf("failed to create corrupted file: %v", err)
}
store := NewStore(tempFile)
_, err := store.Load()
if err == nil {
t.Fatal("expected Load to fail for corrupted JSON file")
}
}
func TestStore_SaveAndLoad(t *testing.T) {
tempFile := filepath.Join(t.TempDir(), "state.json")
store := NewStore(tempFile)
original := &Snapshot{
NodeID: "node-test-123",
CurrentVersion: "20260531-001",
CurrentChecksum: "chk-active-xyz",
BlockedVersion: "20260531-002",
BlockedChecksum: "chk-blocked-abc",
BlockedReason: "invalid upstream domain name",
LastError: "configuration reload timeout",
OpenrestyStatus: "unhealthy",
}
if err := store.Save(original); err != nil {
t.Fatalf("expected Save to succeed, got: %v", err)
}
loaded, err := store.Load()
if err != nil {
t.Fatalf("expected Load to succeed, got: %v", err)
}
if loaded.NodeID != original.NodeID ||
loaded.CurrentVersion != original.CurrentVersion ||
loaded.CurrentChecksum != original.CurrentChecksum ||
loaded.BlockedVersion != original.BlockedVersion ||
loaded.BlockedChecksum != original.BlockedChecksum ||
loaded.BlockedReason != original.BlockedReason ||
loaded.LastError != original.LastError ||
loaded.OpenrestyStatus != original.OpenrestyStatus {
t.Errorf("loaded snapshot does not match original: %+v vs %+v", loaded, original)
}
}
func TestStore_ConcurrencySafety(t *testing.T) {
tempFile := filepath.Join(t.TempDir(), "state.json")
store := NewStore(tempFile)
var wg sync.WaitGroup
workers := 20
iterations := 50
// Run concurrent writers and readers
for i := 0; i < workers; i++ {
wg.Add(1)
go func(workerID int) {
defer wg.Done()
for j := 0; j < iterations; j++ {
// Concurrently save
snap := &Snapshot{
NodeID: fmt.Sprintf("node-%d", workerID),
CurrentVersion: fmt.Sprintf("v-%d", j),
}
if err := store.Save(snap); err != nil {
t.Errorf("Save failed under concurrency: %v", err)
}
// Concurrently load
if _, err := store.Load(); err != nil {
t.Errorf("Load failed under concurrency: %v", err)
}
// Concurrently ensure ID
if _, err := store.EnsureNodeID(); err != nil {
t.Errorf("EnsureNodeID failed under concurrency: %v", err)
}
}
}(i)
}
wg.Wait()
}
+328
View File
@@ -0,0 +1,328 @@
package sync
import (
"archive/zip"
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"path"
"path/filepath"
"strings"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
)
type pagesSourceDocument struct {
Routes []pagesSourceRoute `json:"routes"`
}
type pagesSourceRoute struct {
UpstreamType string `json:"upstream_type"`
PagesDeployment *pagesDeploymentSource `json:"pages_deployment"`
}
type pagesDeploymentSource struct {
DeploymentID uint `json:"deployment_id"`
Checksum string `json:"checksum"`
}
type pagesDeploymentMarker struct {
DeploymentID uint `json:"deployment_id"`
Checksum string `json:"checksum"`
}
func (s *Service) syncPagesDeployments(ctx context.Context, config *protocol.ActiveConfigResponse) error {
deployments, err := referencedPagesDeployments(config)
if err != nil {
return err
}
if len(deployments) == 0 {
return nil
}
if strings.TrimSpace(s.pagesDir) == "" {
return errors.New("pages_dir is required when active config references Pages deployments")
}
for _, deployment := range deployments {
if err := s.ensurePagesDeployment(ctx, deployment); err != nil {
return err
}
}
return nil
}
func (s *Service) ensurePagesDeployment(ctx context.Context, deployment pagesDeploymentSource) error {
currentDir := pagesCurrentDir(s.pagesDir, deployment.DeploymentID)
if markerMatches(currentDir, deployment) {
return nil
}
packageBytes, err := s.client.DownloadPagesDeploymentPackage(ctx, deployment.DeploymentID)
if err != nil {
return fmt.Errorf("download Pages deployment %d: %w", deployment.DeploymentID, err)
}
if got := checksumBytes(packageBytes); got != deployment.Checksum {
return fmt.Errorf("Pages deployment %d checksum mismatch: expected %s, got %s", deployment.DeploymentID, deployment.Checksum, got)
}
releaseDir := pagesReleaseDir(s.pagesDir, deployment.DeploymentID, deployment.Checksum)
if !markerMatches(releaseDir, deployment) {
if err := extractPagesPackage(packageBytes, releaseDir, deployment); err != nil {
return err
}
}
return switchPagesCurrentDir(s.pagesDir, deployment.DeploymentID, releaseDir)
}
func referencedPagesDeployments(config *protocol.ActiveConfigResponse) ([]pagesDeploymentSource, error) {
if config == nil || strings.TrimSpace(config.SourceConfigJSON) == "" {
return nil, nil
}
var doc pagesSourceDocument
if err := json.Unmarshal([]byte(config.SourceConfigJSON), &doc); err != nil {
return nil, fmt.Errorf("decode Pages references: %w", err)
}
seen := make(map[uint]struct{})
result := make([]pagesDeploymentSource, 0)
for _, route := range doc.Routes {
if strings.ToLower(strings.TrimSpace(route.UpstreamType)) != "pages" || route.PagesDeployment == nil {
continue
}
deploymentID := route.PagesDeployment.DeploymentID
checksum := strings.TrimSpace(route.PagesDeployment.Checksum)
if deploymentID == 0 || checksum == "" {
return nil, errors.New("Pages deployment snapshot is incomplete")
}
if _, ok := seen[deploymentID]; ok {
continue
}
seen[deploymentID] = struct{}{}
result = append(result, pagesDeploymentSource{DeploymentID: deploymentID, Checksum: checksum})
}
return result, nil
}
func findCommonRootPrefix(files []*zip.File) (string, error) {
var firstFilePath string
hasMultipleFiles := false
for _, item := range files {
relativePath, skip, err := normalizePagesArchivePath(item.Name)
if err != nil {
return "", err
}
if skip {
continue
}
normalizedPath := filepath.ToSlash(relativePath)
if firstFilePath == "" {
firstFilePath = normalizedPath
} else {
hasMultipleFiles = true
}
}
if firstFilePath == "" {
return "", nil
}
parts := strings.Split(firstFilePath, "/")
if len(parts) <= 1 {
return "", nil
}
commonPrefix := parts[0] + "/"
if hasMultipleFiles {
for _, item := range files {
relativePath, skip, err := normalizePagesArchivePath(item.Name)
if err != nil {
return "", err
}
if skip {
continue
}
normalizedPath := filepath.ToSlash(relativePath)
if !strings.HasPrefix(normalizedPath, commonPrefix) {
return "", nil
}
}
}
return commonPrefix, nil
}
func extractPagesPackage(packageBytes []byte, releaseDir string, deployment pagesDeploymentSource) error {
tmpDir := releaseDir + ".tmp"
_ = os.RemoveAll(tmpDir)
if err := os.MkdirAll(tmpDir, 0o755); err != nil {
return err
}
reader, err := zip.NewReader(bytes.NewReader(packageBytes), int64(len(packageBytes)))
if err != nil {
_ = os.RemoveAll(tmpDir)
return fmt.Errorf("open Pages zip: %w", err)
}
commonPrefix, err := findCommonRootPrefix(reader.File)
if err != nil {
_ = os.RemoveAll(tmpDir)
return err
}
for _, item := range reader.File {
relativePath, skip, err := normalizePagesArchivePath(item.Name)
if err != nil {
_ = os.RemoveAll(tmpDir)
return err
}
if skip {
continue
}
if commonPrefix != "" {
slashPath := filepath.ToSlash(relativePath)
if strings.HasPrefix(slashPath, commonPrefix) {
relativePath = filepath.FromSlash(strings.TrimPrefix(slashPath, commonPrefix))
}
}
if item.FileInfo().Mode()&os.ModeSymlink != 0 {
_ = os.RemoveAll(tmpDir)
return fmt.Errorf("Pages package contains unsupported symlink: %s", relativePath)
}
if err := extractPagesFile(item, filepath.Join(tmpDir, relativePath)); err != nil {
_ = os.RemoveAll(tmpDir)
return err
}
}
if err := writePagesMarker(tmpDir, deployment); err != nil {
_ = os.RemoveAll(tmpDir)
return err
}
_ = os.RemoveAll(releaseDir)
return os.Rename(tmpDir, releaseDir)
}
func extractPagesFile(item *zip.File, targetPath string) error {
if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
return err
}
source, err := item.Open()
if err != nil {
return err
}
defer source.Close()
target, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, item.FileInfo().Mode().Perm())
if err != nil {
return err
}
defer target.Close()
_, err = io.Copy(target, source)
return err
}
func switchPagesCurrentDir(baseDir string, deploymentID uint, releaseDir string) error {
currentDir := pagesCurrentDir(baseDir, deploymentID)
previousDir := currentDir + ".previous"
_ = os.RemoveAll(previousDir)
if err := os.MkdirAll(filepath.Dir(currentDir), 0o755); err != nil {
return err
}
if _, err := os.Stat(currentDir); err == nil {
if err := os.Rename(currentDir, previousDir); err != nil {
return err
}
}
if err := copyPagesDir(releaseDir, currentDir); err != nil {
_ = os.RemoveAll(currentDir)
if _, restoreErr := os.Stat(previousDir); restoreErr == nil {
_ = os.Rename(previousDir, currentDir)
}
return err
}
_ = os.RemoveAll(previousDir)
return nil
}
func copyPagesDir(sourceDir string, targetDir string) error {
return filepath.WalkDir(sourceDir, func(sourcePath string, entry os.DirEntry, err error) error {
if err != nil {
return err
}
relativePath, err := filepath.Rel(sourceDir, sourcePath)
if err != nil || relativePath == "." {
return err
}
targetPath := filepath.Join(targetDir, relativePath)
if entry.IsDir() {
return os.MkdirAll(targetPath, 0o755)
}
info, err := entry.Info()
if err != nil {
return err
}
input, err := os.Open(sourcePath)
if err != nil {
return err
}
defer input.Close()
if err := os.MkdirAll(filepath.Dir(targetPath), 0o755); err != nil {
return err
}
output, err := os.OpenFile(targetPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, info.Mode().Perm())
if err != nil {
return err
}
defer output.Close()
_, err = io.Copy(output, input)
return err
})
}
func normalizePagesArchivePath(raw string) (string, bool, error) {
name := strings.TrimSpace(filepath.ToSlash(raw))
if name == "" || strings.HasSuffix(name, "/") {
return "", true, nil
}
if strings.HasPrefix(name, "/") {
return "", false, fmt.Errorf("Pages package contains absolute path: %s", raw)
}
cleaned := path.Clean(name)
if cleaned == "." {
return "", true, nil
}
if cleaned == ".." || strings.HasPrefix(cleaned, "../") || strings.Contains(cleaned, "/../") {
return "", false, fmt.Errorf("Pages package path escapes deployment root: %s", raw)
}
return filepath.FromSlash(cleaned), false, nil
}
func markerMatches(dir string, deployment pagesDeploymentSource) bool {
data, err := os.ReadFile(filepath.Join(dir, ".openflare-pages.json"))
if err != nil {
return false
}
var marker pagesDeploymentMarker
if err := json.Unmarshal(data, &marker); err != nil {
return false
}
return marker.DeploymentID == deployment.DeploymentID && marker.Checksum == deployment.Checksum
}
func writePagesMarker(dir string, deployment pagesDeploymentSource) error {
data, err := json.Marshal(pagesDeploymentMarker{
DeploymentID: deployment.DeploymentID,
Checksum: deployment.Checksum,
})
if err != nil {
return err
}
return os.WriteFile(filepath.Join(dir, ".openflare-pages.json"), data, 0o644)
}
func pagesCurrentDir(baseDir string, deploymentID uint) string {
return filepath.Join(baseDir, "deployments", fmt.Sprintf("%d", deploymentID), "current")
}
func pagesReleaseDir(baseDir string, deploymentID uint, checksum string) string {
return filepath.Join(baseDir, "deployments", fmt.Sprintf("%d", deploymentID), "releases", checksum)
}
func checksumBytes(data []byte) string {
sum := sha256.Sum256(data)
return hex.EncodeToString(sum[:])
}
+533
View File
@@ -0,0 +1,533 @@
package sync
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"log/slog"
"sort"
"strings"
openrestyrender "github.com/Rain-kl/Wavelet/pkg/render/openresty"
"github.com/Rain-kl/Wavelet/internal/apps/agent/nginx"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
"github.com/Rain-kl/Wavelet/internal/apps/agent/state"
)
const (
ApplyResultSuccess = "success"
ApplyResultWarning = "warning"
ApplyResultFailed = "failed"
)
type ConfigClient interface {
GetActiveConfig(ctx context.Context) (*protocol.ActiveConfigResponse, error)
DownloadPagesDeploymentPackage(ctx context.Context, deploymentID uint) ([]byte, error)
ReportApplyLog(ctx context.Context, payload protocol.ApplyLogPayload) error
SyncWAFIPGroups(ctx context.Context, payload protocol.WAFIPGroupSyncRequest) (*protocol.WAFIPGroupSyncResponse, error)
}
type NginxManager interface {
Apply(ctx context.Context, mainConfig string, routeConfig string, supportFiles []protocol.SupportFile) nginx.ApplyOutcome
EnsureRuntime(ctx context.Context, recreate bool) error
EnsureSafeFallbackRuntime(ctx context.Context, reason string) error
CurrentChecksum() (string, error)
WAFIPGroupChecksums() (map[string]string, error)
SyncWAFIPGroups(groups []protocol.WAFIPGroup) error
}
type Service struct {
client ConfigClient
nginxManager NginxManager
stateStore *state.Store
pagesDir string
}
func (s *Service) SetPagesDir(path string) {
s.pagesDir = strings.TrimSpace(path)
}
func New(client ConfigClient, nginxManager NginxManager, stateStore *state.Store) *Service {
return &Service{
client: client,
nginxManager: nginxManager,
stateStore: stateStore,
}
}
func (s *Service) SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error {
return s.sync(ctx, false, target)
}
func (s *Service) SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error {
return s.sync(ctx, true, target)
}
func (s *Service) sync(ctx context.Context, startup bool, target *protocol.ActiveConfigMeta) error {
mode := "periodic"
if startup {
mode = "startup"
}
snapshot, err := s.stateStore.Load()
if err != nil {
return err
}
currentChecksum, err := s.nginxManager.CurrentChecksum()
if err != nil {
return err
}
if target != nil {
target.Version = strings.TrimSpace(target.Version)
target.Checksum = strings.TrimSpace(target.Checksum)
}
if target == nil || target.Version == "" || target.Checksum == "" {
if !startup {
slog.Debug("skipping sync because heartbeat returned no active config summary", "mode", mode)
return nil
}
slog.Debug("sync startup fallback: active config summary unavailable, fetching active config directly")
config, fetchErr := s.client.GetActiveConfig(ctx)
if fetchErr != nil {
slog.Error("fetch active config failed", "mode", mode, "error", fetchErr)
return fetchErr
}
target = &protocol.ActiveConfigMeta{
Version: config.Version,
Checksum: config.Checksum,
}
return s.applyIfNeeded(ctx, mode, startup, snapshot, currentChecksum, target, config)
}
if currentChecksum == target.Checksum {
if startup {
config, fetchErr := s.client.GetActiveConfig(ctx)
if fetchErr != nil {
slog.Error("fetch active config failed", "mode", mode, "error", fetchErr)
return fetchErr
}
return s.applyIfNeeded(ctx, mode, startup, snapshot, currentChecksum, target, config)
}
slog.Debug("local openresty config already up to date", "mode", mode, "version", target.Version)
shouldReport := shouldReportNoopApply(snapshot, target.Version, target.Checksum)
if shouldReport {
if err = s.reportNoopApply(ctx, snapshot.NodeID, target.Version, target.Checksum, "", "", 0); err != nil {
return err
}
}
snapshot.CurrentVersion = target.Version
snapshot.CurrentChecksum = target.Checksum
clearBlockedTarget(snapshot)
snapshot.LastError = ""
slog.Debug("sync finished without changes", "mode", mode, "version", target.Version)
return s.stateStore.Save(snapshot)
}
if isBlockedTarget(snapshot, target.Version, target.Checksum) {
slog.Warn("skipping blocked config version after previous failed apply", "mode", mode, "version", target.Version, "checksum", target.Checksum)
if startup {
if err = s.ensureRuntimeForCurrentConfig(ctx, mode, snapshot, currentChecksum); err != nil {
return err
}
return s.stateStore.Save(snapshot)
}
return nil
}
if hasBlockedTarget(snapshot) {
clearBlockedTarget(snapshot)
}
if snapshot.CurrentVersion == target.Version && snapshot.CurrentChecksum == target.Checksum && !startup {
slog.Debug("skipping config fetch because state already records target version/checksum", "version", target.Version, "checksum", target.Checksum)
return s.stateStore.Save(snapshot)
}
config, err := s.client.GetActiveConfig(ctx)
if err != nil {
slog.Error("fetch active config failed", "mode", mode, "error", err)
return err
}
return s.applyIfNeeded(ctx, mode, startup, snapshot, currentChecksum, target, config)
}
func (s *Service) ForceSyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error {
snapshot, err := s.stateStore.Load()
if err != nil {
return err
}
if hasBlockedTarget(snapshot) {
clearBlockedTarget(snapshot)
_ = s.stateStore.Save(snapshot)
}
currentChecksum, err := s.nginxManager.CurrentChecksum()
if err != nil {
return err
}
config, err := s.client.GetActiveConfig(ctx)
if err != nil {
slog.Error("fetch active config failed", "mode", "force", "error", err)
return err
}
return s.applyIfNeeded(ctx, "force", true, snapshot, currentChecksum, target, config)
}
func (s *Service) WAFIPGroupChecksums() (map[string]string, error) {
if s.nginxManager == nil {
return map[string]string{}, nil
}
return s.nginxManager.WAFIPGroupChecksums()
}
func (s *Service) ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error {
if len(groups) == 0 || s.nginxManager == nil {
return nil
}
return s.nginxManager.SyncWAFIPGroups(groups)
}
func (s *Service) applyIfNeeded(ctx context.Context, mode string, startup bool, snapshot *state.Snapshot, currentChecksum string, target *protocol.ActiveConfigMeta, config *protocol.ActiveConfigResponse) error {
if currentChecksum == config.Checksum && !startup {
slog.Debug("local openresty config already up to date", "mode", mode, "version", config.Version)
shouldReport := shouldReportNoopApply(snapshot, config.Version, config.Checksum)
if shouldReport {
rendered, renderErr := renderActiveConfig(config)
if renderErr != nil {
return renderErr
}
if err := s.reportNoopApply(ctx, snapshot.NodeID, config.Version, config.Checksum, checksumString(rendered.mainConfig), checksumString(rendered.routeConfig), len(rendered.supportFiles)); err != nil {
return err
}
}
snapshot.CurrentVersion = config.Version
snapshot.CurrentChecksum = config.Checksum
clearBlockedTarget(snapshot)
snapshot.LastError = ""
slog.Debug("sync finished without changes", "mode", mode, "version", config.Version)
return s.stateStore.Save(snapshot)
}
if target != nil && (target.Version != config.Version || target.Checksum != config.Checksum) {
slog.Warn("active config changed between heartbeat and fetch", "heartbeat_version", target.Version, "heartbeat_checksum", target.Checksum, "fetched_version", config.Version, "fetched_checksum", config.Checksum)
}
if isBlockedTarget(snapshot, config.Version, config.Checksum) {
slog.Warn("skipping blocked config after fetch because the same version previously failed", "mode", mode, "version", config.Version, "checksum", config.Checksum)
if startup {
if err := s.ensureRuntimeForCurrentConfig(ctx, mode, snapshot, currentChecksum); err != nil {
return err
}
return s.stateStore.Save(snapshot)
}
return nil
}
if hasBlockedTarget(snapshot) {
clearBlockedTarget(snapshot)
}
if snapshot.CurrentVersion == config.Version && snapshot.CurrentChecksum == config.Checksum && !startup {
slog.Debug("skipping apply because state already records target version/checksum", "version", config.Version, "checksum", config.Checksum)
return s.stateStore.Save(snapshot)
}
rendered, err := renderActiveConfig(config)
if err != nil {
return err
}
if err := s.syncPagesDeployments(ctx, config); err != nil {
return err
}
mainConfigChecksum := checksumString(rendered.mainConfig)
routeConfigChecksum := checksumString(rendered.routeConfig)
slog.Info("applying new openresty config", "mode", mode, "from_version", snapshot.CurrentVersion, "to_version", config.Version, "old_checksum", currentChecksum, "new_checksum", config.Checksum)
outcome := s.nginxManager.Apply(ctx, rendered.mainConfig, rendered.routeConfig, rendered.supportFiles)
message := strings.TrimSpace(outcome.Message)
if outcome.Status == "" {
outcome.Status = nginx.ApplyStatusFatal
if message == "" {
message = "openresty apply returned empty outcome"
}
}
reportResult := ApplyResultFailed
switch outcome.Status {
case nginx.ApplyStatusSuccess:
slog.Info("openresty config applied successfully", "mode", mode, "version", config.Version)
snapshot.CurrentVersion = config.Version
snapshot.CurrentChecksum = config.Checksum
clearBlockedTarget(snapshot)
snapshot.LastError = ""
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = ""
reportResult = ApplyResultSuccess
if message == "" {
message = "apply success"
}
case nginx.ApplyStatusWarning:
if message == "" {
message = "apply rolled back to previous config"
}
slog.Warn("openresty config apply rolled back", "mode", mode, "version", config.Version, "message", message)
markBlockedTarget(snapshot, config.Version, config.Checksum, message)
snapshot.LastError = message
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = message
reportResult = ApplyResultWarning
default:
if message == "" {
message = "openresty apply failed"
}
slog.Error("apply openresty config failed", "mode", mode, "version", config.Version, "message", message)
markBlockedTarget(snapshot, config.Version, config.Checksum, message)
snapshot.LastError = message
snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
snapshot.OpenrestyMessage = message
}
if err := s.stateStore.Save(snapshot); err != nil {
return err
}
if err := s.client.ReportApplyLog(ctx, protocol.ApplyLogPayload{
NodeID: snapshot.NodeID,
Version: config.Version,
Result: reportResult,
Message: message,
Checksum: config.Checksum,
MainConfigChecksum: mainConfigChecksum,
RouteConfigChecksum: routeConfigChecksum,
SupportFileCount: len(rendered.supportFiles),
}); err != nil {
slog.Error("report apply log failed", "version", config.Version, "result", reportResult, "error", err)
return err
}
if reportResult == ApplyResultFailed {
slog.Warn("failed apply log reported", "version", config.Version)
return outcomeError(config.Version, message)
}
if err := s.syncReferencedWAFIPGroups(ctx, rendered.supportFiles); err != nil {
slog.Error("sync referenced waf ip groups failed", "version", config.Version, "error", err)
return err
}
slog.Debug("apply log reported", "version", config.Version, "result", reportResult)
return nil
}
func (s *Service) syncReferencedWAFIPGroups(ctx context.Context, supportFiles []protocol.SupportFile) error {
ids := referencedWAFIPGroupIDs(supportFiles)
if len(ids) == 0 {
return nil
}
checksums, err := s.WAFIPGroupChecksums()
if err != nil {
return err
}
response, err := s.client.SyncWAFIPGroups(ctx, protocol.WAFIPGroupSyncRequest{
IDs: ids,
Checksums: checksums,
})
if err != nil {
return err
}
if response == nil || len(response.Groups) == 0 {
return nil
}
return s.ApplyWAFIPGroups(ctx, response.Groups)
}
type renderedActiveConfig struct {
mainConfig string
routeConfig string
supportFiles []protocol.SupportFile
}
func renderActiveConfig(config *protocol.ActiveConfigResponse) (*renderedActiveConfig, error) {
if config == nil {
return nil, errors.New("active config is nil")
}
sourceJSON := strings.TrimSpace(config.SourceConfigJSON)
if sourceJSON == "" {
return nil, errors.New("active config source_config_json is empty")
}
rendered, err := openrestyrender.RenderJSON(sourceJSON, toOpenRestySupportFiles(config.SupportFiles))
if err != nil {
return nil, err
}
files := fromOpenRestySupportFiles(rendered.SupportFiles)
files = append(files, protocol.SupportFile{Path: openrestyrender.SourceConfigFileName, Content: sourceJSON})
return &renderedActiveConfig{
mainConfig: rendered.MainConfig,
routeConfig: rendered.RouteConfig,
supportFiles: files,
}, nil
}
func toOpenRestySupportFiles(files []protocol.SupportFile) []openrestyrender.SupportFile {
if len(files) == 0 {
return nil
}
result := make([]openrestyrender.SupportFile, 0, len(files))
for _, file := range files {
result = append(result, openrestyrender.SupportFile{Path: file.Path, Content: file.Content})
}
return result
}
func fromOpenRestySupportFiles(files []openrestyrender.SupportFile) []protocol.SupportFile {
if len(files) == 0 {
return nil
}
result := make([]protocol.SupportFile, 0, len(files))
for _, file := range files {
result = append(result, protocol.SupportFile{Path: file.Path, Content: file.Content})
}
return result
}
func referencedWAFIPGroupIDs(supportFiles []protocol.SupportFile) []uint {
var content string
for _, file := range supportFiles {
if file.Path == "waf_config.json" {
content = strings.TrimSpace(file.Content)
break
}
}
if content == "" {
return []uint{}
}
var payload struct {
RuleGroups []struct {
IPWhitelistGroups []uint `json:"ip_whitelist_group_ids"`
IPBlacklistGroups []uint `json:"ip_blacklist_group_ids"`
} `json:"rule_groups"`
}
if err := json.Unmarshal([]byte(content), &payload); err != nil {
slog.Debug("decode waf_config.json for ip group references failed", "error", err)
return []uint{}
}
seen := make(map[uint]struct{})
for _, group := range payload.RuleGroups {
for _, id := range group.IPWhitelistGroups {
if id > 0 {
seen[id] = struct{}{}
}
}
for _, id := range group.IPBlacklistGroups {
if id > 0 {
seen[id] = struct{}{}
}
}
}
ids := make([]uint, 0, len(seen))
for id := range seen {
ids = append(ids, id)
}
sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
return ids
}
func shouldReportNoopApply(snapshot *state.Snapshot, version string, checksum string) bool {
if snapshot == nil {
return false
}
return strings.TrimSpace(snapshot.CurrentVersion) != strings.TrimSpace(version) ||
strings.TrimSpace(snapshot.CurrentChecksum) != strings.TrimSpace(checksum)
}
func (s *Service) reportNoopApply(ctx context.Context, nodeID string, version string, checksum string, mainConfigChecksum string, routeConfigChecksum string, supportFileCount int) error {
message := "local config already matches active version; apply skipped"
if err := s.client.ReportApplyLog(ctx, protocol.ApplyLogPayload{
NodeID: nodeID,
Version: strings.TrimSpace(version),
Result: ApplyResultSuccess,
Message: message,
Checksum: strings.TrimSpace(checksum),
MainConfigChecksum: strings.TrimSpace(mainConfigChecksum),
RouteConfigChecksum: strings.TrimSpace(routeConfigChecksum),
SupportFileCount: supportFileCount,
}); err != nil {
slog.Error("report noop apply log failed", "version", version, "error", err)
return err
}
slog.Debug("noop apply log reported", "version", version)
return nil
}
func outcomeError(version string, message string) error {
trimmed := strings.TrimSpace(message)
if trimmed == "" {
trimmed = "openresty apply failed"
}
return fmt.Errorf("apply version %s failed: %s", version, trimmed)
}
func (s *Service) ensureRuntimeForCurrentConfig(ctx context.Context, mode string, snapshot *state.Snapshot, currentChecksum string) error {
if strings.TrimSpace(currentChecksum) == "" {
slog.Warn("blocked config cannot be retried and no local checksum is available for runtime recovery", "mode", mode, "blocked_version", snapshot.BlockedVersion)
reason := fmt.Sprintf("blocked config %s has no valid local config available for runtime recovery", strings.TrimSpace(snapshot.BlockedVersion))
if err := s.nginxManager.EnsureSafeFallbackRuntime(ctx, reason); err != nil {
snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
snapshot.OpenrestyMessage = err.Error()
_ = s.stateStore.Save(snapshot)
return err
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = "safe default fallback runtime started"
return nil
}
slog.Info("ensuring runtime with current local config while active target remains blocked", "mode", mode, "current_version", snapshot.CurrentVersion, "current_checksum", currentChecksum, "blocked_version", snapshot.BlockedVersion)
if err := s.nginxManager.EnsureRuntime(ctx, true); err != nil {
if strings.TrimSpace(snapshot.CurrentChecksum) == "" {
reason := fmt.Sprintf("blocked config %s has no historical config and current local config cannot start: %v", strings.TrimSpace(snapshot.BlockedVersion), err)
if fallbackErr := s.nginxManager.EnsureSafeFallbackRuntime(ctx, reason); fallbackErr == nil {
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
snapshot.OpenrestyMessage = "safe default fallback runtime started"
return nil
} else {
err = fmt.Errorf("%v; fallback recovery failed: %w", err, fallbackErr)
}
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
snapshot.OpenrestyMessage = err.Error()
_ = s.stateStore.Save(snapshot)
return err
}
snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
if strings.TrimSpace(snapshot.OpenrestyMessage) == strings.TrimSpace(snapshot.BlockedReason) {
snapshot.OpenrestyMessage = ""
}
return nil
}
func markBlockedTarget(snapshot *state.Snapshot, version string, checksum string, reason string) {
if snapshot == nil {
return
}
snapshot.BlockedVersion = strings.TrimSpace(version)
snapshot.BlockedChecksum = strings.TrimSpace(checksum)
snapshot.BlockedReason = strings.TrimSpace(reason)
}
func clearBlockedTarget(snapshot *state.Snapshot) {
if snapshot == nil {
return
}
snapshot.BlockedVersion = ""
snapshot.BlockedChecksum = ""
snapshot.BlockedReason = ""
}
func hasBlockedTarget(snapshot *state.Snapshot) bool {
return snapshot != nil && (strings.TrimSpace(snapshot.BlockedVersion) != "" || strings.TrimSpace(snapshot.BlockedChecksum) != "")
}
func isBlockedTarget(snapshot *state.Snapshot, version string, checksum string) bool {
if snapshot == nil {
return false
}
return strings.TrimSpace(snapshot.BlockedVersion) == strings.TrimSpace(version) &&
strings.TrimSpace(snapshot.BlockedChecksum) == strings.TrimSpace(checksum) &&
(strings.TrimSpace(version) != "" || strings.TrimSpace(checksum) != "")
}
func checksumString(content string) string {
sum := sha256.Sum256([]byte(content))
return hex.EncodeToString(sum[:])
}
+923
View File
@@ -0,0 +1,923 @@
package sync
import (
"archive/zip"
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"os"
"path/filepath"
"strings"
"testing"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/nginx"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
"github.com/Rain-kl/Wavelet/internal/apps/agent/state"
)
type fakeExecutor struct {
testErr error
reloadErr error
}
func testPagesSourceConfigJSON(deploymentID uint, checksum string) string {
return fmt.Sprintf(`{"routes":[{"id":1,"site_name":"pages","domain":"pages.example.com","domains":["pages.example.com"],"origin_url":"openflare-pages://project/1","upstreams":["openflare-pages://project/1"],"enabled":true,"upstream_type":"pages","pages_deployment":{"project_id":1,"project_slug":"pages","deployment_id":%d,"deployment_number":1,"checksum":"%s","entry_file":"index.html","spa_fallback_enabled":true,"local_root":"__OPENFLARE_PAGES_DIR__/deployments/%d/current"}}],"openresty_config":{"worker_processes":"auto","worker_connections":1024,"worker_rlimit_nofile":65535,"events_multi_accept_enabled":true,"keepalive_timeout":20,"keepalive_requests":1000,"client_header_timeout":15,"client_body_timeout":15,"client_max_body_size":"64m","large_client_header_buffers":"4 16k","send_timeout":30,"proxy_connect_timeout":3,"proxy_send_timeout":60,"proxy_read_timeout":60,"websocket_enabled":true,"proxy_request_buffering":false,"proxy_buffering_enabled":true,"proxy_buffers":"16 16k","proxy_buffer_size":"8k","proxy_busy_buffers_size":"64k","gzip_enabled":true,"gzip_min_length":1024,"gzip_comp_level":5,"cache_enabled":false,"cache_levels":"1:2","cache_inactive":"30m","cache_max_size":"1g","cache_key_template":"$scheme$host$request_uri","cache_lock_enabled":true,"cache_lock_timeout":"5s","cache_use_stale":"error timeout updating http_500 http_502 http_503 http_504","main_config_template":"worker_processes {{OpenRestyWorkerProcesses}};"},"waf":{"rule_groups":[],"bindings":[]}}`, deploymentID, checksum, deploymentID)
}
type fakeClient struct {
config protocol.ActiveConfigResponse
reports []protocol.ApplyLogPayload
pagesPackages map[uint][]byte
fetchCalls int
}
type fakeManager struct {
applyOutcome nginx.ApplyOutcome
currentChecksum string
currentChecksumErr error
ensureErr error
fallbackErr error
ensureCalls []bool
fallbackReasons []string
applyMainContents []string
applyRouteContents []string
applyFiles [][]protocol.SupportFile
}
func testSourceConfigJSON(workerProcesses string, listen int) string {
return fmt.Sprintf(`{"routes":[{"id":1,"site_name":"example","domain":"example.com","domains":["example.com"],"origin_url":"http://127.0.0.1:%d","upstreams":["http://127.0.0.1:%d"],"enabled":true}],"openresty_config":{"worker_processes":"%s","worker_connections":1024,"worker_rlimit_nofile":65535,"events_multi_accept_enabled":true,"keepalive_timeout":20,"keepalive_requests":1000,"client_header_timeout":15,"client_body_timeout":15,"client_max_body_size":"64m","large_client_header_buffers":"4 16k","send_timeout":30,"proxy_connect_timeout":3,"proxy_send_timeout":60,"proxy_read_timeout":60,"websocket_enabled":true,"proxy_request_buffering":false,"proxy_buffering_enabled":true,"proxy_buffers":"16 16k","proxy_buffer_size":"8k","proxy_busy_buffers_size":"64k","gzip_enabled":true,"gzip_min_length":1024,"gzip_comp_level":5,"cache_enabled":false,"cache_levels":"1:2","cache_inactive":"30m","cache_max_size":"1g","cache_key_template":"$scheme$host$request_uri","cache_lock_enabled":true,"cache_lock_timeout":"5s","cache_use_stale":"error timeout updating http_500 http_502 http_503 http_504","main_config_template":"worker_processes {{OpenRestyWorkerProcesses}};"},"waf":{"rule_groups":[],"bindings":[]}}`, listen, listen, workerProcesses)
}
func (f *fakeExecutor) Test(ctx context.Context) error {
return f.testErr
}
func (f *fakeExecutor) Reload(ctx context.Context) error {
return f.reloadErr
}
func (f *fakeExecutor) EnsureRuntime(ctx context.Context, recreate bool) error {
return nil
}
func (f *fakeExecutor) CheckHealth(ctx context.Context) error {
return f.testErr
}
func (f *fakeExecutor) Restart(ctx context.Context) error {
return f.reloadErr
}
func (f *fakeClient) GetActiveConfig(ctx context.Context) (*protocol.ActiveConfigResponse, error) {
f.fetchCalls++
return &f.config, nil
}
func (f *fakeClient) DownloadPagesDeploymentPackage(ctx context.Context, deploymentID uint) ([]byte, error) {
if f.pagesPackages == nil {
return nil, fmt.Errorf("missing Pages package %d", deploymentID)
}
return f.pagesPackages[deploymentID], nil
}
func (f *fakeClient) ReportApplyLog(ctx context.Context, payload protocol.ApplyLogPayload) error {
f.reports = append(f.reports, payload)
return nil
}
func (f *fakeClient) SyncWAFIPGroups(ctx context.Context, payload protocol.WAFIPGroupSyncRequest) (*protocol.WAFIPGroupSyncResponse, error) {
return &protocol.WAFIPGroupSyncResponse{}, nil
}
func (m *fakeManager) Apply(ctx context.Context, mainConfig string, routeConfig string, supportFiles []protocol.SupportFile) nginx.ApplyOutcome {
m.applyMainContents = append(m.applyMainContents, mainConfig)
m.applyRouteContents = append(m.applyRouteContents, routeConfig)
m.applyFiles = append(m.applyFiles, append([]protocol.SupportFile(nil), supportFiles...))
if m.applyOutcome.Status == "" {
return nginx.ApplyOutcome{Status: nginx.ApplyStatusSuccess}
}
return m.applyOutcome
}
func (m *fakeManager) EnsureRuntime(ctx context.Context, recreate bool) error {
m.ensureCalls = append(m.ensureCalls, recreate)
return m.ensureErr
}
func (m *fakeManager) EnsureSafeFallbackRuntime(ctx context.Context, reason string) error {
m.fallbackReasons = append(m.fallbackReasons, reason)
return m.fallbackErr
}
func (m *fakeManager) CurrentChecksum() (string, error) {
return m.currentChecksum, m.currentChecksumErr
}
func (m *fakeManager) WAFIPGroupChecksums() (map[string]string, error) {
return map[string]string{}, nil
}
func (m *fakeManager) SyncWAFIPGroups(groups []protocol.WAFIPGroup) error {
return nil
}
func TestSyncOnceSuccess(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-001",
Checksum: "checksum-1",
SourceConfigJSON: testSourceConfigJSON("auto", 80),
SupportFiles: []protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
snapshot, _ := stateStore.Load()
snapshot.NodeID = nodeID
if err = stateStore.Save(snapshot); err != nil {
t.Fatalf("failed to save initial state: %v", err)
}
routePath := filepath.Join(t.TempDir(), "routes.conf")
service := New(client, &nginx.Manager{
MainConfigPath: filepath.Join(filepath.Dir(routePath), "nginx.conf"),
RouteConfigPath: routePath,
Executor: &fakeExecutor{},
}, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
data, err := os.ReadFile(routePath)
if err != nil {
t.Fatalf("failed to read route config: %v", err)
}
if !strings.Contains(string(data), "listen 80;") || !strings.Contains(string(data), "server_name example.com;") {
t.Fatal("expected rendered config to be written to route file")
}
mainData, err := os.ReadFile(filepath.Join(filepath.Dir(routePath), "nginx.conf"))
if err != nil {
t.Fatalf("failed to read main config: %v", err)
}
if string(mainData) != "worker_processes auto;" {
t.Fatal("expected main config to be written")
}
snapshot, err = stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.CurrentVersion != "20260309-001" || snapshot.CurrentChecksum != "checksum-1" {
t.Fatal("expected state store to persist current version and checksum")
}
if len(client.reports) != 1 || client.reports[0].Result != ApplyResultSuccess {
t.Fatal("expected successful apply report to be sent")
}
if client.reports[0].Checksum != "checksum-1" {
t.Fatalf("expected config checksum to be reported, got %q", client.reports[0].Checksum)
}
if client.reports[0].MainConfigChecksum == "" || client.reports[0].RouteConfigChecksum == "" {
t.Fatal("expected main and route config checksums to be reported")
}
if client.reports[0].SupportFileCount != 3 {
t.Fatalf("expected support file count to be reported, got %d", client.reports[0].SupportFileCount)
}
}
func TestSyncOnceDownloadsPagesDeploymentBeforeApply(t *testing.T) {
packageBytes := testPagesPackage(t, map[string]string{"index.html": "hello"})
checksum := testBytesChecksum(packageBytes)
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-101",
Checksum: "pages-config-checksum",
SourceConfigJSON: testPagesSourceConfigJSON(7, checksum),
CreatedAt: time.Now().Format(time.RFC3339),
},
pagesPackages: map[uint][]byte{7: packageBytes},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
snapshot, _ := stateStore.Load()
snapshot.NodeID = nodeID
if err = stateStore.Save(snapshot); err != nil {
t.Fatalf("save state failed: %v", err)
}
manager := &fakeManager{currentChecksum: "old-checksum"}
service := New(client, manager, stateStore)
pagesDir := t.TempDir()
service.SetPagesDir(pagesDir)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{Version: "20260309-101", Checksum: "pages-config-checksum"}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
data, err := os.ReadFile(filepath.Join(pagesDir, "deployments", "7", "current", "index.html"))
if err != nil {
t.Fatalf("expected Pages file to be extracted: %v", err)
}
if string(data) != "hello" {
t.Fatalf("unexpected Pages file content: %s", string(data))
}
if len(manager.applyRouteContents) != 1 || !strings.Contains(manager.applyRouteContents[0], "__OPENFLARE_PAGES_DIR__/deployments/7/current") {
t.Fatalf("expected Pages placeholder in rendered route config, got %#v", manager.applyRouteContents)
}
}
func TestSyncOnceRejectsPagesZipSlipBeforeApply(t *testing.T) {
packageBytes := testPagesPackage(t, map[string]string{"../escape.html": "bad", "index.html": "ok"})
checksum := testBytesChecksum(packageBytes)
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-102",
Checksum: "pages-config-checksum",
SourceConfigJSON: testPagesSourceConfigJSON(8, checksum),
CreatedAt: time.Now().Format(time.RFC3339),
},
pagesPackages: map[uint][]byte{8: packageBytes},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
if _, err := stateStore.EnsureNodeID(); err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
manager := &fakeManager{currentChecksum: "old-checksum"}
service := New(client, manager, stateStore)
service.SetPagesDir(t.TempDir())
err := service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{Version: "20260309-102", Checksum: "pages-config-checksum"})
if err == nil || !strings.Contains(err.Error(), "escapes deployment root") {
t.Fatalf("expected zip-slip rejection, got %v", err)
}
if len(manager.applyRouteContents) != 0 {
t.Fatalf("OpenResty apply must not run after Pages package rejection")
}
}
func TestSyncOnceRollbackOnNginxFailure(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-002",
Checksum: "checksum-2",
SourceConfigJSON: testSourceConfigJSON("2", 81),
SupportFiles: []protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-001",
CurrentChecksum: "checksum-1",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
service := New(client, &fakeManager{
applyOutcome: nginx.ApplyOutcome{
Status: nginx.ApplyStatusFatal,
Message: "openresty failed after rollback",
},
}, stateStore)
err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
})
if err == nil {
t.Fatal("expected SyncOnce to fail when apply outcome is fatal")
}
snapshot, loadErr := stateStore.Load()
if loadErr != nil {
t.Fatalf("failed to load state: %v", loadErr)
}
if snapshot.CurrentVersion != "20260309-001" {
t.Fatal("expected failed sync not to overwrite current version")
}
if snapshot.BlockedVersion != "20260309-002" || snapshot.BlockedChecksum != "checksum-2" {
t.Fatalf("expected failed target version to be blocked, got %+v", snapshot)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusUnhealthy {
t.Fatalf("expected unhealthy openresty status, got %q", snapshot.OpenrestyStatus)
}
if len(client.reports) != 1 || client.reports[0].Result != ApplyResultFailed {
t.Fatal("expected failed apply report to be sent")
}
if client.reports[0].Checksum != "checksum-2" {
t.Fatalf("expected failed report to retain target checksum, got %q", client.reports[0].Checksum)
}
if client.reports[0].MainConfigChecksum == "" || client.reports[0].RouteConfigChecksum == "" {
t.Fatal("expected failed report to include main and route config checksums")
}
if client.reports[0].SupportFileCount != 3 {
t.Fatalf("expected failed report to include support file count, got %d", client.reports[0].SupportFileCount)
}
}
func TestSyncOnceReportsWarningWhenRollbackKeepsOpenrestyHealthy(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-002",
Checksum: "checksum-2",
SourceConfigJSON: testSourceConfigJSON("2", 81),
SupportFiles: []protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-001",
CurrentChecksum: "checksum-1",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
service := New(client, &fakeManager{
applyOutcome: nginx.ApplyOutcome{
Status: nginx.ApplyStatusWarning,
Message: "apply failed, rolled back to previous config",
},
}, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
}); err != nil {
t.Fatalf("expected warning outcome to keep sync successful, got %v", err)
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.CurrentVersion != "20260309-001" || snapshot.CurrentChecksum != "checksum-1" {
t.Fatal("expected warning apply to keep previous version state")
}
if snapshot.BlockedVersion != "20260309-002" || snapshot.BlockedChecksum != "checksum-2" {
t.Fatalf("expected rolled-back target version to be blocked, got %+v", snapshot)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy {
t.Fatalf("expected healthy openresty after rollback, got %q", snapshot.OpenrestyStatus)
}
if snapshot.LastError == "" {
t.Fatal("expected rollback warning to be recorded")
}
if len(client.reports) != 1 || client.reports[0].Result != ApplyResultWarning {
t.Fatal("expected warning apply report to be sent")
}
}
func TestSyncOnStartupRecreatesRuntimeWhenChecksumMatches(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-003",
Checksum: "checksum-3",
SourceConfigJSON: testSourceConfigJSON("auto", 82),
SupportFiles: []protocol.SupportFile{{Path: "1.crt", Content: "cert"}},
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{NodeID: nodeID}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: "checksum-3"}
service := New(client, manager, stateStore)
if err = service.SyncOnStartup(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
}); err != nil {
t.Fatalf("SyncOnStartup failed: %v", err)
}
if len(manager.applyMainContents) != 1 {
t.Fatal("expected startup sync to re-render and apply local config")
}
if len(client.reports) != 1 || client.reports[0].Result != ApplyResultSuccess {
t.Fatal("expected startup sync to report apply success when state is refreshed")
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.CurrentChecksum != "checksum-3" || snapshot.CurrentVersion != "20260309-003" {
t.Fatal("expected snapshot to be refreshed from active config")
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy || snapshot.OpenrestyMessage != "" {
t.Fatal("expected startup sync to mark openresty healthy")
}
}
func TestSyncOnceReportsNoopWhenVersionChangesButChecksumMatches(t *testing.T) {
client := &fakeClient{}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-002",
CurrentChecksum: "checksum-3",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: "checksum-3"}
service := New(client, manager, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-003",
Checksum: "checksum-3",
}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
if client.fetchCalls != 0 {
t.Fatalf("expected checksum match to skip config fetch, got %d", client.fetchCalls)
}
if len(manager.applyMainContents) != 0 {
t.Fatal("expected checksum match to skip apply")
}
if len(client.reports) != 1 || client.reports[0].Result != ApplyResultSuccess {
t.Fatalf("expected noop apply success report, got %+v", client.reports)
}
if client.reports[0].Version != "20260309-003" || client.reports[0].Checksum != "checksum-3" {
t.Fatalf("unexpected noop apply report: %+v", client.reports[0])
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.CurrentVersion != "20260309-003" || snapshot.CurrentChecksum != "checksum-3" {
t.Fatalf("expected state to refresh active version, got %+v", snapshot)
}
}
func TestSyncOnceDoesNotRepeatNoopReportWhenStateAlreadyMatches(t *testing.T) {
client := &fakeClient{}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-003",
CurrentChecksum: "checksum-3",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: "checksum-3"}
service := New(client, manager, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-003",
Checksum: "checksum-3",
}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
if len(client.reports) != 0 {
t.Fatalf("expected matching state to skip duplicate noop report, got %+v", client.reports)
}
}
func TestSyncOnStartupRecordsRuntimeFailure(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-004",
Checksum: "checksum-4",
SourceConfigJSON: testSourceConfigJSON("4", 83),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{NodeID: nodeID}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{
currentChecksum: "checksum-4",
applyOutcome: nginx.ApplyOutcome{Status: nginx.ApplyStatusFatal, Message: context.DeadlineExceeded.Error()},
}
service := New(client, manager, stateStore)
if err = service.SyncOnStartup(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
}); err == nil {
t.Fatal("expected SyncOnStartup to fail when runtime recreation fails")
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusUnhealthy {
t.Fatalf("expected unhealthy openresty status, got %q", snapshot.OpenrestyStatus)
}
if snapshot.OpenrestyMessage == "" {
t.Fatal("expected runtime error message to be recorded")
}
}
func TestSyncOnceSkipsPreviouslyBlockedVersion(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-006",
Checksum: "checksum-6",
SourceConfigJSON: testSourceConfigJSON("6", 86),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-005",
CurrentChecksum: "checksum-5",
BlockedVersion: "20260309-006",
BlockedChecksum: "checksum-6",
BlockedReason: "apply failed, rolled back to previous config",
LastError: "apply failed, rolled back to previous config",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: "checksum-5"}
service := New(client, manager, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-006",
Checksum: "checksum-6",
}); err != nil {
t.Fatalf("expected blocked version to be skipped, got %v", err)
}
if client.fetchCalls != 0 {
t.Fatalf("expected blocked version to skip fetch, got %d", client.fetchCalls)
}
if len(manager.applyMainContents) != 0 {
t.Fatal("expected blocked version to skip apply")
}
if len(client.reports) != 0 {
t.Fatal("expected blocked version to skip reporting duplicate apply result")
}
}
func TestSyncOnStartupKeepsBlockedVersionSuppressedUntilNewTargetArrives(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-007",
Checksum: "checksum-7",
SourceConfigJSON: testSourceConfigJSON("7", 87),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-005",
CurrentChecksum: "checksum-5",
BlockedVersion: "20260309-007",
BlockedChecksum: "checksum-7",
BlockedReason: "apply failed, rolled back to previous config",
OpenrestyStatus: protocol.OpenrestyStatusUnhealthy,
OpenrestyMessage: "apply failed, rolled back to previous config",
LastError: "apply failed, rolled back to previous config",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: "checksum-5"}
service := New(client, manager, stateStore)
if err = service.SyncOnStartup(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-007",
Checksum: "checksum-7",
}); err != nil {
t.Fatalf("expected blocked startup target to be skipped, got %v", err)
}
if len(manager.ensureCalls) != 1 || !manager.ensureCalls[0] {
t.Fatal("expected startup skip to ensure runtime with current local config")
}
if client.fetchCalls != 0 {
t.Fatalf("expected blocked startup target to skip fetch, got %d", client.fetchCalls)
}
if len(client.reports) != 0 {
t.Fatal("expected blocked startup target to skip duplicate apply report")
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.BlockedVersion != "20260309-007" || snapshot.BlockedChecksum != "checksum-7" {
t.Fatalf("expected blocked target to remain recorded, got %+v", snapshot)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy {
t.Fatalf("expected startup runtime recovery to mark openresty healthy, got %q", snapshot.OpenrestyStatus)
}
}
func TestSyncOnStartupStartsFallbackWhenBlockedVersionHasNoLocalConfig(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-007",
Checksum: "checksum-7",
SourceConfigJSON: testSourceConfigJSON("7", 87),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
BlockedVersion: "20260309-007",
BlockedChecksum: "checksum-7",
BlockedReason: "apply failed, but fallback runtime started",
OpenrestyStatus: protocol.OpenrestyStatusUnhealthy,
OpenrestyMessage: "apply failed, but fallback runtime started",
LastError: "apply failed, but fallback runtime started",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{}
service := New(client, manager, stateStore)
if err = service.SyncOnStartup(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-007",
Checksum: "checksum-7",
}); err != nil {
t.Fatalf("expected blocked startup target to start fallback, got %v", err)
}
if len(manager.fallbackReasons) != 1 {
t.Fatalf("expected fallback runtime to be started once, got %d", len(manager.fallbackReasons))
}
if client.fetchCalls != 0 {
t.Fatalf("expected blocked startup target to skip fetch, got %d", client.fetchCalls)
}
if len(client.reports) != 0 {
t.Fatal("expected blocked startup target to skip duplicate apply report")
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.BlockedVersion != "20260309-007" || snapshot.BlockedChecksum != "checksum-7" {
t.Fatalf("expected blocked target to remain recorded, got %+v", snapshot)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy {
t.Fatalf("expected fallback startup recovery to mark openresty healthy, got %q", snapshot.OpenrestyStatus)
}
if snapshot.OpenrestyMessage != "safe default fallback runtime started" {
t.Fatalf("expected fallback status message, got %q", snapshot.OpenrestyMessage)
}
}
func TestSyncOnStartupStartsFallbackWhenResidualConfigCannotRecover(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-007",
Checksum: "checksum-7",
SourceConfigJSON: testSourceConfigJSON("7", 87),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
BlockedVersion: "20260309-007",
BlockedChecksum: "checksum-7",
BlockedReason: "apply failed, but fallback runtime started",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{
currentChecksum: "residual-checksum",
ensureErr: context.DeadlineExceeded,
}
service := New(client, manager, stateStore)
if err = service.SyncOnStartup(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-007",
Checksum: "checksum-7",
}); err != nil {
t.Fatalf("expected residual config failure to start fallback, got %v", err)
}
if len(manager.ensureCalls) != 1 {
t.Fatalf("expected residual config to be tested once, got %d", len(manager.ensureCalls))
}
if len(manager.fallbackReasons) != 1 {
t.Fatalf("expected fallback runtime to be started once, got %d", len(manager.fallbackReasons))
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy {
t.Fatalf("expected fallback startup recovery to mark openresty healthy, got %q", snapshot.OpenrestyStatus)
}
if snapshot.BlockedVersion != "20260309-007" || snapshot.BlockedChecksum != "checksum-7" {
t.Fatalf("expected blocked target to remain recorded, got %+v", snapshot)
}
}
func TestSyncOnceClearsBlockedTargetWhenNewVersionArrives(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-008",
Checksum: "checksum-8",
SourceConfigJSON: testSourceConfigJSON("8", 88),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: "20260309-005",
CurrentChecksum: "checksum-5",
BlockedVersion: "20260309-007",
BlockedChecksum: "checksum-7",
BlockedReason: "apply failed, rolled back to previous config",
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{}
service := New(client, manager, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: "20260309-008",
Checksum: "checksum-8",
}); err != nil {
t.Fatalf("expected new target version to be applied, got %v", err)
}
if client.fetchCalls != 1 {
t.Fatalf("expected new target to trigger fetch, got %d", client.fetchCalls)
}
if len(manager.applyMainContents) != 1 {
t.Fatal("expected new target to trigger apply")
}
snapshot, err := stateStore.Load()
if err != nil {
t.Fatalf("failed to load state: %v", err)
}
if snapshot.BlockedVersion != "" || snapshot.BlockedChecksum != "" {
t.Fatalf("expected blocked target to be cleared after new version succeeds, got %+v", snapshot)
}
if snapshot.CurrentVersion != "20260309-008" || snapshot.CurrentChecksum != "checksum-8" {
t.Fatalf("expected current version to move to new target, got %+v", snapshot)
}
}
func TestSyncOnceSkipsFetchWhenHeartbeatChecksumMatches(t *testing.T) {
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-005",
Checksum: "checksum-5",
SourceConfigJSON: testSourceConfigJSON("auto", 84),
CreatedAt: time.Now().Format(time.RFC3339),
},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
if err = stateStore.Save(&state.Snapshot{
NodeID: nodeID,
CurrentVersion: client.config.Version,
CurrentChecksum: client.config.Checksum,
}); err != nil {
t.Fatalf("failed to seed state: %v", err)
}
manager := &fakeManager{currentChecksum: client.config.Checksum}
service := New(client, manager, stateStore)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{
Version: client.config.Version,
Checksum: client.config.Checksum,
}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
if client.fetchCalls != 0 {
t.Fatalf("expected no active config fetch when heartbeat checksum matches, got %d", client.fetchCalls)
}
if len(client.reports) != 0 {
t.Fatal("expected no apply log when no config change is needed")
}
}
func testPagesPackage(t *testing.T, files map[string]string) []byte {
t.Helper()
var buffer bytes.Buffer
writer := zip.NewWriter(&buffer)
for name, content := range files {
file, err := writer.Create(name)
if err != nil {
t.Fatalf("create zip file failed: %v", err)
}
if _, err := file.Write([]byte(content)); err != nil {
t.Fatalf("write zip file failed: %v", err)
}
}
if err := writer.Close(); err != nil {
t.Fatalf("close zip failed: %v", err)
}
return buffer.Bytes()
}
func testBytesChecksum(data []byte) string {
sum := sha256.Sum256(data)
return hex.EncodeToString(sum[:])
}
func TestSyncOnceDownloadsPagesDeploymentWithTopLevelFolder(t *testing.T) {
packageBytes := testPagesPackage(t, map[string]string{
"Speed-Test-source/index.html": "hello html",
"Speed-Test-source/assets/app.js": "hello js",
})
checksum := testBytesChecksum(packageBytes)
client := &fakeClient{
config: protocol.ActiveConfigResponse{
Version: "20260309-105",
Checksum: "pages-config-checksum",
SourceConfigJSON: testPagesSourceConfigJSON(77, checksum),
CreatedAt: time.Now().Format(time.RFC3339),
},
pagesPackages: map[uint][]byte{77: packageBytes},
}
stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
nodeID, err := stateStore.EnsureNodeID()
if err != nil {
t.Fatalf("EnsureNodeID failed: %v", err)
}
snapshot, _ := stateStore.Load()
snapshot.NodeID = nodeID
if err = stateStore.Save(snapshot); err != nil {
t.Fatalf("save state failed: %v", err)
}
manager := &fakeManager{currentChecksum: "old-checksum"}
service := New(client, manager, stateStore)
pagesDir := t.TempDir()
service.SetPagesDir(pagesDir)
if err = service.SyncOnce(context.Background(), &protocol.ActiveConfigMeta{Version: "20260309-105", Checksum: "pages-config-checksum"}); err != nil {
t.Fatalf("SyncOnce failed: %v", err)
}
data, err := os.ReadFile(filepath.Join(pagesDir, "deployments", "77", "current", "index.html"))
if err != nil {
t.Fatalf("expected Pages index.html file to be extracted: %v", err)
}
if string(data) != "hello html" {
t.Fatalf("unexpected Pages index.html content: %s", string(data))
}
jsData, err := os.ReadFile(filepath.Join(pagesDir, "deployments", "77", "current", "assets", "app.js"))
if err != nil {
t.Fatalf("expected Pages assets/app.js file to be extracted: %v", err)
}
if string(jsData) != "hello js" {
t.Fatalf("unexpected Pages assets/app.js content: %s", string(jsData))
}
}
@@ -0,0 +1,51 @@
//go:build !windows
package updater
import (
"fmt"
"log/slog"
"os"
"syscall"
)
func replaceAndRestart(execPath string, tmpPath string) error {
backupPath := execPath + ".bak"
if err := removeBackupBinary(backupPath); err != nil {
return err
}
if err := os.Rename(execPath, backupPath); err != nil {
renameErr := err
if err := os.Remove(tmpPath); err != nil && !os.IsNotExist(err) {
slog.Error("remove tmp binary failed", "path", tmpPath, "error", err)
return fmt.Errorf("backup current binary: %w; remove tmp binary: %v", renameErr, err)
}
return fmt.Errorf("backup current binary: %w", renameErr)
}
if err := os.Rename(tmpPath, execPath); err != nil {
replaceErr := err
if err := os.Rename(backupPath, execPath); err != nil {
slog.Error("restore backup binary failed", "path", backupPath, "error", err)
return fmt.Errorf("replace binary: %w; restore backup binary: %v", replaceErr, err)
}
return fmt.Errorf("replace binary: %w", replaceErr)
}
if err := removeBackupBinary(backupPath); err != nil {
return err
}
if err := syscall.Exec(execPath, os.Args, os.Environ()); err != nil {
return fmt.Errorf("exec restart: %w", err)
}
return fmt.Errorf("unreachable after exec")
}
func removeBackupBinary(path string) error {
if err := os.Remove(path); err != nil {
if os.IsNotExist(err) {
return nil
}
slog.Error("remove backup binary failed", "path", path, "error", err)
return err
}
return nil
}
@@ -0,0 +1,15 @@
//go:build !windows
package updater
import (
"path/filepath"
"testing"
)
func TestRemoveBackupBinaryIgnoresMissingFile(t *testing.T) {
backupPath := filepath.Join(t.TempDir(), "openflare-agent.bak")
if err := removeBackupBinary(backupPath); err != nil {
t.Fatalf("expected missing backup cleanup to be ignored: %v", err)
}
}
@@ -0,0 +1,53 @@
//go:build windows
package updater
import (
"fmt"
"os"
"os/exec"
"strings"
)
func replaceAndRestart(execPath string, tmpPath string) error {
backupPath := execPath + ".bak"
scriptPath := execPath + ".update.cmd"
script := fmt.Sprintf(`@echo off
setlocal
:waitloop
move /Y "%s" "%s" >nul 2>nul
if errorlevel 1 (
ping 127.0.0.1 -n 2 >nul
goto waitloop
)
move /Y "%s" "%s" >nul 2>nul
if errorlevel 1 exit /b 1
start "" %s
del /Q "%s" >nul 2>nul
del /Q "%%~f0" >nul 2>nul
`, execPath, backupPath, tmpPath, execPath, buildWindowsCommandLine(execPath, os.Args[1:]), backupPath)
if err := os.WriteFile(scriptPath, []byte(script), 0o700); err != nil {
os.Remove(tmpPath)
return fmt.Errorf("write restart script: %w", err)
}
cmd := exec.Command("cmd", "/C", "start", "", scriptPath)
if err := cmd.Start(); err != nil {
os.Remove(scriptPath)
os.Remove(tmpPath)
return fmt.Errorf("schedule restart: %w", err)
}
os.Exit(0)
return nil
}
func buildWindowsCommandLine(execPath string, args []string) string {
parts := []string{quoteWindowsArg(execPath)}
for _, arg := range args {
parts = append(parts, quoteWindowsArg(arg))
}
return strings.Join(parts, " ")
}
func quoteWindowsArg(value string) string {
return `"` + strings.ReplaceAll(value, `"`, `""`) + `"`
}
+366
View File
@@ -0,0 +1,366 @@
package updater
import (
"context"
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"os"
"runtime"
"strings"
"time"
"github.com/Rain-kl/Wavelet/pkg/utils"
"github.com/Rain-kl/Wavelet/internal/apps/agent/agent"
"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
)
const maxChecksumAssetSize = 64 * 1024
var replaceAndRestartFunc = replaceAndRestart
type Service struct {
httpClient *http.Client
lastCheckKey string
}
func New() *Service {
return &Service{
httpClient: &http.Client{Timeout: 30 * time.Second},
}
}
type githubRelease struct {
TagName string `json:"tag_name"`
Prerelease bool `json:"prerelease"`
Draft bool `json:"draft"`
Assets []githubAsset `json:"assets"`
}
type githubAsset struct {
Name string `json:"name"`
BrowserDownloadURL string `json:"browser_download_url"`
}
func (s *Service) CheckAndUpdate(ctx context.Context, repo string, options agent.UpdateOptions) error {
release, err := s.getRelease(ctx, repo, options)
if err != nil {
return fmt.Errorf("check latest release: %w", err)
}
if release == nil || release.TagName == "" {
return nil
}
remoteVersion := normalizeVersion(release.TagName)
localVersion := normalizeVersion(config.Version)
checkKey := buildReleaseCheckKey(options, remoteVersion)
if remoteVersion == localVersion {
return nil
}
if !options.Force && checkKey != "" && checkKey == s.lastCheckKey {
return nil
}
if !isNewer(localVersion, remoteVersion) {
s.lastCheckKey = checkKey
return nil
}
slog.Info("agent update available", "from", localVersion, "to", remoteVersion)
assetName := assetNameForGOOSGOARCH(runtime.GOOS, runtime.GOARCH)
checksumAssetName := assetName + ".sha256"
var downloadURL string
var checksumURL string
for _, asset := range release.Assets {
switch asset.Name {
case assetName:
downloadURL = asset.BrowserDownloadURL
case checksumAssetName:
checksumURL = asset.BrowserDownloadURL
}
}
if downloadURL == "" {
s.lastCheckKey = checkKey
return fmt.Errorf("no matching asset %q in release %s", assetName, release.TagName)
}
if checksumURL == "" {
return fmt.Errorf("no matching checksum asset %q in release %s", checksumAssetName, release.TagName)
}
expectedChecksum, err := s.downloadChecksum(ctx, checksumURL, assetName)
if err != nil {
return fmt.Errorf("download checksum: %w", err)
}
execPath, err := os.Executable()
if err != nil {
return fmt.Errorf("get executable path: %w", err)
}
if err = s.downloadAndRestart(ctx, downloadURL, expectedChecksum, execPath); err != nil {
return fmt.Errorf("download and restart: %w", err)
}
s.lastCheckKey = checkKey
return nil
}
func (s *Service) getRelease(ctx context.Context, repo string, options agent.UpdateOptions) (*githubRelease, error) {
tagName := strings.TrimSpace(options.TagName)
if tagName != "" {
return s.getReleaseByTag(ctx, repo, tagName)
}
if strings.EqualFold(strings.TrimSpace(options.Channel), "preview") {
return s.getLatestPreviewRelease(ctx, repo)
}
return s.getLatestStableRelease(ctx, repo)
}
func (s *Service) getLatestStableRelease(ctx context.Context, repo string) (*githubRelease, error) {
url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo)
return s.fetchReleaseFromURL(ctx, url)
}
func (s *Service) getLatestPreviewRelease(ctx context.Context, repo string) (*githubRelease, error) {
url := fmt.Sprintf("https://api.github.com/repos/%s/releases?per_page=20", repo)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, err
}
req.Header.Set("Accept", "application/vnd.github+json")
resp, err := s.httpClient.Do(req)
if err != nil {
return nil, err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("github api returned %s", resp.Status)
}
var releases []githubRelease
if err = json.NewDecoder(resp.Body).Decode(&releases); err != nil {
return nil, err
}
for _, release := range releases {
if release.Draft || !release.Prerelease {
continue
}
releaseCopy := release
return &releaseCopy, nil
}
return nil, nil
}
func (s *Service) getReleaseByTag(ctx context.Context, repo string, tag string) (*githubRelease, error) {
url := fmt.Sprintf("https://api.github.com/repos/%s/releases/tags/%s", repo, strings.TrimSpace(tag))
return s.fetchReleaseFromURL(ctx, url)
}
func (s *Service) fetchReleaseFromURL(ctx context.Context, url string) (*githubRelease, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, err
}
req.Header.Set("Accept", "application/vnd.github+json")
resp, err := s.httpClient.Do(req)
if err != nil {
return nil, err
}
defer func(Body io.ReadCloser) {
err := Body.Close()
if err != nil {
slog.Error("failed to close response body", "error", err)
}
}(resp.Body)
if resp.StatusCode == http.StatusNotFound {
return nil, nil
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("github api returned %s", resp.Status)
}
return decodeRelease(resp.Body)
}
func decodeRelease(reader io.Reader) (*githubRelease, error) {
var release githubRelease
if err := json.NewDecoder(reader).Decode(&release); err != nil {
return nil, err
}
return &release, nil
}
func (s *Service) downloadChecksum(ctx context.Context, url string, assetName string) (string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return "", err
}
resp, err := s.httpClient.Do(req)
if err != nil {
return "", err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("checksum download returned %s", resp.Status)
}
content, err := io.ReadAll(io.LimitReader(resp.Body, maxChecksumAssetSize+1))
if err != nil {
return "", err
}
if len(content) > maxChecksumAssetSize {
return "", fmt.Errorf("checksum asset exceeds %d bytes", maxChecksumAssetSize)
}
checksum, err := parseSHA256Checksum(string(content), assetName)
if err != nil {
return "", err
}
return checksum, nil
}
func parseSHA256Checksum(content string, assetName string) (string, error) {
assetName = strings.TrimSpace(assetName)
for _, line := range strings.Split(content, "\n") {
line = strings.TrimSpace(line)
if line == "" || strings.HasPrefix(line, "#") {
continue
}
if checksum, ok := parseSHA256Line(line, assetName); ok {
return checksum, nil
}
}
if assetName == "" {
return "", fmt.Errorf("checksum asset does not contain a valid sha256 digest")
}
return "", fmt.Errorf("checksum asset does not contain a sha256 digest for %q", assetName)
}
func parseSHA256Line(line string, assetName string) (string, bool) {
fields := strings.Fields(line)
if len(fields) == 1 && isSHA256Hex(fields[0]) {
return strings.ToLower(fields[0]), true
}
if len(fields) >= 2 && isSHA256Hex(fields[0]) {
fileName := strings.TrimPrefix(strings.TrimSpace(fields[1]), "*")
if assetName == "" || fileName == assetName {
return strings.ToLower(fields[0]), true
}
}
prefix := "SHA256("
if strings.HasPrefix(line, prefix) {
closing := strings.Index(line, ")")
if closing > len(prefix) && closing+1 < len(line) {
fileName := strings.TrimSpace(line[len(prefix):closing])
rest := strings.TrimSpace(line[closing+1:])
rest = strings.TrimPrefix(rest, "=")
rest = strings.TrimSpace(rest)
if isSHA256Hex(rest) && (assetName == "" || fileName == assetName) {
return strings.ToLower(rest), true
}
}
}
return "", false
}
func isSHA256Hex(value string) bool {
value = strings.TrimSpace(value)
if len(value) != sha256.Size*2 {
return false
}
_, err := hex.DecodeString(value)
return err == nil
}
func (s *Service) downloadAndRestart(ctx context.Context, url string, expectedChecksum string, targetPath string) error {
expectedChecksum = strings.ToLower(strings.TrimSpace(expectedChecksum))
if !isSHA256Hex(expectedChecksum) {
return fmt.Errorf("invalid expected sha256 checksum")
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return err
}
resp, err := s.httpClient.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("download returned %s", resp.Status)
}
tmpPath := targetPath + ".update"
if runtime.GOOS == "windows" && !strings.HasSuffix(strings.ToLower(tmpPath), ".exe") {
tmpPath += ".exe"
}
tmpFile, err := os.OpenFile(tmpPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
if err != nil {
return err
}
hasher := sha256.New()
if _, err = io.Copy(io.MultiWriter(tmpFile, hasher), resp.Body); err != nil {
tmpFile.Close()
os.Remove(tmpPath)
return err
}
if err = tmpFile.Close(); err != nil {
os.Remove(tmpPath)
return err
}
actualChecksum := hex.EncodeToString(hasher.Sum(nil))
if actualChecksum != expectedChecksum {
os.Remove(tmpPath)
return fmt.Errorf("sha256 checksum mismatch: expected %s, got %s", expectedChecksum, actualChecksum)
}
if err = os.Chmod(tmpPath, 0o755); err != nil && runtime.GOOS != "windows" {
os.Remove(tmpPath)
return fmt.Errorf("set executable permission: %w", err)
}
slog.Info("agent binary updated, restarting")
return replaceAndRestartFunc(targetPath, tmpPath)
}
func assetNameForGOOSGOARCH(goos string, goarch string) string {
name := fmt.Sprintf("openflare-agent-%s-%s", goos, goarch)
if goos == "windows" {
return name + ".exe"
}
return name
}
func normalizeVersion(v string) string {
v = strings.TrimSpace(v)
v = strings.TrimPrefix(v, "v")
return v
}
func isNewer(local, remote string) bool {
return compareVersions(local, remote) < 0
}
func buildReleaseCheckKey(options agent.UpdateOptions, remoteVersion string) string {
channel := strings.TrimSpace(options.Channel)
if channel == "" {
channel = "stable"
}
if tagName := strings.TrimSpace(options.TagName); tagName != "" {
return channel + ":" + tagName
}
return channel + ":" + remoteVersion
}
func compareVersions(local string, remote string) int {
return utils.CompareVersions(local, remote)
}
+242
View File
@@ -0,0 +1,242 @@
package updater
import (
"context"
"crypto/sha256"
"encoding/hex"
"io"
"net/http"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"github.com/Rain-kl/Wavelet/internal/apps/agent/agent"
"github.com/Rain-kl/Wavelet/internal/apps/agent/config"
)
type roundTripFunc func(req *http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
func TestGetLatestPreviewRelease(t *testing.T) {
service := &Service{
httpClient: &http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
if req.URL.String() != "https://api.github.com/repos/Rain-kl/OpenFlare/releases?per_page=20" {
t.Fatalf("unexpected request url: %s", req.URL.String())
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`[
{"tag_name":"v1.0.0","prerelease":false},
{"tag_name":"v1.1.0-rc.1","prerelease":true}
]`)),
}, nil
}),
},
}
release, err := service.getRelease(context.Background(), "Rain-kl/OpenFlare", agent.UpdateOptions{Channel: "preview"})
if err != nil {
t.Fatalf("expected preview release query to succeed: %v", err)
}
if release == nil || release.TagName != "v1.1.0-rc.1" {
t.Fatalf("unexpected preview release: %#v", release)
}
}
func TestGetReleaseByTag(t *testing.T) {
service := &Service{
httpClient: &http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
if req.URL.String() != "https://api.github.com/repos/Rain-kl/OpenFlare/releases/tags/v1.1.0-rc.1" {
t.Fatalf("unexpected request url: %s", req.URL.String())
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"tag_name":"v1.1.0-rc.1","prerelease":true}`)),
}, nil
}),
},
}
release, err := service.getRelease(context.Background(), "Rain-kl/OpenFlare", agent.UpdateOptions{Channel: "preview", TagName: "v1.1.0-rc.1", Force: true})
if err != nil {
t.Fatalf("expected tag release query to succeed: %v", err)
}
if release == nil || release.TagName != "v1.1.0-rc.1" {
t.Fatalf("unexpected tag release: %#v", release)
}
}
func TestCheckAndUpdateRequiresChecksumAsset(t *testing.T) {
originalVersion := config.Version
config.Version = "v1.0.0"
t.Cleanup(func() {
config.Version = originalVersion
})
assetName := assetNameForGOOSGOARCH(runtime.GOOS, runtime.GOARCH)
service := &Service{
httpClient: &http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
if req.URL.String() != "https://api.github.com/repos/Rain-kl/OpenFlare/releases/latest" {
t.Fatalf("unexpected request url: %s", req.URL.String())
}
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{
"tag_name":"v1.0.1",
"assets":[
{"name":"` + assetName + `","browser_download_url":"https://example.test/agent"}
]
}`)),
}, nil
}),
},
}
err := service.CheckAndUpdate(context.Background(), "Rain-kl/OpenFlare", agent.UpdateOptions{})
if err == nil || !strings.Contains(err.Error(), "no matching checksum asset") {
t.Fatalf("expected missing checksum asset error, got %v", err)
}
}
func TestParseSHA256Checksum(t *testing.T) {
checksum := strings.Repeat("a", sha256.Size*2)
testCases := []struct {
name string
content string
asset string
want string
}{
{name: "single digest", content: checksum + "\n", asset: "openflare-agent-linux-amd64", want: checksum},
{name: "sha256sum format", content: checksum + " openflare-agent-linux-amd64\n", asset: "openflare-agent-linux-amd64", want: checksum},
{name: "bsd format", content: "SHA256(openflare-agent-linux-amd64)= " + checksum + "\n", asset: "openflare-agent-linux-amd64", want: checksum},
{name: "selects matching file", content: strings.Repeat("b", sha256.Size*2) + " other\n" + checksum + " openflare-agent-linux-amd64\n", asset: "openflare-agent-linux-amd64", want: checksum},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
got, err := parseSHA256Checksum(testCase.content, testCase.asset)
if err != nil {
t.Fatalf("expected checksum parse to succeed: %v", err)
}
if got != testCase.want {
t.Fatalf("unexpected checksum: got %s want %s", got, testCase.want)
}
})
}
}
func TestDownloadAndRestartVerifiesChecksum(t *testing.T) {
payload := []byte("new-agent-binary")
sum := sha256.Sum256(payload)
expectedChecksum := hex.EncodeToString(sum[:])
targetPath := filepath.Join(t.TempDir(), "openflare-agent")
if err := os.WriteFile(targetPath, []byte("old-agent-binary"), 0o755); err != nil {
t.Fatalf("write target: %v", err)
}
var replacedTarget string
var replacedTemp string
originalReplace := replaceAndRestartFunc
replaceAndRestartFunc = func(execPath string, tmpPath string) error {
replacedTarget = execPath
replacedTemp = tmpPath
return nil
}
t.Cleanup(func() {
replaceAndRestartFunc = originalReplace
})
service := &Service{
httpClient: &http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(string(payload))),
}, nil
}),
},
}
if err := service.downloadAndRestart(context.Background(), "https://example.test/agent", expectedChecksum, targetPath); err != nil {
t.Fatalf("expected verified download to succeed: %v", err)
}
if replacedTarget != targetPath {
t.Fatalf("unexpected replace target: %s", replacedTarget)
}
if replacedTemp == "" {
t.Fatal("expected replacement temp path to be recorded")
}
if _, err := os.Stat(replacedTemp); err != nil {
t.Fatalf("expected verified temp binary to remain for replacement: %v", err)
}
}
func TestDownloadAndRestartRejectsChecksumMismatch(t *testing.T) {
targetPath := filepath.Join(t.TempDir(), "openflare-agent")
if err := os.WriteFile(targetPath, []byte("old-agent-binary"), 0o755); err != nil {
t.Fatalf("write target: %v", err)
}
originalReplace := replaceAndRestartFunc
replaceAndRestartFunc = func(execPath string, tmpPath string) error {
t.Fatal("replace should not run on checksum mismatch")
return nil
}
t.Cleanup(func() {
replaceAndRestartFunc = originalReplace
})
service := &Service{
httpClient: &http.Client{
Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: http.StatusOK,
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader("tampered")),
}, nil
}),
},
}
err := service.downloadAndRestart(context.Background(), "https://example.test/agent", strings.Repeat("0", sha256.Size*2), targetPath)
if err == nil || !strings.Contains(err.Error(), "sha256 checksum mismatch") {
t.Fatalf("expected checksum mismatch error, got %v", err)
}
if _, err = os.Stat(targetPath + ".update"); !os.IsNotExist(err) {
t.Fatalf("expected temp update file to be removed, stat err=%v", err)
}
}
func TestIsNewerSupportsPrerelease(t *testing.T) {
testCases := []struct {
name string
local string
remote string
expected bool
}{
{name: "stable newer than prerelease", local: "1.2.3-rc.1", remote: "1.2.3", expected: true},
{name: "same stable not newer", local: "1.2.3", remote: "1.2.3-rc.1", expected: false},
{name: "higher prerelease sequence", local: "1.2.3-rc.1", remote: "1.2.3-rc.2", expected: true},
{name: "higher minor", local: "1.2.3", remote: "1.3.0-rc.1", expected: true},
}
for _, testCase := range testCases {
t.Run(testCase.name, func(t *testing.T) {
if actual := isNewer(testCase.local, testCase.remote); actual != testCase.expected {
t.Fatalf("unexpected compare result: local=%s remote=%s actual=%v expected=%v", testCase.local, testCase.remote, actual, testCase.expected)
}
})
}
}
+82
View File
@@ -0,0 +1,82 @@
package wsclient
import (
"context"
"time"
"github.com/Rain-kl/Wavelet/internal/apps/agent/protocol"
shared "github.com/Rain-kl/Wavelet/pkg/wsclient"
)
type WSMessage = shared.WSMessage
type MessageHandler = shared.MessageHandler
type Client struct {
sharedClient *shared.Client
}
type Connection struct {
sharedConn *shared.Connection
}
func New(baseURL string, token string, timeout time.Duration) *Client {
return &Client{
sharedClient: shared.New(shared.Config{
BaseURL: baseURL,
Token: token,
Timeout: timeout,
HeaderKey: "X-Agent-Token",
WSPath: "/api/v1/agent/ws",
}),
}
}
func (c *Client) SetToken(token string) {
c.sharedClient.SetToken(token)
}
func (c *Client) URL() string {
return c.sharedClient.URL()
}
func (c *Client) Connect(ctx context.Context) (protocol.WebSocketConnection, error) {
conn, err := c.sharedClient.Connect(ctx)
if err != nil {
return nil, err
}
return &Connection{sharedConn: conn}, nil
}
func (conn *Connection) URL() string {
if conn == nil || conn.sharedConn == nil {
return ""
}
return conn.sharedConn.URL
}
func (conn *Connection) SendStatus(payload protocol.NodePayload) error {
return conn.sharedConn.SendMessage(protocol.WSMessageTypeStatus, payload)
}
func (conn *Connection) SendPong() error {
return conn.sharedConn.SendMessage(protocol.WSMessageTypePong, nil)
}
func (conn *Connection) Receive() (protocol.WSMessage, error) {
var message protocol.WSMessage
if err := conn.sharedConn.Receive(&message); err != nil {
return message, err
}
return message, nil
}
func (conn *Connection) RunReceiveLoop(ctx context.Context, handler shared.MessageHandler) error {
return conn.sharedConn.RunReceiveLoop(ctx, handler)
}
func (conn *Connection) Close() error {
if conn == nil || conn.sharedConn == nil {
return nil
}
return conn.sharedConn.Close()
}
+10
View File
@@ -0,0 +1,10 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package cap 提供人机验证中间件
package cap
const (
errCapTokenMissing = "验证码验证失败,缺少验证码凭证" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
errCapTokenInvalidOrExpired = "验证码校验失败或已过期,请重试" //nolint:gosec // false positive: this is an error message, not hardcoded credentials
)
+196
View File
@@ -0,0 +1,196 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
// Package cap provides CAPTCHA and proof-of-work (PoW) verification services.
package cap
import (
"context"
"crypto/sha256"
"encoding/hex"
"strconv"
"strings"
"sync"
"time"
"github.com/Rain-kl/Wavelet/internal/config"
"github.com/Rain-kl/Wavelet/internal/db"
pkgcap "github.com/Rain-kl/Wavelet/pkg/cap"
)
const (
redeemTokenIDLength = 8 // 兑换 Token ID 字节长度
redeemVerTokenLength = 15 // 兑换验证 Token 字节长度
tokenPartsCount = 2 // 兑换 Token 由两部分组成
valuePartsCount = 2 // 存储值由 scope 和过期时间组成
)
// Manager orchestrates challenge generation and solution validation.
type Manager struct {
secret []byte
store pkgcap.Store
}
// NewManager creates a new CAPTCHA Manager.
func NewManager(secret []byte, store pkgcap.Store) *Manager {
return &Manager{
secret: secret,
store: store,
}
}
// Generate creates a challenge response.
func (m *Manager) Generate(ctx context.Context, scope string) (*pkgcap.ChallengeResponse, error) {
settings, err := CurrentSettings(ctx)
if err != nil {
return nil, err
}
challengeConfig := pkgcap.ChallengeConfig{
Count: settings.ChallengeCount,
Size: settings.ChallengeSize,
Difficulty: settings.ChallengeDifficulty,
Expires: settings.ChallengeTTL,
}
return pkgcap.GenerateChallenge(m.secret, challengeConfig, scope)
}
// RedeemResponse is returned to the client on redeem.
type RedeemResponse struct {
Success bool `json:"success"`
Token string `json:"token,omitempty"`
Expires int64 `json:"expires,omitempty"`
Error string `json:"error,omitempty"`
}
// Redeem verifies PoW solutions and returns a one-time redeem token.
func (m *Manager) Redeem(ctx context.Context, token string, solutions []int, scope string) (*RedeemResponse, error) {
sigHex := pkgcap.JwtSigHex(token)
if sigHex == "" {
return &RedeemResponse{Success: false, Error: "invalid_token"}, nil
}
nonceKey := "cap:nonce:" + sigHex
payload, err := pkgcap.VerifyChallengeSolutions(token, solutions, m.secret, scope)
if err != nil {
return &RedeemResponse{Success: false, Error: err.Error()}, nil //nolint:nilerr // validation errors are returned as response, not system errors
}
now := time.Now().UnixNano() / int64(time.Millisecond)
nonceTTL := time.Duration(payload.Expires-now) * time.Millisecond
if nonceTTL < time.Second {
nonceTTL = time.Second
}
set, err := m.store.SetNX(ctx, nonceKey, "1", nonceTTL)
if err != nil {
return &RedeemResponse{Success: false, Error: "nonce_store_error"}, err
}
if !set {
return &RedeemResponse{Success: false, Error: "already_redeemed"}, nil
}
settings, err := CurrentSettings(ctx)
if err != nil {
return &RedeemResponse{Success: false, Error: "settings_load_error"}, err
}
id := pkgcap.RandomHex(redeemTokenIDLength)
verToken := pkgcap.RandomHex(redeemVerTokenLength)
verHashBytes := sha256.Sum256([]byte(verToken))
verHashHex := hex.EncodeToString(verHashBytes[:])
tokenKey := "cap:token:" + id + ":" + verHashHex
tokenExpires := time.Now().Add(settings.TokenTTL)
storeVal := strconv.FormatInt(tokenExpires.UnixNano(), 10) + "|" + scope
if err := m.store.Set(ctx, tokenKey, storeVal, settings.TokenTTL); err != nil {
return &RedeemResponse{Success: false, Error: "token_store_error"}, err
}
return &RedeemResponse{
Success: true,
Token: id + ":" + verToken,
Expires: tokenExpires.UnixNano() / int64(time.Millisecond),
}, nil
}
// VerifyToken validates and consumes the redeem token (single-use).
func (m *Manager) VerifyToken(ctx context.Context, token string, expectedScope string) (bool, error) {
if token == "" {
return false, nil
}
parts := strings.Split(token, ":")
if len(parts) != tokenPartsCount {
return false, nil
}
id := parts[0]
verToken := parts[1]
verHashBytes := sha256.Sum256([]byte(verToken))
verHashHex := hex.EncodeToString(verHashBytes[:])
tokenKey := "cap:token:" + id + ":" + verHashHex
val, exists, err := sGetAndDelete(ctx, m.store, tokenKey)
if err != nil {
return false, err
}
if !exists {
return false, nil
}
valParts := strings.Split(val, "|")
if len(valParts) != valuePartsCount {
return false, nil
}
expNano, err := strconv.ParseInt(valParts[0], 10, 64)
if err != nil {
return false, nil //nolint:nilerr // invalid format is treated as validation failure
}
tokenScope := valParts[1]
if expectedScope != "" && tokenScope != expectedScope {
return false, nil
}
if time.Now().UnixNano() > expNano {
return false, nil
}
return true, nil
}
func sGetAndDelete(ctx context.Context, store pkgcap.Store, key string) (string, bool, error) {
if store == nil {
return "", false, nil
}
return store.GetAndDelete(ctx, key)
}
var (
defaultManager *Manager
once sync.Once
)
// GetDefaultManager yields the global singleton CAPTCHA manager.
func GetDefaultManager() *Manager {
once.Do(func() {
secret := []byte("default-captcha-secret-key-at-least-16-bytes")
if config.Config != nil && config.Config.App.SessionSecret != "" {
secret = []byte(config.Config.App.SessionSecret)
}
var store pkgcap.Store
if config.Config != nil && config.Config.Redis.Enabled && db.Redis != nil {
store = pkgcap.NewRedisStore(db.Redis)
} else {
store = pkgcap.NewMemoryStore(1 * time.Minute)
}
defaultManager = NewManager(secret, store)
})
return defaultManager
}
+162
View File
@@ -0,0 +1,162 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package cap
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
pkgcap "github.com/Rain-kl/Wavelet/pkg/cap"
)
func installTestManagerSettings(t *testing.T) func() {
t.Helper()
return InstallTestRuntimeSettings(RuntimeSettings{
ChallengeCount: 3,
ChallengeSize: 32,
ChallengeDifficulty: 3,
ChallengeTTL: 5 * time.Second,
TokenTTL: 10 * time.Second,
})
}
func TestCapFullFlow(t *testing.T) {
cleanup := installTestManagerSettings(t)
defer cleanup()
secret := []byte("a-very-long-secret-key-at-least-16-bytes")
store := pkgcap.NewMemoryStore(1 * time.Minute)
manager := NewManager(secret, store)
scope := "test-scope"
ctx := context.Background()
resp, err := manager.Generate(ctx, scope)
if err != nil {
t.Fatalf("Generate() error = %v", err)
}
if resp.Challenge.C != 3 {
t.Fatalf("Generate().Challenge.C = %d, want %d", resp.Challenge.C, 3)
}
solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
redeemResp, err := manager.Redeem(ctx, resp.Token, solutions, scope)
if err != nil {
t.Fatalf("Redeem() error = %v", err)
}
if !redeemResp.Success {
t.Fatalf("Redeem().Success = false, error = %s", redeemResp.Error)
}
if redeemResp.Token == "" {
t.Fatal("Redeem().Token is empty")
}
valid, err := manager.VerifyToken(ctx, redeemResp.Token, scope)
if err != nil {
t.Fatalf("VerifyToken() error = %v", err)
}
if !valid {
t.Fatal("VerifyToken() = false, want true")
}
validAgain, err := manager.VerifyToken(ctx, redeemResp.Token, scope)
if err != nil {
t.Fatalf("VerifyToken() second call error = %v", err)
}
if validAgain {
t.Fatal("VerifyToken() second call = true, want false")
}
}
func TestRedeemConcurrentRace(t *testing.T) {
const goroutines = 50
cleanup := installTestManagerSettings(t)
defer cleanup()
secret := []byte("race-test-secret-key-at-least-16-bytes")
store := pkgcap.NewMemoryStore(1 * time.Minute)
manager := NewManager(secret, store)
ctx := context.Background()
resp, err := manager.Generate(ctx, "login")
if err != nil {
t.Fatalf("Generate() error = %v", err)
}
solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
var (
wg sync.WaitGroup
success atomic.Int32
barrier = make(chan struct{})
)
for range goroutines {
wg.Add(1)
go func() {
defer wg.Done()
<-barrier
r, _ := manager.Redeem(ctx, resp.Token, solutions, "login")
if r != nil && r.Success {
success.Add(1)
}
}()
}
close(barrier)
wg.Wait()
if got := success.Load(); got != 1 {
t.Fatalf("successful Redeem count = %d, want %d", got, 1)
}
}
func TestVerifyTokenConcurrentRace(t *testing.T) {
const goroutines = 50
cleanup := installTestManagerSettings(t)
defer cleanup()
secret := []byte("race-test-secret-key-at-least-16-bytes")
store := pkgcap.NewMemoryStore(1 * time.Minute)
manager := NewManager(secret, store)
ctx := context.Background()
resp, err := manager.Generate(ctx, "login")
if err != nil {
t.Fatalf("Generate() error = %v", err)
}
solutions := pkgcap.Solve(resp.Token, resp.Challenge.C, resp.Challenge.S, resp.Challenge.D)
redeemResp, err := manager.Redeem(ctx, resp.Token, solutions, "login")
if err != nil || !redeemResp.Success {
t.Fatalf("Redeem() error = %v, resp = %+v", err, redeemResp)
}
var (
wg sync.WaitGroup
success atomic.Int32
barrier = make(chan struct{})
)
for range goroutines {
wg.Add(1)
go func() {
defer wg.Done()
<-barrier
ok, _ := manager.VerifyToken(ctx, redeemResp.Token, "login")
if ok {
success.Add(1)
}
}()
}
close(barrier)
wg.Wait()
if got := success.Load(); got != 1 {
t.Fatalf("successful VerifyToken count = %d, want %d", got, 1)
}
}
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package cap
import (
"github.com/gin-gonic/gin"
"github.com/Rain-kl/Wavelet/internal/common/response"
)
// VerifyMiddleware returns a Gin middleware that checks and consumes the X-Cap-Token header.
func VerifyMiddleware(mgr *Manager, scope string) gin.HandlerFunc {
return func(c *gin.Context) {
if !ProtectionEnabled(c.Request.Context()) {
c.Next()
return
}
token := c.GetHeader("X-Cap-Token")
if token == "" {
response.AbortUnauthorized(c, errCapTokenMissing)
return
}
valid, err := mgr.VerifyToken(c.Request.Context(), token, scope)
if err != nil || !valid {
response.AbortUnauthorized(c, errCapTokenInvalidOrExpired)
return
}
c.Next()
}
}

Some files were not shown because too many files have changed in this diff Show More