mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-08 00:26:37 +08:00
refactor(backend): rename OpenFlare directory to lowercase openflare
This commit is contained in:
ryan
@@ -0,0 +1,211 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package websocket manages persistent WebSocket connections between the OpenFlare server and its agents.
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
// AgentWSConnectedLastSeenValue is the sentinel last_seen_at value when agent WS is 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 {
|
||||
wsClientCore
|
||||
remoteAddr string
|
||||
onStatus AgentStatusHandler
|
||||
}
|
||||
|
||||
type agentHub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]*agentClient
|
||||
}
|
||||
|
||||
var defaultAgentHub = &agentHub{clients: make(map[string]*agentClient)}
|
||||
|
||||
// ServeAgent handles an upgraded agent websocket connection.
|
||||
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)
|
||||
return
|
||||
}
|
||||
|
||||
client := &agentClient{
|
||||
wsClientCore: wsClientCore{
|
||||
nodeID: nodeID,
|
||||
conn: conn,
|
||||
send: make(chan Message, wsChannelBuf),
|
||||
done: make(chan struct{}),
|
||||
},
|
||||
remoteAddr: c.Request.RemoteAddr,
|
||||
onStatus: onStatus,
|
||||
}
|
||||
defaultAgentHub.register(client)
|
||||
defer defaultAgentHub.unregister(client)
|
||||
|
||||
slog.Debug("agent ws connected", "node_id", nodeID, "remote", client.remoteAddr)
|
||||
|
||||
go client.writePump()
|
||||
client.readPump()
|
||||
}
|
||||
|
||||
func (h *agentHub) register(client *agentClient) {
|
||||
h.mu.Lock()
|
||||
if existing := h.clients[client.nodeID]; existing != nil {
|
||||
existing.close()
|
||||
}
|
||||
h.clients[client.nodeID] = client
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func (h *agentHub) unregister(client *agentClient) {
|
||||
h.mu.Lock()
|
||||
if current := h.clients[client.nodeID]; current == client {
|
||||
delete(h.clients, client.nodeID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
client.close()
|
||||
}
|
||||
|
||||
// IsAgentConnected reports whether an agent websocket is active.
|
||||
func IsAgentConnected(nodeID string) bool {
|
||||
defaultAgentHub.mu.RLock()
|
||||
client := defaultAgentHub.clients[nodeID]
|
||||
defaultAgentHub.mu.RUnlock()
|
||||
if client == nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case <-client.done:
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return broadcastAgent(agentMessageTypeWAFIPGroups, payload)
|
||||
}
|
||||
|
||||
// BroadcastActiveConfig pushes active config metadata to all connected agents.
|
||||
func BroadcastActiveConfig(payload any) int {
|
||||
return broadcastAgent(agentMessageTypeActiveConfig, payload)
|
||||
}
|
||||
|
||||
func broadcastAgent(messageType string, payload any) int {
|
||||
if payload == nil {
|
||||
return 0
|
||||
}
|
||||
message := Message{Type: messageType, Payload: payload}
|
||||
defaultAgentHub.mu.RLock()
|
||||
clients := make([]*agentClient, 0, len(defaultAgentHub.clients))
|
||||
for _, client := range defaultAgentHub.clients {
|
||||
clients = append(clients, client)
|
||||
}
|
||||
defaultAgentHub.mu.RUnlock()
|
||||
|
||||
success := 0
|
||||
for _, client := range clients {
|
||||
if client.enqueue(message) {
|
||||
success++
|
||||
}
|
||||
}
|
||||
return success
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
return client.enqueue(message)
|
||||
}
|
||||
|
||||
func (c *agentClient) readPump() {
|
||||
defer c.close()
|
||||
|
||||
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 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
|
||||
}
|
||||
|
||||
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:
|
||||
slog.Debug("agent ws unsupported message type", "node_id", c.nodeID, "type", inbound.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func agentWSReadTimeout() time.Duration {
|
||||
timeout := wsReadDeadline
|
||||
if timeout < minAgentWSReadTimeout {
|
||||
return minAgentWSReadTimeout
|
||||
}
|
||||
return timeout
|
||||
}
|
||||
|
||||
func (c *agentClient) writePump() {
|
||||
runWritePump(c.nodeID, c.conn, c.done, c.send, c.close, "agent ws")
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// wsClientCore holds the state and lifecycle shared by all WebSocket client
|
||||
// variants (agent/relay/flared). Embed it; call close exactly-once semantics
|
||||
// are guaranteed via once.
|
||||
type wsClientCore struct {
|
||||
nodeID string
|
||||
conn *websocket.Conn
|
||||
send chan Message
|
||||
done chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// close tears down the connection at most once.
|
||||
func (c *wsClientCore) close() {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.once.Do(func() {
|
||||
close(c.done)
|
||||
if c.conn != nil {
|
||||
_ = c.conn.Close()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// enqueue best-effort delivers message; it never blocks and fails fast when
|
||||
// the client is closed or its send buffer is full.
|
||||
func (c *wsClientCore) enqueue(message Message) bool {
|
||||
// 先确定性检查 closed：若与发送合并在同一个 select，两个 case 同时就绪时
|
||||
// Go 会随机选择，close 后仍可能投递成功。
|
||||
select {
|
||||
case <-c.done:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case c.send <- message:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWSClientCoreCloseIsIdempotent(t *testing.T) {
|
||||
core := &wsClientCore{
|
||||
send: make(chan Message, 1),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
for range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
core.close()
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
select {
|
||||
case <-core.done:
|
||||
default:
|
||||
t.Fatal("close did not signal done")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWSClientCoreEnqueueFailsAfterClose(t *testing.T) {
|
||||
core := &wsClientCore{
|
||||
send: make(chan Message, 1),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
core.close()
|
||||
// 循环多次：若 close 检查与发送合并在同一个 select，两 case 同时就绪时
|
||||
// Go 随机选择，单次调用可能碰巧通过。
|
||||
for range 50 {
|
||||
if core.enqueue(Message{Type: messageTypePing}) {
|
||||
t.Fatal("enqueue must fail after close")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWSClientCoreEnqueueNeverBlocks(t *testing.T) {
|
||||
core := &wsClientCore{
|
||||
send: make(chan Message, 1), // 缓冲小于消息数，验证不阻塞
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
defer core.close()
|
||||
for range 3 {
|
||||
if !core.enqueue(Message{Type: messageTypePing}) && len(core.send) == 0 {
|
||||
t.Fatal("enqueue failed with empty buffer")
|
||||
}
|
||||
}
|
||||
if core.enqueue(Message{Type: messageTypePing}) {
|
||||
t.Fatal("enqueue must fail when buffer full")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
const (
|
||||
messageTypePing = "ping"
|
||||
messageTypePong = "pong"
|
||||
messageTypeNotify = "notify"
|
||||
)
|
||||
|
||||
// Message is a JSON-framed websocket payload.
|
||||
type Message struct {
|
||||
Type string `json:"type"`
|
||||
Payload any `json:"payload,omitempty"`
|
||||
}
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(_ *http.Request) bool { return true },
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import "time"
|
||||
|
||||
const (
|
||||
wsChannelBuf = 16
|
||||
wsPingInterval = 30 * time.Second
|
||||
wsReadDeadline = 90 * time.Second
|
||||
wsWriteDeadline = 10 * time.Second
|
||||
minAgentWSReadTimeout = 30 * time.Second
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
// FlaredWSConnectedLastSeenValue is the sentinel last_seen_at value when flared WS is connected.
|
||||
FlaredWSConnectedLastSeenValue = "__OPENFLARE_FLARED_WS_CONNECTED__"
|
||||
|
||||
flaredMessageTypeActiveConfig = "active_config"
|
||||
flaredMessageTypeForceSync = "force_sync"
|
||||
flaredMessageTypePong = "pong"
|
||||
)
|
||||
|
||||
type flaredClient struct {
|
||||
wsClientCore
|
||||
}
|
||||
|
||||
type flaredHub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]*flaredClient
|
||||
}
|
||||
|
||||
var defaultFlaredHub = &flaredHub{clients: make(map[string]*flaredClient)}
|
||||
|
||||
// ServeFlared handles an upgraded flared websocket connection.
|
||||
func ServeFlared(c *gin.Context, nodeID string) {
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
slog.Debug("flared ws upgrade failed", "node_id", nodeID, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
client := &flaredClient{
|
||||
wsClientCore: wsClientCore{
|
||||
nodeID: nodeID,
|
||||
conn: conn,
|
||||
send: make(chan Message, wsChannelBuf),
|
||||
done: make(chan struct{}),
|
||||
},
|
||||
}
|
||||
defaultFlaredHub.register(client)
|
||||
defer defaultFlaredHub.unregister(client)
|
||||
|
||||
slog.Debug("flared ws connected", "node_id", nodeID, "remote", c.Request.RemoteAddr)
|
||||
|
||||
go client.writePump()
|
||||
client.readPump()
|
||||
}
|
||||
|
||||
func (h *flaredHub) register(client *flaredClient) {
|
||||
h.mu.Lock()
|
||||
if existing := h.clients[client.nodeID]; existing != nil {
|
||||
existing.close()
|
||||
}
|
||||
h.clients[client.nodeID] = client
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func (h *flaredHub) unregister(client *flaredClient) {
|
||||
h.mu.Lock()
|
||||
if current := h.clients[client.nodeID]; current == client {
|
||||
delete(h.clients, client.nodeID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
client.close()
|
||||
}
|
||||
|
||||
// DisconnectFlaredClient forcefully disconnects a flared websocket client.
|
||||
func DisconnectFlaredClient(nodeID string) {
|
||||
defaultFlaredHub.mu.Lock()
|
||||
client := defaultFlaredHub.clients[nodeID]
|
||||
if client != nil {
|
||||
delete(defaultFlaredHub.clients, nodeID)
|
||||
}
|
||||
defaultFlaredHub.mu.Unlock()
|
||||
if client != nil {
|
||||
client.close()
|
||||
}
|
||||
}
|
||||
|
||||
// IsFlaredConnected reports whether a flared websocket is active.
|
||||
func IsFlaredConnected(nodeID string) bool {
|
||||
defaultFlaredHub.mu.RLock()
|
||||
client := defaultFlaredHub.clients[nodeID]
|
||||
defaultFlaredHub.mu.RUnlock()
|
||||
if client == nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case <-client.done:
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// SendFlaredPong enqueues a pong message for the flared node.
|
||||
func SendFlaredPong(nodeID string) bool {
|
||||
defaultFlaredHub.mu.RLock()
|
||||
client := defaultFlaredHub.clients[nodeID]
|
||||
defaultFlaredHub.mu.RUnlock()
|
||||
if client == nil {
|
||||
return false
|
||||
}
|
||||
// 委托 enqueue：closed 检查与发送不能合并在同一个 select（两 case 同时
|
||||
// 就绪时 Go 随机选择，close 后仍可能投递成功）。
|
||||
return client.enqueue(Message{Type: flaredMessageTypePong})
|
||||
}
|
||||
|
||||
func (c *flaredClient) readPump() {
|
||||
runReadPump(c.nodeID, c.conn, c.close, "flared ws", SendFlaredPong, flaredMessageTypePong)
|
||||
}
|
||||
|
||||
func (c *flaredClient) writePump() {
|
||||
runWritePump(c.nodeID, c.conn, c.done, c.send, c.close, "flared ws")
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
func runReadPump(
|
||||
nodeID string,
|
||||
conn *websocket.Conn,
|
||||
closeFn func(),
|
||||
logLabel string,
|
||||
sendPong func(string) bool,
|
||||
clientPongType string,
|
||||
) {
|
||||
defer closeFn()
|
||||
_ = conn.SetReadDeadline(time.Now().Add(wsReadDeadline))
|
||||
conn.SetPongHandler(func(string) error {
|
||||
return conn.SetReadDeadline(time.Now().Add(wsReadDeadline))
|
||||
})
|
||||
|
||||
for {
|
||||
_, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
slog.Debug(logLabel+" read closed", "node_id", nodeID, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
var message Message
|
||||
if err = json.Unmarshal(data, &message); err != nil {
|
||||
slog.Debug(logLabel+" invalid message", "node_id", nodeID, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
switch message.Type {
|
||||
case messageTypePing:
|
||||
_ = sendPong(nodeID)
|
||||
case clientPongType:
|
||||
// Refresh read deadline when the client replies with a JSON pong.
|
||||
// This keeps the connection alive when the WebSocket is proxied
|
||||
// through Cloudflare, which enforces a 100-second idle timeout on
|
||||
// the TCP stream. Without this refresh, the server's 90-second read
|
||||
// deadline expires and terminates the connection even though the
|
||||
// client is actively responding to pings.
|
||||
_ = conn.SetReadDeadline(time.Now().Add(wsReadDeadline))
|
||||
default:
|
||||
slog.Debug(logLabel+" unsupported message", "node_id", nodeID, "type", message.Type)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// RelayWSConnectedLastSeenValue is the sentinel last_seen_at value when relay WS is connected.
|
||||
const RelayWSConnectedLastSeenValue = "__OPENFLARE_WS_CONNECTED__"
|
||||
|
||||
type relayClient struct {
|
||||
wsClientCore
|
||||
}
|
||||
|
||||
type relayHub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]*relayClient
|
||||
}
|
||||
|
||||
var defaultRelayHub = &relayHub{clients: make(map[string]*relayClient)}
|
||||
|
||||
// ServeRelay handles an upgraded relay websocket connection.
|
||||
func ServeRelay(c *gin.Context, nodeID string) {
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
slog.Debug("relay ws upgrade failed", "node_id", nodeID, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
client := &relayClient{
|
||||
wsClientCore: wsClientCore{
|
||||
nodeID: nodeID,
|
||||
conn: conn,
|
||||
send: make(chan Message, wsChannelBuf),
|
||||
done: make(chan struct{}),
|
||||
},
|
||||
}
|
||||
defaultRelayHub.register(client)
|
||||
defer defaultRelayHub.unregister(client)
|
||||
|
||||
slog.Debug("relay ws connected", "node_id", nodeID, "remote", c.Request.RemoteAddr)
|
||||
|
||||
go client.writePump()
|
||||
client.readPump()
|
||||
}
|
||||
|
||||
func (h *relayHub) register(client *relayClient) {
|
||||
h.mu.Lock()
|
||||
if existing := h.clients[client.nodeID]; existing != nil {
|
||||
existing.close()
|
||||
}
|
||||
h.clients[client.nodeID] = client
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func (h *relayHub) unregister(client *relayClient) {
|
||||
h.mu.Lock()
|
||||
if current := h.clients[client.nodeID]; current == client {
|
||||
delete(h.clients, client.nodeID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
client.close()
|
||||
}
|
||||
|
||||
// IsRelayConnected reports whether a relay websocket is active.
|
||||
func IsRelayConnected(nodeID string) bool {
|
||||
defaultRelayHub.mu.RLock()
|
||||
client := defaultRelayHub.clients[nodeID]
|
||||
defaultRelayHub.mu.RUnlock()
|
||||
if client == nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case <-client.done:
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// SendRelayPong enqueues a pong message for the relay node.
|
||||
func SendRelayPong(nodeID string) bool {
|
||||
defaultRelayHub.mu.RLock()
|
||||
client := defaultRelayHub.clients[nodeID]
|
||||
defaultRelayHub.mu.RUnlock()
|
||||
if client == nil {
|
||||
return false
|
||||
}
|
||||
// 委托 enqueue：closed 检查与发送不能合并在同一个 select（两 case 同时
|
||||
// 就绪时 Go 随机选择，close 后仍可能投递成功）。
|
||||
return client.enqueue(Message{Type: messageTypePong})
|
||||
}
|
||||
|
||||
func (c *relayClient) readPump() {
|
||||
runReadPump(c.nodeID, c.conn, c.close, "relay ws", SendRelayPong, messageTypePong)
|
||||
}
|
||||
|
||||
func (c *relayClient) writePump() {
|
||||
runWritePump(c.nodeID, c.conn, c.done, c.send, c.close, "relay ws")
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// runWritePump drains send onto conn until done is closed, emitting
|
||||
// JSON pings at wsPingInterval. Shared by agent/relay/flared clients;
|
||||
// closeFn must be idempotent.
|
||||
func runWritePump(
|
||||
nodeID string,
|
||||
conn *websocket.Conn,
|
||||
done <-chan struct{},
|
||||
send chan Message,
|
||||
closeFn func(),
|
||||
logLabel string,
|
||||
) {
|
||||
ticker := time.NewTicker(wsPingInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case message := <-send:
|
||||
_ = conn.SetWriteDeadline(time.Now().Add(wsWriteDeadline))
|
||||
if err := conn.WriteJSON(message); err != nil {
|
||||
slog.Debug(logLabel+" write failed", "node_id", nodeID, "error", err)
|
||||
closeFn()
|
||||
return
|
||||
}
|
||||
case <-ticker.C:
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case send <- Message{Type: messageTypePing}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user