mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-11 01:36:37 +08:00
refactor(backend): rename OpenFlare directory to lowercase openflare
This commit is contained in:
ryan
@@ -0,0 +1,159 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package wsclient provides WebSocket client abstractions for edge node communication.
|
||||
package wsclient
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
pkgprotocol "Wavelet/openflare/share/protocol"
|
||||
shared "Wavelet/openflare/share/wsclient"
|
||||
)
|
||||
|
||||
// WSMessage is an alias for the shared WebSocket message type.
|
||||
type WSMessage = shared.WSMessage
|
||||
|
||||
// MessageHandler is an alias for the shared WebSocket message handler type.
|
||||
type MessageHandler = shared.MessageHandler
|
||||
|
||||
// Preset represents a predefined connection configuration for a specific edge role.
|
||||
type Preset int
|
||||
|
||||
const (
|
||||
// PresetAgent is the configuration preset for agent connections.
|
||||
PresetAgent Preset = iota
|
||||
// PresetRelay is the configuration preset for relay connections.
|
||||
PresetRelay
|
||||
// PresetFlared is the configuration preset for flared (tunnel) connections.
|
||||
PresetFlared
|
||||
)
|
||||
|
||||
type presetConfig struct {
|
||||
HeaderKey string
|
||||
WSPath string
|
||||
}
|
||||
|
||||
var presets = map[Preset]presetConfig{
|
||||
PresetAgent: {HeaderKey: "X-Agent-Token", WSPath: "/api/v1/agent/ws"},
|
||||
PresetRelay: {HeaderKey: "X-Agent-Token", WSPath: "/api/v1/relay/ws"},
|
||||
PresetFlared: {HeaderKey: "X-Tunnel-Token", WSPath: "/api/v1/tunnel/ws"},
|
||||
}
|
||||
|
||||
// PresetHeaderKey returns the HTTP header key used for authentication with the given preset.
|
||||
func PresetHeaderKey(preset Preset) string {
|
||||
return presets[preset].HeaderKey
|
||||
}
|
||||
|
||||
// PresetWSPath returns the WebSocket path used for the given preset.
|
||||
func PresetWSPath(preset Preset) string {
|
||||
return presets[preset].WSPath
|
||||
}
|
||||
|
||||
// Client is a WebSocket client configured for a specific edge preset.
|
||||
type Client struct {
|
||||
sharedClient *shared.Client
|
||||
}
|
||||
|
||||
// New creates a new Client for the given preset, base URL, token, and timeout.
|
||||
func New(preset Preset, baseURL, token string, timeout time.Duration) *Client {
|
||||
cfg := presets[preset]
|
||||
return &Client{
|
||||
sharedClient: shared.New(shared.Config{
|
||||
BaseURL: baseURL,
|
||||
Token: token,
|
||||
Timeout: timeout,
|
||||
HeaderKey: cfg.HeaderKey,
|
||||
WSPath: cfg.WSPath,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
// SetToken updates the authentication token used by the client.
|
||||
func (c *Client) SetToken(token string) {
|
||||
c.sharedClient.SetToken(token)
|
||||
}
|
||||
|
||||
// URL returns the fully resolved WebSocket URL for this client.
|
||||
func (c *Client) URL() string {
|
||||
return c.sharedClient.URL()
|
||||
}
|
||||
|
||||
// Connection represents an established WebSocket connection to an edge node.
|
||||
type Connection struct {
|
||||
sharedConn *shared.Connection
|
||||
}
|
||||
|
||||
// Connect establishes a WebSocket connection using the client configuration.
|
||||
func (c *Client) Connect(ctx context.Context) (*Connection, error) {
|
||||
conn, err := c.sharedClient.Connect(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Connection{sharedConn: conn}, nil
|
||||
}
|
||||
|
||||
// AgentConnection is a Connection specialized for agent node communication.
|
||||
type AgentConnection struct {
|
||||
Connection
|
||||
}
|
||||
|
||||
// ConnectAgent establishes a WebSocket connection and returns it as an AgentConnection.
|
||||
func (c *Client) ConnectAgent(ctx context.Context) (*AgentConnection, error) {
|
||||
conn, err := c.Connect(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &AgentConnection{Connection: *conn}, nil
|
||||
}
|
||||
|
||||
// URL returns the resolved WebSocket URL of this connection.
|
||||
func (conn *Connection) URL() string {
|
||||
if conn == nil || conn.sharedConn == nil {
|
||||
return ""
|
||||
}
|
||||
return conn.sharedConn.URL
|
||||
}
|
||||
|
||||
// SendPing sends a ping message over the connection.
|
||||
func (conn *Connection) SendPing() error {
|
||||
return conn.sharedConn.SendMessage(pkgprotocol.WSMessageTypePing, nil)
|
||||
}
|
||||
|
||||
// SendPong sends a pong message over the connection.
|
||||
func (conn *Connection) SendPong() error {
|
||||
return conn.sharedConn.SendMessage(pkgprotocol.WSMessageTypePong, nil)
|
||||
}
|
||||
|
||||
// SendMessage sends a typed message with an optional payload over the connection.
|
||||
func (conn *Connection) SendMessage(msgType string, payload any) error {
|
||||
return conn.sharedConn.SendMessage(msgType, payload)
|
||||
}
|
||||
|
||||
// Receive reads the next message from the connection.
|
||||
func (conn *Connection) Receive() (pkgprotocol.WSMessage, error) {
|
||||
var message pkgprotocol.WSMessage
|
||||
if err := conn.sharedConn.Receive(&message); err != nil {
|
||||
return message, err
|
||||
}
|
||||
return message, nil
|
||||
}
|
||||
|
||||
// RunReceiveLoop blocks and dispatches incoming messages to the handler until the context is canceled.
|
||||
func (conn *Connection) RunReceiveLoop(ctx context.Context, handler MessageHandler) error {
|
||||
return conn.sharedConn.RunReceiveLoop(ctx, handler)
|
||||
}
|
||||
|
||||
// Close gracefully closes the WebSocket connection.
|
||||
func (conn *Connection) Close() error {
|
||||
if conn == nil || conn.sharedConn == nil {
|
||||
return nil
|
||||
}
|
||||
return conn.sharedConn.Close()
|
||||
}
|
||||
|
||||
// SendStatus sends a node status payload over the agent connection.
|
||||
func (conn *AgentConnection) SendStatus(payload pkgprotocol.NodePayload) error {
|
||||
return conn.sharedConn.SendMessage(pkgprotocol.WSMessageTypeStatus, payload)
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package wsclient
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestPresetConfig(t *testing.T) {
|
||||
tests := []struct {
|
||||
preset Preset
|
||||
headerKey string
|
||||
wsPath string
|
||||
}{
|
||||
{preset: PresetAgent, headerKey: "X-Agent-Token", wsPath: "/api/v1/agent/ws"},
|
||||
{preset: PresetRelay, headerKey: "X-Agent-Token", wsPath: "/api/v1/relay/ws"},
|
||||
{preset: PresetFlared, headerKey: "X-Tunnel-Token", wsPath: "/api/v1/tunnel/ws"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.wsPath, func(t *testing.T) {
|
||||
if got := PresetHeaderKey(tt.preset); got != tt.headerKey {
|
||||
t.Fatalf("PresetHeaderKey() = %q, want %q", got, tt.headerKey)
|
||||
}
|
||||
if got := PresetWSPath(tt.preset); got != tt.wsPath {
|
||||
t.Fatalf("PresetWSPath() = %q, want %q", got, tt.wsPath)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientURL(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
preset Preset
|
||||
baseURL string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "agent https",
|
||||
preset: PresetAgent,
|
||||
baseURL: "https://example.com",
|
||||
want: "wss://example.com/api/v1/agent/ws",
|
||||
},
|
||||
{
|
||||
name: "relay http with path prefix",
|
||||
preset: PresetRelay,
|
||||
baseURL: "http://example.com/api",
|
||||
want: "ws://example.com/api/api/v1/relay/ws",
|
||||
},
|
||||
{
|
||||
name: "flared wss",
|
||||
preset: PresetFlared,
|
||||
baseURL: "wss://edge.example.com",
|
||||
want: "wss://edge.example.com/api/v1/tunnel/ws",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := New(tt.preset, tt.baseURL, "token", time.Second)
|
||||
if got := client.URL(); got != tt.want {
|
||||
t.Fatalf("URL() = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user