397 lines
9.4 KiB
Go
397 lines
9.4 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"runtime"
|
|
"sync"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
"golang.org/x/net/websocket"
|
|
)
|
|
|
|
// Client Agent 客户端
|
|
type Client struct {
|
|
schedulerHost string
|
|
apiKey string
|
|
nodeID string
|
|
name string
|
|
region string
|
|
heartbeatInterval time.Duration
|
|
reportInterval time.Duration
|
|
logger *zap.Logger
|
|
httpClient *http.Client
|
|
connMutex sync.RWMutex
|
|
wsConn *websocket.Conn
|
|
stats *NodeStats
|
|
statsMutex sync.RWMutex
|
|
commandChan chan Command
|
|
}
|
|
|
|
// NodeStats 节点统计
|
|
type NodeStats struct {
|
|
Online bool `json:"online"`
|
|
WARPConnected bool `json:"warp_connected"`
|
|
CurrentIP string `json:"current_ip"`
|
|
IPRegion string `json:"ip_region"`
|
|
Connections int `json:"connections"`
|
|
TrafficUsedGB float64 `json:"traffic_used_gb"`
|
|
Unlocks map[string]Unlock `json:"unlocks"`
|
|
CPUUsage float64 `json:"cpu_usage"`
|
|
MemoryUsage float64 `json:"memory_usage"`
|
|
NetworkInMbps float64 `json:"network_in_mbps"`
|
|
NetworkOutMbps float64 `json:"network_out_mbps"`
|
|
LastUpdate time.Time `json:"last_update"`
|
|
}
|
|
|
|
// Unlock 解锁状态
|
|
type Unlock struct {
|
|
Unlocked bool `json:"unlocked"`
|
|
Region string `json:"region"`
|
|
}
|
|
|
|
// Command 调度中心指令
|
|
type Command struct {
|
|
Action string `json:"action"`
|
|
Params map[string]interface{} `json:"params"`
|
|
}
|
|
|
|
// NewClient 创建 Agent 客户端
|
|
func NewClient(
|
|
schedulerHost string,
|
|
apiKey string,
|
|
nodeID string,
|
|
name string,
|
|
region string,
|
|
heartbeatInterval time.Duration,
|
|
reportInterval time.Duration,
|
|
logger *zap.Logger,
|
|
) *Client {
|
|
return &Client{
|
|
schedulerHost: schedulerHost,
|
|
apiKey: apiKey,
|
|
nodeID: nodeID,
|
|
name: name,
|
|
region: region,
|
|
heartbeatInterval: heartbeatInterval,
|
|
reportInterval: reportInterval,
|
|
logger: logger,
|
|
httpClient: &http.Client{
|
|
Timeout: 10 * time.Second,
|
|
},
|
|
stats: &NodeStats{
|
|
Online: true,
|
|
Unlocks: make(map[string]Unlock),
|
|
LastUpdate: time.Now(),
|
|
},
|
|
commandChan: make(chan Command, 10),
|
|
}
|
|
}
|
|
|
|
// StartHeartbeat 启动心跳
|
|
func (c *Client) StartHeartbeat(ctx context.Context) {
|
|
ticker := time.NewTicker(c.heartbeatInterval)
|
|
defer ticker.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
c.sendHeartbeat(ctx)
|
|
}
|
|
}
|
|
}
|
|
|
|
// sendHeartbeat 发送心跳
|
|
func (c *Client) sendHeartbeat(ctx context.Context) {
|
|
c.statsMutex.RLock()
|
|
stats := *c.stats
|
|
c.statsMutex.RUnlock()
|
|
|
|
// 收集系统指标
|
|
stats.CPUUsage = c.getCPUUsage()
|
|
stats.MemoryUsage = c.getMemoryUsage()
|
|
stats.Connections = c.getConnections()
|
|
|
|
payload := map[string]interface{}{
|
|
"node_id": c.nodeID,
|
|
"online": stats.Online,
|
|
"warp_connected": stats.WARPConnected,
|
|
"current_ip": stats.CurrentIP,
|
|
"ip_region": stats.IPRegion,
|
|
"connections": stats.Connections,
|
|
"traffic_used_gb": stats.TrafficUsedGB,
|
|
"unlocks": stats.Unlocks,
|
|
"cpu_usage": stats.CPUUsage,
|
|
"memory_usage": stats.MemoryUsage,
|
|
"network_in_mbps": stats.NetworkInMbps,
|
|
"network_out_mbps": stats.NetworkOutMbps,
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/api/v1/agent/heartbeat", c.schedulerHost)
|
|
|
|
body, err := c.post(ctx, url, payload)
|
|
if err != nil {
|
|
c.logger.Error("发送心跳失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
// 解析响应，获取指令
|
|
var resp struct {
|
|
Message string `json:"message"`
|
|
Commands []Command `json:"commands"`
|
|
}
|
|
|
|
if err := json.Unmarshal(body, &resp); err != nil {
|
|
c.logger.Error("解析心跳响应失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
// 处理指令
|
|
for _, cmd := range resp.Commands {
|
|
c.logger.Info("收到调度中心指令",
|
|
zap.String("action", cmd.Action),
|
|
zap.Any("params", cmd.Params),
|
|
)
|
|
select {
|
|
case c.commandChan <- cmd:
|
|
default:
|
|
c.logger.Warn("指令队列已满")
|
|
}
|
|
}
|
|
}
|
|
|
|
// ReportUnlock 上报解锁状态
|
|
func (c *Client) ReportUnlock(ctx context.Context, unlocks map[string]struct {
|
|
Unlocked bool `json:"unlocked"`
|
|
Region string `json:"region"`
|
|
}) {
|
|
// 更新本地状态
|
|
c.statsMutex.Lock()
|
|
for service, status := range unlocks {
|
|
c.stats.Unlocks[service] = Unlock{
|
|
Unlocked: status.Unlocked,
|
|
Region: status.Region,
|
|
}
|
|
}
|
|
c.statsMutex.Unlock()
|
|
|
|
payload := map[string]interface{}{
|
|
"node_id": c.nodeID,
|
|
"unlocks": unlocks,
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/api/v1/agent/unlock/report", c.schedulerHost)
|
|
|
|
_, err := c.post(ctx, url, payload)
|
|
if err != nil {
|
|
c.logger.Error("上报解锁状态失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
c.logger.Info("解锁状态已上报")
|
|
}
|
|
|
|
// ReportIPChange 上报 IP 变更
|
|
func (c *Client) ReportIPChange(ctx context.Context, oldIP, newIP string, success bool, reason string) {
|
|
payload := map[string]interface{}{
|
|
"node_id": c.nodeID,
|
|
"old_ip": oldIP,
|
|
"new_ip": newIP,
|
|
"success": success,
|
|
"reason": reason,
|
|
}
|
|
|
|
url := fmt.Sprintf("%s/api/v1/agent/ip/change/result", c.schedulerHost)
|
|
|
|
_, err := c.post(ctx, url, payload)
|
|
if err != nil {
|
|
c.logger.Error("上报 IP 变更失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
c.logger.Info("IP 变更已上报",
|
|
zap.String("old_ip", oldIP),
|
|
zap.String("new_ip", newIP),
|
|
)
|
|
}
|
|
|
|
// UpdateStats 更新统计信息
|
|
func (c *Client) UpdateStats(stats *NodeStats) {
|
|
c.statsMutex.Lock()
|
|
defer c.statsMutex.Unlock()
|
|
|
|
if stats.CurrentIP != "" {
|
|
c.stats.CurrentIP = stats.CurrentIP
|
|
}
|
|
if stats.IPRegion != "" {
|
|
c.stats.IPRegion = stats.IPRegion
|
|
}
|
|
if stats.WARPConnected != c.stats.WARPConnected {
|
|
c.stats.WARPConnected = stats.WARPConnected
|
|
}
|
|
c.stats.Connections = stats.Connections
|
|
c.stats.LastUpdate = time.Now()
|
|
}
|
|
|
|
// GetCommands 获取指令通道
|
|
func (c *Client) GetCommands() <-chan Command {
|
|
return c.commandChan
|
|
}
|
|
|
|
// post 发送 POST 请求
|
|
func (c *Client) post(ctx context.Context, url string, payload interface{}) ([]byte, error) {
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-API-Key", c.apiKey)
|
|
|
|
resp, err := c.httpClient.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
return io.ReadAll(resp.Body)
|
|
}
|
|
|
|
// getCPUUsage 获取 CPU 使用率
|
|
func (c *Client) getCPUUsage() float64 {
|
|
// 简化实现，实际应该使用 gopsutil
|
|
return 0.0
|
|
}
|
|
|
|
// getMemoryUsage 获取内存使用率
|
|
func (c *Client) getMemoryUsage() float64 {
|
|
var m runtime.MemStats
|
|
runtime.ReadMemStats(&m)
|
|
return float64(m.Sys) / 1024 / 1024
|
|
}
|
|
|
|
// getConnections 获取连接数
|
|
func (c *Client) getConnections() int {
|
|
// TODO: 从 SOCKS5 服务器获取实际连接数
|
|
return 0
|
|
}
|
|
|
|
// CommandHandler 指令处理器
|
|
type CommandHandler struct {
|
|
client *Client
|
|
warpClient WarpClient
|
|
logger *zap.Logger
|
|
}
|
|
|
|
// WarpClient WARP 客户端接口
|
|
type WarpClient interface {
|
|
RefreshIP(ctx context.Context) (string, error)
|
|
GetCurrentIP(ctx context.Context) (string, error)
|
|
Connect(ctx context.Context) error
|
|
Disconnect(ctx context.Context) error
|
|
}
|
|
|
|
// NewCommandHandler 创建指令处理器
|
|
func NewCommandHandler(client *Client, warpClient WarpClient, logger *zap.Logger) *CommandHandler {
|
|
return &CommandHandler{
|
|
client: client,
|
|
warpClient: warpClient,
|
|
logger: logger,
|
|
}
|
|
}
|
|
|
|
// Start 启动指令处理
|
|
func (h *CommandHandler) Start(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case cmd := <-h.client.GetCommands():
|
|
h.handleCommand(ctx, cmd)
|
|
}
|
|
}
|
|
}
|
|
|
|
// handleCommand 处理指令
|
|
func (h *CommandHandler) handleCommand(ctx context.Context, cmd Command) {
|
|
switch cmd.Action {
|
|
case "refresh_ip":
|
|
h.handleRefreshIP(ctx, cmd)
|
|
case "connect_warp":
|
|
h.handleConnectWARP(ctx, cmd)
|
|
case "disconnect_warp":
|
|
h.handleDisconnectWARP(ctx, cmd)
|
|
default:
|
|
h.logger.Warn("未知指令", zap.String("action", cmd.Action))
|
|
}
|
|
}
|
|
|
|
// handleRefreshIP 处理刷新 IP 指令
|
|
func (h *CommandHandler) handleRefreshIP(ctx context.Context, cmd Command) {
|
|
h.logger.Info("执行刷新 IP 指令")
|
|
|
|
reason, _ := cmd.Params["reason"].(string)
|
|
|
|
// 获取旧 IP
|
|
oldIP, _ := h.warpClient.GetCurrentIP(ctx)
|
|
|
|
// 刷新 IP
|
|
newIP, err := h.warpClient.RefreshIP(ctx)
|
|
if err != nil {
|
|
h.logger.Error("刷新 IP 失败", zap.Error(err))
|
|
h.client.ReportIPChange(ctx, oldIP, "", false, reason)
|
|
return
|
|
}
|
|
|
|
// 上报结果
|
|
h.client.ReportIPChange(ctx, oldIP, newIP, true, reason)
|
|
|
|
// 更新本地状态
|
|
h.client.UpdateStats(&NodeStats{
|
|
CurrentIP: newIP,
|
|
WARPConnected: true,
|
|
})
|
|
}
|
|
|
|
// handleConnectWARP 处理连接 WARP 指令
|
|
func (h *CommandHandler) handleConnectWARP(ctx context.Context, cmd Command) {
|
|
h.logger.Info("执行连接 WARP 指令")
|
|
|
|
if err := h.warpClient.Connect(ctx); err != nil {
|
|
h.logger.Error("连接 WARP 失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
ip, _ := h.warpClient.GetCurrentIP(ctx)
|
|
h.client.UpdateStats(&NodeStats{
|
|
WARPConnected: true,
|
|
CurrentIP: ip,
|
|
})
|
|
}
|
|
|
|
// handleDisconnectWARP 处理断开 WARP 指令
|
|
func (h *CommandHandler) handleDisconnectWARP(ctx context.Context, cmd Command) {
|
|
h.logger.Info("执行断开 WARP 指令")
|
|
|
|
if err := h.warpClient.Disconnect(ctx); err != nil {
|
|
h.logger.Error("断开 WARP 失败", zap.Error(err))
|
|
return
|
|
}
|
|
|
|
h.client.UpdateStats(&NodeStats{
|
|
WARPConnected: false,
|
|
CurrentIP: "",
|
|
})
|
|
} |