migrate

Wavelet/internal/apps/openflare/relay/errs.go

@@ -0,0 +1,9 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+const (
+	errAgentTokenInvalid     = "无权进行此操作，Agent Token 无效"
+	errRelayNodeTypeMismatch = "此节点不是 TunnelRelay 类型"
+)

Wavelet/internal/apps/openflare/relay/helpers.go

@@ -0,0 +1,112 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+import (
+	"net"
+	"strings"
+
+	"github.com/Rain-kl/Wavelet/internal/model"
+)
+
+func normalizeRelayStatus(status string) string {
+	switch strings.ToLower(strings.TrimSpace(status)) {
+	case "healthy":
+		return "healthy"
+	case "unhealthy":
+		return "unhealthy"
+	default:
+		return "unknown"
+	}
+}
+
+func normalizeReleaseChannel(channel string) string {
+	if strings.ToLower(strings.TrimSpace(channel)) == "preview" {
+		return "preview"
+	}
+	return "stable"
+}
+
+func resolveReportedNodeIP(reportedIP string, remoteAddr string) string {
+	reported := normalizeNodeIP(reportedIP)
+	remote := normalizeRemoteAddr(remoteAddr)
+	if reported == "" {
+		return remote
+	}
+	if isPublicNodeIP(reported) {
+		return reported
+	}
+	if isPublicNodeIP(remote) {
+		return remote
+	}
+	return reported
+}
+
+func normalizeNodeIP(raw string) string {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return ""
+	}
+	if host, _, err := net.SplitHostPort(raw); err == nil {
+		raw = host
+	}
+	raw = strings.Trim(raw, "[]")
+	return raw
+}
+
+func normalizeRemoteAddr(remoteAddr string) string {
+	remoteAddr = strings.TrimSpace(remoteAddr)
+	if remoteAddr == "" {
+		return ""
+	}
+	host, _, err := net.SplitHostPort(remoteAddr)
+	if err != nil {
+		return normalizeNodeIP(remoteAddr)
+	}
+	return normalizeNodeIP(host)
+}
+
+func isPublicNodeIP(raw string) bool {
+	ip := net.ParseIP(strings.TrimSpace(raw))
+	if ip == nil {
+		return false
+	}
+	if ip.IsLoopback() || ip.IsPrivate() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsUnspecified() {
+		return false
+	}
+	return true
+}
+
+func buildRelayConfig(node *model.OpenFlareNode) *Config {
+	if node == nil {
+		return nil
+	}
+	return &Config{
+		BindPort:         node.RelayBindPort,
+		VhostHTTPPort:    node.RelayVhostHTTPPort,
+		AuthToken:        node.RelayAuthToken,
+		LogLevel:         "info",
+		WebServerEnabled: node.RelayWebServerEnabled,
+	}
+}
+
+// BuildSettings returns runtime settings shared by relay and flared clients.
+func BuildSettings(node *model.OpenFlareNode, updateNow bool, updateChannel, updateTag string) *Settings {
+	autoUpdate := false
+	if node != nil {
+		autoUpdate = node.AutoUpdateEnabled
+	}
+	if strings.TrimSpace(updateChannel) == "" {
+		updateChannel = "stable"
+	}
+	return &Settings{
+		HeartbeatInterval:       model.AgentHeartbeatInterval,
+		WebsocketUpgradeEnabled: model.AgentWebsocketUpgradeEnabled,
+		AutoUpdate:              autoUpdate,
+		UpdateRepo:              model.AgentUpdateRepo,
+		UpdateNow:               updateNow,
+		UpdateChannel:           updateChannel,
+		UpdateTag:               strings.TrimSpace(updateTag),
+	}
+}

Wavelet/internal/apps/openflare/relay/logics.go

@@ -0,0 +1,117 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/Rain-kl/Wavelet/internal/db"
+	"github.com/Rain-kl/Wavelet/internal/model"
+)
+
+const nodeStatusOnline = "online"
+
+// HeartbeatPayload is sent by OpenFlareRelay on each heartbeat.
+type HeartbeatPayload struct {
+	Version     string `json:"version"`
+	ExtVersion  string `json:"frp_version"`
+	RelayStatus string `json:"relay_status"`
+	Name        string `json:"name"`
+	IP          string `json:"ip"`
+}
+
+// Config is the frps configuration sent to the relay.
+type Config struct {
+	BindPort         int    `json:"bind_port"`
+	VhostHTTPPort    int    `json:"vhost_http_port"`
+	AuthToken        string `json:"auth_token"`
+	LogLevel         string `json:"log_level"`
+	WebServerEnabled bool   `json:"web_server_enabled"`
+}
+
+// Settings contains runtime settings for relay and flared clients.
+type Settings 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"`
+}
+
+// HeartbeatResponse is returned from a relay heartbeat.
+type HeartbeatResponse struct {
+	RelayConfig   *Config   `json:"relay_config"`
+	RelaySettings *Settings `json:"relay_settings"`
+}
+
+// Heartbeat processes a relay heartbeat, updates node status, and returns config.
+func Heartbeat(ctx context.Context, node *model.OpenFlareNode, payload HeartbeatPayload) (*HeartbeatResponse, error) {
+	if node == nil {
+		return nil, fmt.Errorf("relay node is nil")
+	}
+
+	payload.Version = strings.TrimSpace(payload.Version)
+	payload.ExtVersion = strings.TrimSpace(payload.ExtVersion)
+	payload.RelayStatus = normalizeRelayStatus(payload.RelayStatus)
+	payload.Name = strings.TrimSpace(payload.Name)
+	payload.IP = strings.TrimSpace(payload.IP)
+
+	previous := *node
+	updateNow := node.UpdateRequested
+	updateChannel := normalizeReleaseChannel(node.UpdateChannel)
+	updateTag := strings.TrimSpace(node.UpdateTag)
+
+	now := time.Now().UTC()
+	changes := map[string]any{
+		"version":          payload.Version,
+		"ext_version":      payload.ExtVersion,
+		"relay_status":     payload.RelayStatus,
+		"last_seen_at":     now,
+		"status":           nodeStatusOnline,
+		"update_requested": false,
+		"update_channel":   "stable",
+		"update_tag":       "",
+	}
+	if payload.Name != "" && strings.TrimSpace(node.Name) == "" {
+		changes["name"] = payload.Name
+		node.Name = payload.Name
+	}
+	if payload.IP != "" && !node.IPManualOverride {
+		changes["ip"] = payload.IP
+		node.IP = payload.IP
+	}
+	if !previous.UpdateRequested {
+		delete(changes, "update_requested")
+	}
+	if previous.UpdateChannel == "stable" {
+		delete(changes, "update_channel")
+	}
+	if previous.UpdateTag == "" {
+		delete(changes, "update_tag")
+	}
+
+	node.Version = payload.Version
+	node.ExtVersion = payload.ExtVersion
+	node.RelayStatus = payload.RelayStatus
+	node.UpdateRequested = false
+	node.UpdateChannel = "stable"
+	node.UpdateTag = ""
+	lastSeen := now
+	node.LastSeenAt = &lastSeen
+	node.Status = nodeStatusOnline
+
+	if err := db.DB(ctx).Model(node).Updates(changes).Error; err != nil {
+		return nil, fmt.Errorf("update relay heartbeat: %w", err)
+	}
+
+	return &HeartbeatResponse{
+		RelayConfig:   buildRelayConfig(node),
+		RelaySettings: BuildSettings(node, updateNow, updateChannel, updateTag),
+	}, nil
+}

Wavelet/internal/apps/openflare/relay/middleware.go

@@ -0,0 +1,51 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+import (
+	"context"
+	"errors"
+	"strings"
+
+	"github.com/Rain-kl/Wavelet/internal/apps/openflare/compat"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
+)
+
+const ctxRelayNodeKey = "relay_node"
+
+// RelayAuth authenticates relay requests using X-Agent-Token and verifies tunnel_relay type.
+func RelayAuth() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		token := strings.TrimSpace(c.GetHeader("X-Agent-Token"))
+		node, err := authenticateAccessToken(c.Request.Context(), token)
+		if err != nil {
+			compat.Unauthorized(c, errAgentTokenInvalid)
+			c.Abort()
+			return
+		}
+		if node.NodeType != "tunnel_relay" {
+			compat.Forbidden(c, errRelayNodeTypeMismatch)
+			c.Abort()
+			return
+		}
+		c.Set(ctxRelayNodeKey, node)
+		c.Next()
+	}
+}
+
+func authenticateAccessToken(ctx context.Context, token string) (*model.OpenFlareNode, error) {
+	if token == "" {
+		return nil, errors.New("missing agent token")
+	}
+	node, err := model.GetOpenFlareNodeByAccessToken(ctx, token)
+	if err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, errors.New("invalid agent token")
+		}
+		return nil, err
+	}
+	return node, nil
+}

Wavelet/internal/apps/openflare/relay/middleware_test.go

@@ -0,0 +1,106 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/Rain-kl/Wavelet/internal/db"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/gin-gonic/gin"
+	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gorm.io/gorm"
+)
+
+func setupRelayMiddlewareTestDB(t *testing.T) func() {
+	t.Helper()
+
+	sqliteDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{
+		DisableForeignKeyConstraintWhenMigrating: true,
+	})
+	require.NoError(t, err)
+	require.NoError(t, sqliteDB.AutoMigrate(&model.OpenFlareNode{}))
+	db.SetDB(sqliteDB)
+
+	return func() {
+		db.SetDB(nil)
+	}
+}
+
+func seedRelayNode(t *testing.T, nodeType, accessToken string) *model.OpenFlareNode {
+	t.Helper()
+	ctx := context.Background()
+	node := &model.OpenFlareNode{
+		NodeID:      "relay-test-node",
+		Name:        "relay-test",
+		Status:      "pending",
+		NodeType:    nodeType,
+		AccessToken: accessToken,
+	}
+	require.NoError(t, model.CreateOpenFlareNode(ctx, node))
+	return node
+}
+
+func TestRelayAuthMissingToken(t *testing.T) {
+	cleanup := setupRelayMiddlewareTestDB(t)
+	defer cleanup()
+
+	gin.SetMode(gin.TestMode)
+	engine := gin.New()
+	engine.GET("/relay/test", RelayAuth(), func(c *gin.Context) {
+		c.Status(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/relay/test", nil)
+	rec := httptest.NewRecorder()
+	engine.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusUnauthorized, rec.Code)
+}
+
+func TestRelayAuthRejectsWrongNodeType(t *testing.T) {
+	cleanup := setupRelayMiddlewareTestDB(t)
+	defer cleanup()
+	seedRelayNode(t, "edge_node", "edge-token-relay")
+
+	gin.SetMode(gin.TestMode)
+	engine := gin.New()
+	engine.GET("/relay/test", RelayAuth(), func(c *gin.Context) {
+		c.Status(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/relay/test", nil)
+	req.Header.Set("X-Agent-Token", "edge-token-relay")
+	rec := httptest.NewRecorder()
+	engine.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusForbidden, rec.Code)
+}
+
+func TestRelayAuthAcceptsTunnelRelay(t *testing.T) {
+	cleanup := setupRelayMiddlewareTestDB(t)
+	defer cleanup()
+	node := seedRelayNode(t, "tunnel_relay", "relay-token-valid")
+
+	gin.SetMode(gin.TestMode)
+	engine := gin.New()
+	engine.GET("/relay/test", RelayAuth(), func(c *gin.Context) {
+		authNode, ok := c.Get(ctxRelayNodeKey)
+		require.True(t, ok)
+		assert.Equal(t, node.NodeID, authNode.(*model.OpenFlareNode).NodeID)
+		c.Status(http.StatusOK)
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/relay/test", nil)
+	req.Header.Set("X-Agent-Token", "relay-token-valid")
+	rec := httptest.NewRecorder()
+	engine.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+}

Wavelet/internal/apps/openflare/relay/routers.go

@@ -0,0 +1,45 @@
+// Copyright 2026 Arctel.net
+// SPDX-License-Identifier: Apache-2.0
+
+package relay
+
+import (
+	"github.com/Rain-kl/Wavelet/internal/apps/openflare/compat"
+	ofws "github.com/Rain-kl/Wavelet/internal/apps/openflare/websocket"
+	"github.com/Rain-kl/Wavelet/internal/model"
+	"github.com/gin-gonic/gin"
+)
+
+// PostHeartbeat handles POST /relay/heartbeat.
+func PostHeartbeat(c *gin.Context) {
+	var payload HeartbeatPayload
+	if !compat.BindJSON(c, &payload) {
+		return
+	}
+	payload.IP = resolveReportedNodeIP(payload.IP, c.Request.RemoteAddr)
+
+	authNode, ok := c.Get(ctxRelayNodeKey)
+	if !ok {
+		compat.Unauthorized(c, errAgentTokenInvalid)
+		return
+	}
+	node := authNode.(*model.OpenFlareNode)
+
+	result, err := Heartbeat(c.Request.Context(), node, payload)
+	if err != nil {
+		compat.Fail(c, err.Error())
+		return
+	}
+	compat.OK(c, result)
+}
+
+// GetWebSocket handles GET /relay/ws.
+func GetWebSocket(c *gin.Context) {
+	authNode, ok := c.Get(ctxRelayNodeKey)
+	if !ok {
+		compat.Unauthorized(c, errAgentTokenInvalid)
+		return
+	}
+	node := authNode.(*model.OpenFlareNode)
+	ofws.ServeRelay(c, node.NodeID)
+}