This commit is contained in:
ryan
2026-06-18 17:08:21 +08:00
parent dfc480c73a
commit 01c28bb88b
11 changed files with 494 additions and 129 deletions
@@ -167,7 +167,7 @@ func AgentWebSocketHandler(c *gin.Context) {
compat.Unauthorized(c, errInvalidAgentToken)
return
}
websocket.ServeAgent(c, authNode.NodeID)
websocket.ServeAgent(c, authNode.NodeID, HandleWSStatus)
}
func okWithExtras(c *gin.Context, data any, extras gin.H) {
@@ -0,0 +1,57 @@
// Copyright 2026 Arctel.net
// SPDX-License-Identifier: Apache-2.0
package agent
import (
"context"
"encoding/json"
"log/slog"
ofws "github.com/Rain-kl/Wavelet/internal/apps/openflare/websocket"
"github.com/Rain-kl/Wavelet/internal/model"
)
// HandleWSStatus processes an agent websocket status payload (replaces HTTP heartbeat in WS mode).
func HandleWSStatus(ctx context.Context, nodeID, remoteAddr string, rawPayload json.RawMessage) {
var payload NodePayload
if err := json.Unmarshal(rawPayload, &payload); err != nil {
slog.Debug("agent ws status payload decode failed", "node_id", nodeID, "error", err)
return
}
authNode, err := model.GetOpenFlareNodeByNodeID(ctx, nodeID)
if err != nil {
slog.Debug("agent ws status reload node failed", "node_id", nodeID, "error", err)
return
}
payload.IP = resolveReportedNodeIP(payload.IP, remoteAddr)
response, err := HeartbeatNode(ctx, authNode, payload)
if err != nil {
slog.Debug("agent ws status handling failed", "node_id", nodeID, "error", err)
return
}
settingsSent := false
if response.AgentSettings != nil {
settingsSent = ofws.SendAgentSettings(nodeID, response.AgentSettings)
}
activeConfigSent := false
if response.ActiveConfig != nil {
activeConfigSent = ofws.SendAgentActiveConfig(nodeID, response.ActiveConfig)
}
wafIPGroupsSent := false
if len(response.WAFIPGroups) > 0 {
wafIPGroupsSent = ofws.SendAgentWAFIPGroups(nodeID, response.WAFIPGroups)
}
slog.Debug("agent ws status processed",
"node_id", nodeID,
"current_version", payload.CurrentVersion,
"openresty_status", payload.OpenrestyStatus,
"settings_sent", settingsSent,
"active_config_sent", activeConfigSent,
"waf_ip_groups_sent", wafIPGroupsSent,
)
}
@@ -4,8 +4,10 @@
package dashboard
import (
"strings"
"time"
ofws "github.com/Rain-kl/Wavelet/internal/apps/openflare/websocket"
"github.com/Rain-kl/Wavelet/internal/model"
)
@@ -29,7 +31,23 @@ func computeNodeStatus(node *model.OpenFlareNode) string {
}
func nodeViewLastSeenAt(node *model.OpenFlareNode) any {
if node == nil || node.LastSeenAt == nil {
if node == nil {
return time.Time{}
}
nodeType := strings.TrimSpace(node.NodeType)
if nodeType == "" {
nodeType = "edge_node"
}
if nodeType == "tunnel_relay" && ofws.IsRelayConnected(node.NodeID) {
return ofws.RelayWSConnectedLastSeenValue
}
if nodeType == "tunnel_client" && ofws.IsFlaredConnected(node.NodeID) {
return ofws.FlaredWSConnectedLastSeenValue
}
if ofws.IsAgentConnected(node.NodeID) {
return ofws.AgentWSConnectedLastSeenValue
}
if node.LastSeenAt == nil {
return time.Time{}
}
return *node.LastSeenAt
@@ -16,6 +16,7 @@ import (
"strings"
"time"
ofws "github.com/Rain-kl/Wavelet/internal/apps/openflare/websocket"
"github.com/Rain-kl/Wavelet/internal/model"
)
@@ -183,7 +184,23 @@ func computeNodeStatus(node *model.OpenFlareNode) string {
}
func nodeViewLastSeenAt(node *model.OpenFlareNode) any {
if node == nil || node.LastSeenAt == nil {
if node == nil {
return time.Time{}
}
nodeType := strings.TrimSpace(node.NodeType)
if nodeType == "" {
nodeType = "edge_node"
}
if nodeType == "tunnel_relay" && ofws.IsRelayConnected(node.NodeID) {
return ofws.RelayWSConnectedLastSeenValue
}
if nodeType == "tunnel_client" && ofws.IsFlaredConnected(node.NodeID) {
return ofws.FlaredWSConnectedLastSeenValue
}
if ofws.IsAgentConnected(node.NodeID) {
return ofws.AgentWSConnectedLastSeenValue
}
if node.LastSeenAt == nil {
return time.Time{}
}
return *node.LastSeenAt
@@ -4,6 +4,7 @@
package websocket
import (
"context"
"encoding/json"
"log/slog"
"sync"
@@ -15,18 +16,26 @@ import (
const (
// AgentWSConnectedLastSeenValue is the sentinel last_seen_at value when agent WS is connected.
AgentWSConnectedLastSeenValue = "__OPENFLARE_AGENT_WS_CONNECTED__"
AgentWSConnectedLastSeenValue = "__OPENFLARE_WS_CONNECTED__"
agentMessageTypeStatus = "status"
agentMessageTypeSettings = "settings"
agentMessageTypeActiveConfig = "active_config"
agentMessageTypeForceSyncConfig = "force_sync_config"
agentMessageTypeWAFIPGroups = "waf_ip_groups"
)
// AgentStatusHandler processes inbound agent websocket status payloads.
type AgentStatusHandler func(ctx context.Context, nodeID, remoteAddr string, payload json.RawMessage)
type agentClient struct {
nodeID string
conn *websocket.Conn
send chan Message
done chan struct{}
once sync.Once
nodeID string
remoteAddr string
conn *websocket.Conn
send chan Message
done chan struct{}
onStatus AgentStatusHandler
once sync.Once
}
func (c *agentClient) close() {
@@ -47,7 +56,7 @@ type agentHub struct {
var defaultAgentHub = &agentHub{clients: make(map[string]*agentClient)}
// ServeAgent handles an upgraded agent websocket connection.
func ServeAgent(c *gin.Context, nodeID string) {
func ServeAgent(c *gin.Context, nodeID string, onStatus AgentStatusHandler) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
slog.Debug("agent ws upgrade failed", "node_id", nodeID, "error", err)
@@ -55,15 +64,17 @@ func ServeAgent(c *gin.Context, nodeID string) {
}
client := &agentClient{
nodeID: nodeID,
conn: conn,
send: make(chan Message, 16),
done: make(chan struct{}),
nodeID: nodeID,
remoteAddr: c.Request.RemoteAddr,
conn: conn,
send: make(chan Message, 16),
done: make(chan struct{}),
onStatus: onStatus,
}
defaultAgentHub.register(client)
defer defaultAgentHub.unregister(client)
slog.Debug("agent ws connected", "node_id", nodeID, "remote", c.Request.RemoteAddr)
slog.Debug("agent ws connected", "node_id", nodeID, "remote", client.remoteAddr)
go client.writePump()
client.readPump()
@@ -103,6 +114,21 @@ func IsAgentConnected(nodeID string) bool {
}
}
// SendAgentSettings pushes agent settings to a connected agent.
func SendAgentSettings(nodeID string, payload any) bool {
return sendAgentMessage(nodeID, Message{Type: agentMessageTypeSettings, Payload: payload})
}
// SendAgentActiveConfig pushes active config metadata to a connected agent.
func SendAgentActiveConfig(nodeID string, payload any) bool {
return sendAgentMessage(nodeID, Message{Type: agentMessageTypeActiveConfig, Payload: payload})
}
// SendAgentWAFIPGroups pushes WAF IP group updates to a connected agent.
func SendAgentWAFIPGroups(nodeID string, payload any) bool {
return sendAgentMessage(nodeID, Message{Type: agentMessageTypeWAFIPGroups, Payload: payload})
}
// BroadcastWAFIPGroups pushes changed WAF IP groups to all connected agents.
func BroadcastWAFIPGroups(payload any) int {
if payload == nil {
@@ -127,56 +153,62 @@ func BroadcastWAFIPGroups(payload any) int {
// SendForceSyncConfig notifies an agent to force sync configuration.
func SendForceSyncConfig(nodeID string, payload any) bool {
return sendAgentMessage(nodeID, Message{Type: agentMessageTypeForceSyncConfig, Payload: payload})
}
func sendAgentMessage(nodeID string, message Message) bool {
defaultAgentHub.mu.RLock()
client := defaultAgentHub.clients[nodeID]
defaultAgentHub.mu.RUnlock()
if client == nil {
return false
}
select {
case <-client.done:
return false
case client.send <- Message{Type: agentMessageTypeForceSyncConfig, Payload: payload}:
return true
default:
return false
}
return client.enqueue(message)
}
func (c *agentClient) readPump() {
defer c.close()
_ = c.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
c.conn.SetPongHandler(func(string) error {
return c.conn.SetReadDeadline(time.Now().Add(90 * time.Second))
})
for {
_ = c.conn.SetReadDeadline(time.Now().Add(agentWSReadTimeout()))
_, data, err := c.conn.ReadMessage()
if err != nil {
slog.Debug("agent ws read closed", "node_id", c.nodeID, "error", err)
return
}
var message Message
if err = json.Unmarshal(data, &message); err != nil {
var inbound struct {
Type string `json:"type"`
Payload json.RawMessage `json:"payload,omitempty"`
}
if err = json.Unmarshal(data, &inbound); err != nil {
slog.Debug("agent ws invalid message", "node_id", c.nodeID, "error", err)
continue
}
switch message.Type {
slog.Debug("agent ws message received", "node_id", c.nodeID, "type", inbound.Type)
switch inbound.Type {
case agentMessageTypeStatus:
if c.onStatus != nil {
c.onStatus(context.Background(), c.nodeID, c.remoteAddr, inbound.Payload)
}
case messageTypePing:
_ = c.enqueue(Message{Type: messageTypePong})
case messageTypePong:
default:
_ = c.enqueue(Message{Type: messageTypeNotify, Payload: gin.H{
"echo": true,
"type": message.Type,
"payload": message.Payload,
}})
slog.Debug("agent ws unsupported message type", "node_id", c.nodeID, "type", inbound.Type)
}
}
}
func agentWSReadTimeout() time.Duration {
timeout := 90 * time.Second
if timeout < 30*time.Second {
return 30 * time.Second
}
return timeout
}
func (c *agentClient) writePump() {
ticker := time.NewTicker(30 * time.Second)
defer ticker.Stop()