chore: update dependencies and improve bot configuration

- Upgrade Go version to 1.24.0 and update toolchain.
- Update various dependencies in go.mod and go.sum, including:
  - Upgrade `fasthttp/websocket` to v1.5.12
  - Upgrade `fsnotify/fsnotify` to v1.9.0
  - Upgrade `valyala/fasthttp` to v1.58.0
  - Add new dependencies for `bytedance/sonic` and `google/uuid`.
- Refactor bot configuration in config.toml to support multiple bot protocols, including "milky" and "onebot11".
- Modify internal configuration structures to accommodate new bot settings.
- Enhance event dispatcher with metrics tracking and asynchronous processing capabilities.
- Implement WebSocket connection management with heartbeat and reconnection logic.
- Update server handling for bot management and event publishing.
This commit is contained in:
lafay
2026-01-05 00:40:09 +08:00
parent ac0dfb64c9
commit 44fe05ff62
30 changed files with 6311 additions and 182 deletions

313
pkg/net/httpclient.go Normal file
View File

@@ -0,0 +1,313 @@
package net
import (
"context"
"fmt"
"sync"
"time"
"cellbot/internal/engine"
"cellbot/internal/protocol"
"github.com/bytedance/sonic"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
)
// HTTPClient HTTP客户端用于正向HTTP连接
type HTTPClient struct {
client *fasthttp.Client
logger *zap.Logger
eventBus *engine.EventBus
botID string
baseURL string
timeout time.Duration
retryCount int
}
// HTTPClientConfig HTTP客户端配置
type HTTPClientConfig struct {
BotID string
BaseURL string
Timeout time.Duration
RetryCount int
}
// NewHTTPClient 创建HTTP客户端
func NewHTTPClient(config HTTPClientConfig, logger *zap.Logger, eventBus *engine.EventBus) *HTTPClient {
if config.Timeout == 0 {
config.Timeout = 30 * time.Second
}
if config.RetryCount == 0 {
config.RetryCount = 3
}
return &HTTPClient{
client: &fasthttp.Client{
ReadTimeout: config.Timeout,
WriteTimeout: config.Timeout,
MaxConnsPerHost: 100,
},
logger: logger.Named("http-client"),
eventBus: eventBus,
botID: config.BotID,
baseURL: config.BaseURL,
timeout: config.Timeout,
retryCount: config.RetryCount,
}
}
// SendAction 发送动作请求正向HTTP
func (hc *HTTPClient) SendAction(ctx context.Context, action protocol.Action) (map[string]interface{}, error) {
// 序列化动作为JSON
data, err := sonic.Marshal(action)
if err != nil {
return nil, fmt.Errorf("failed to marshal action: %w", err)
}
req := fasthttp.AcquireRequest()
resp := fasthttp.AcquireResponse()
defer fasthttp.ReleaseRequest(req)
defer fasthttp.ReleaseResponse(resp)
url := hc.baseURL + "/action"
req.SetRequestURI(url)
req.Header.SetMethod("POST")
req.Header.SetContentType("application/json")
req.SetBody(data)
hc.logger.Debug("Sending action",
zap.String("url", url),
zap.String("action", string(action.GetType())))
// 重试机制
var lastErr error
for i := 0; i <= hc.retryCount; i++ {
if i > 0 {
hc.logger.Info("Retrying action request",
zap.Int("attempt", i),
zap.Int("max", hc.retryCount))
time.Sleep(time.Duration(i) * time.Second)
}
err := hc.client.DoTimeout(req, resp, hc.timeout)
if err != nil {
lastErr = fmt.Errorf("request failed: %w", err)
continue
}
if resp.StatusCode() != fasthttp.StatusOK {
lastErr = fmt.Errorf("unexpected status code: %d", resp.StatusCode())
continue
}
// 解析响应
var result map[string]interface{}
if err := sonic.Unmarshal(resp.Body(), &result); err != nil {
lastErr = fmt.Errorf("failed to parse response: %w", err)
continue
}
hc.logger.Info("Action sent successfully",
zap.String("action", string(action.GetType())))
return result, nil
}
return nil, fmt.Errorf("action failed after %d retries: %w", hc.retryCount, lastErr)
}
// PollEvents 轮询事件正向HTTP
func (hc *HTTPClient) PollEvents(ctx context.Context, interval time.Duration) error {
ticker := time.NewTicker(interval)
defer ticker.Stop()
hc.logger.Info("Starting event polling",
zap.Duration("interval", interval))
for {
select {
case <-ticker.C:
if err := hc.fetchEvents(ctx); err != nil {
hc.logger.Error("Failed to fetch events", zap.Error(err))
}
case <-ctx.Done():
hc.logger.Info("Event polling stopped")
return ctx.Err()
}
}
}
// fetchEvents 获取事件
func (hc *HTTPClient) fetchEvents(ctx context.Context) error {
req := fasthttp.AcquireRequest()
resp := fasthttp.AcquireResponse()
defer fasthttp.ReleaseRequest(req)
defer fasthttp.ReleaseResponse(resp)
url := hc.baseURL + "/events"
req.SetRequestURI(url)
req.Header.SetMethod("GET")
err := hc.client.DoTimeout(req, resp, hc.timeout)
if err != nil {
return fmt.Errorf("request failed: %w", err)
}
if resp.StatusCode() != fasthttp.StatusOK {
return fmt.Errorf("unexpected status code: %d", resp.StatusCode())
}
// 解析事件列表
var events []protocol.BaseEvent
if err := sonic.Unmarshal(resp.Body(), &events); err != nil {
return fmt.Errorf("failed to parse events: %w", err)
}
// 发布事件到事件总线
for i := range events {
hc.logger.Debug("Event received",
zap.String("type", string(events[i].Type)),
zap.String("detail_type", events[i].DetailType))
hc.eventBus.Publish(&events[i])
}
return nil
}
// HTTPWebhookServer HTTP Webhook服务器用于反向HTTP连接
type HTTPWebhookServer struct {
server *fasthttp.Server
logger *zap.Logger
eventBus *engine.EventBus
handlers map[string]*WebhookHandler
mu sync.RWMutex
}
// WebhookHandler Webhook处理器
type WebhookHandler struct {
BotID string
Secret string
Validator func([]byte, string) bool
}
// NewHTTPWebhookServer 创建HTTP Webhook服务器
func NewHTTPWebhookServer(logger *zap.Logger, eventBus *engine.EventBus) *HTTPWebhookServer {
return &HTTPWebhookServer{
logger: logger.Named("webhook-server"),
eventBus: eventBus,
handlers: make(map[string]*WebhookHandler),
}
}
// RegisterWebhook 注册Webhook处理器
func (hws *HTTPWebhookServer) RegisterWebhook(path string, handler *WebhookHandler) {
hws.mu.Lock()
defer hws.mu.Unlock()
hws.handlers[path] = handler
hws.logger.Info("Webhook registered",
zap.String("path", path),
zap.String("bot_id", handler.BotID))
}
// UnregisterWebhook 注销Webhook处理器
func (hws *HTTPWebhookServer) UnregisterWebhook(path string) {
hws.mu.Lock()
defer hws.mu.Unlock()
delete(hws.handlers, path)
hws.logger.Info("Webhook unregistered", zap.String("path", path))
}
// Start 启动Webhook服务器
func (hws *HTTPWebhookServer) Start(addr string) error {
hws.server = &fasthttp.Server{
Handler: hws.handleWebhook,
}
hws.logger.Info("Starting webhook server", zap.String("address", addr))
go func() {
if err := hws.server.ListenAndServe(addr); err != nil {
hws.logger.Error("Webhook server error", zap.Error(err))
}
}()
return nil
}
// Stop 停止Webhook服务器
func (hws *HTTPWebhookServer) Stop() error {
if hws.server != nil {
hws.logger.Info("Stopping webhook server")
return hws.server.Shutdown()
}
return nil
}
// handleWebhook 处理Webhook请求
func (hws *HTTPWebhookServer) handleWebhook(ctx *fasthttp.RequestCtx) {
path := string(ctx.Path())
hws.mu.RLock()
handler, exists := hws.handlers[path]
hws.mu.RUnlock()
if !exists {
ctx.Error("Webhook not found", fasthttp.StatusNotFound)
return
}
// 验证签名(如果配置了)
if handler.Secret != "" && handler.Validator != nil {
signature := string(ctx.Request.Header.Peek("X-Signature"))
if !handler.Validator(ctx.PostBody(), signature) {
hws.logger.Warn("Invalid webhook signature",
zap.String("path", path),
zap.String("bot_id", handler.BotID))
ctx.Error("Invalid signature", fasthttp.StatusUnauthorized)
return
}
}
// 解析事件
var event protocol.BaseEvent
if err := sonic.Unmarshal(ctx.PostBody(), &event); err != nil {
hws.logger.Error("Failed to parse webhook event",
zap.Error(err),
zap.String("path", path))
ctx.Error("Invalid event format", fasthttp.StatusBadRequest)
return
}
// 设置BotID
if event.SelfID == "" {
event.SelfID = handler.BotID
}
// 设置时间戳
if event.Timestamp == 0 {
event.Timestamp = time.Now().Unix()
}
// 确保Data字段不为nil
if event.Data == nil {
event.Data = make(map[string]interface{})
}
hws.logger.Info("Webhook event received",
zap.String("path", path),
zap.String("bot_id", handler.BotID),
zap.String("type", string(event.Type)),
zap.String("detail_type", event.DetailType))
// 发布到事件总线
hws.eventBus.Publish(&event)
// 返回成功响应
ctx.SetContentType("application/json")
ctx.SetBodyString(`{"success":true}`)
}

View File

@@ -1,12 +1,13 @@
package net
import (
"fmt"
"net"
"strconv"
"cellbot/internal/engine"
"cellbot/internal/protocol"
"github.com/bytedance/sonic"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
)
@@ -18,17 +19,20 @@ type Server struct {
logger *zap.Logger
botManager *protocol.BotManager
eventBus *engine.EventBus
wsManager *WebSocketManager
server *fasthttp.Server
}
// NewServer 创建HTTP服务器
func NewServer(host string, port int, logger *zap.Logger, botManager *protocol.BotManager, eventBus *engine.EventBus) *Server {
wsManager := NewWebSocketManager(logger, eventBus)
return &Server{
host: host,
port: port,
logger: logger.Named("server"),
botManager: botManager,
eventBus: eventBus,
wsManager: wsManager,
}
}
@@ -105,28 +109,55 @@ func (s *Server) handleHealth(ctx *fasthttp.RequestCtx) {
ctx.SetBodyString(`{"status":"ok"}`)
}
// BotInfo 机器人信息结构
type BotInfo struct {
ID string `json:"id"`
Name string `json:"name"`
Version string `json:"version"`
Status string `json:"status"`
SelfID string `json:"self_id"`
Connected bool `json:"connected"`
}
// handleBots 获取机器人列表
func (s *Server) handleBots(ctx *fasthttp.RequestCtx) {
bots := s.botManager.GetAll()
ctx.SetContentType("application/json")
if len(bots) == 0 {
ctx.SetBodyString(`{"bots":[]}`)
// 构建机器人信息列表
botInfos := make([]BotInfo, 0, len(bots))
for _, bot := range bots {
botInfos = append(botInfos, BotInfo{
ID: bot.GetID(),
Name: bot.Name(),
Version: bot.Version(),
Status: string(bot.GetStatus()),
SelfID: bot.GetSelfID(),
Connected: bot.IsConnected(),
})
}
// 序列化为JSON
response := map[string]interface{}{
"bots": botInfos,
"count": len(botInfos),
}
data, err := sonic.Marshal(response)
if err != nil {
s.logger.Error("Failed to marshal bots response", zap.Error(err))
ctx.Error("Internal Server Error", fasthttp.StatusInternalServerError)
return
}
// 简化实现,实际应该序列化完整信息
response := `{"bots":[`
for i, bot := range bots {
if i > 0 {
response += ","
}
response += fmt.Sprintf(`{"id":"%s","name":"%s","status":"%s"}`,
bot.GetID(), bot.Name(), bot.GetStatus())
}
response += `]}`
ctx.SetBody(data)
}
ctx.SetBodyString(response)
// CreateBotRequest 创建机器人请求
type CreateBotRequest struct {
ID string `json:"id"`
Protocol string `json:"protocol"`
Config map[string]interface{} `json:"config"`
}
// handleCreateBot 创建机器人
@@ -137,7 +168,73 @@ func (s *Server) handleCreateBot(ctx *fasthttp.RequestCtx) {
}
ctx.SetContentType("application/json")
ctx.SetBodyString(`{"message":"Bot creation not implemented yet"}`)
// 解析请求体
var req CreateBotRequest
if err := sonic.Unmarshal(ctx.PostBody(), &req); err != nil {
s.logger.Error("Failed to parse create bot request", zap.Error(err))
response := map[string]interface{}{
"success": false,
"error": "Invalid request body",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
// 验证必需字段
if req.ID == "" {
response := map[string]interface{}{
"success": false,
"error": "Bot ID is required",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
if req.Protocol == "" {
response := map[string]interface{}{
"success": false,
"error": "Protocol is required",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
// 检查机器人是否已存在
if _, exists := s.botManager.Get(req.ID); exists {
response := map[string]interface{}{
"success": false,
"error": "Bot with this ID already exists",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusConflict)
ctx.SetBody(data)
return
}
// TODO: 根据协议类型创建相应的机器人实例
// 这里需要协议工厂来创建不同类型的机器人
// 目前返回成功但提示需要实现协议适配器
s.logger.Info("Bot creation requested",
zap.String("bot_id", req.ID),
zap.String("protocol", req.Protocol))
response := map[string]interface{}{
"success": true,
"message": "Bot creation queued (protocol adapter implementation required)",
"bot_id": req.ID,
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusAccepted)
ctx.SetBody(data)
}
// handlePublishEvent 发布事件
@@ -147,19 +244,169 @@ func (s *Server) handlePublishEvent(ctx *fasthttp.RequestCtx) {
return
}
// 解析请求体并发布事件
// 这里简化实现,实际应该解析JSON并创建Event
ctx.SetContentType("application/json")
ctx.SetBodyString(`{"message":"Event published"}`)
// 解析请求体为事件对象
var event protocol.BaseEvent
if err := sonic.Unmarshal(ctx.PostBody(), &event); err != nil {
s.logger.Error("Failed to parse event", zap.Error(err))
response := map[string]interface{}{
"success": false,
"error": "Invalid event format",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
// 验证事件类型
if event.Type == "" {
response := map[string]interface{}{
"success": false,
"error": "Event type is required",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
// 如果没有时间戳,使用当前时间
if event.Timestamp == 0 {
event.Timestamp = ctx.Time().Unix()
}
// 确保Data字段不为nil
if event.Data == nil {
event.Data = make(map[string]interface{})
}
s.logger.Info("Publishing event",
zap.String("type", string(event.Type)),
zap.String("detail_type", event.DetailType),
zap.String("self_id", event.SelfID))
// 发布到事件总线
s.eventBus.Publish(&event)
response := map[string]interface{}{
"success": true,
"message": "Event published successfully",
"timestamp": event.Timestamp,
}
data, _ := sonic.Marshal(response)
ctx.SetBody(data)
}
// handleSubscribeEvent 订阅事件
// handleSubscribeEvent 订阅事件WebSocket升级
func (s *Server) handleSubscribeEvent(ctx *fasthttp.RequestCtx) {
if string(ctx.Method()) != "GET" {
// 检查是否为WebSocket升级请求
if !ctx.IsGet() {
ctx.Error("Method Not Allowed", fasthttp.StatusMethodNotAllowed)
return
}
ctx.SetContentType("application/json")
ctx.SetBodyString(`{"message":"Event subscription not implemented yet"}`)
// 检查是否为WebSocket升级请求
if string(ctx.Request.Header.Peek("Upgrade")) != "websocket" {
ctx.SetContentType("application/json")
response := map[string]interface{}{
"success": false,
"error": "WebSocket upgrade required",
"message": "This endpoint requires WebSocket connection",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusBadRequest)
ctx.SetBody(data)
return
}
// 获取订阅者ID可选
subscriberID := string(ctx.QueryArgs().Peek("subscriber_id"))
if subscriberID == "" {
subscriberID = "event-subscriber"
}
// 获取要订阅的事件类型(可选,为空则订阅所有类型)
eventTypeStr := string(ctx.QueryArgs().Peek("event_type"))
s.logger.Info("WebSocket event subscription request",
zap.String("subscriber_id", subscriberID),
zap.String("event_type", eventTypeStr),
zap.String("remote_addr", ctx.RemoteAddr().String()))
// 升级为WebSocket连接
wsConn, err := s.wsManager.UpgradeWebSocket(ctx)
if err != nil {
s.logger.Error("Failed to upgrade WebSocket", zap.Error(err))
ctx.SetContentType("application/json")
response := map[string]interface{}{
"success": false,
"error": "Failed to upgrade WebSocket connection",
}
data, _ := sonic.Marshal(response)
ctx.SetStatusCode(fasthttp.StatusInternalServerError)
ctx.SetBody(data)
return
}
// 订阅所有主要事件类型
eventTypes := []protocol.EventType{
protocol.EventTypeMessage,
protocol.EventTypeNotice,
protocol.EventTypeRequest,
protocol.EventTypeMeta,
protocol.EventTypeMessageSent,
protocol.EventTypeNoticeSent,
protocol.EventTypeRequestSent,
}
// 如果指定了事件类型,只订阅该类型
if eventTypeStr != "" {
eventTypes = []protocol.EventType{protocol.EventType(eventTypeStr)}
}
// 为每种事件类型创建订阅
for _, eventType := range eventTypes {
eventChan := s.eventBus.Subscribe(eventType, nil) // nil filter means accept all events
// 启动goroutine监听事件并发送到WebSocket
go func(ch chan protocol.Event, et protocol.EventType) {
for {
select {
case event, ok := <-ch:
if !ok {
return
}
// 序列化事件为JSON
data, err := sonic.Marshal(event)
if err != nil {
s.logger.Error("Failed to marshal event", zap.Error(err))
continue
}
// 发送到WebSocket连接
if err := wsConn.SendMessage(data); err != nil {
s.logger.Error("Failed to send event to WebSocket",
zap.String("conn_id", wsConn.ID),
zap.Error(err))
// 连接可能已断开,取消订阅
s.eventBus.Unsubscribe(et, ch)
return
}
case <-wsConn.ctx.Done():
// 连接关闭,取消订阅
s.eventBus.Unsubscribe(et, ch)
return
}
}
}(eventChan, eventType)
}
s.logger.Info("WebSocket event subscription established",
zap.String("conn_id", wsConn.ID),
zap.String("subscriber_id", subscriberID))
}

244
pkg/net/sse.go Normal file
View File

@@ -0,0 +1,244 @@
package net
import (
"bufio"
"context"
"fmt"
"net"
"net/http"
"strings"
"time"
"go.uber.org/zap"
)
// SSEClient Server-Sent Events 客户端
type SSEClient struct {
url string
accessToken string
eventChan chan []byte
logger *zap.Logger
reconnectDelay time.Duration
maxReconnect int
ctx context.Context
cancel context.CancelFunc
eventFilter string
}
// SSEClientConfig SSE 客户端配置
type SSEClientConfig struct {
URL string
AccessToken string
ReconnectDelay time.Duration
MaxReconnect int
EventFilter string
BufferSize int
}
// NewSSEClient 创建 SSE 客户端
func NewSSEClient(config SSEClientConfig, logger *zap.Logger) *SSEClient {
ctx, cancel := context.WithCancel(context.Background())
if config.ReconnectDelay == 0 {
config.ReconnectDelay = 5 * time.Second
}
if config.MaxReconnect == 0 {
config.MaxReconnect = -1
}
if config.BufferSize == 0 {
config.BufferSize = 100
}
return &SSEClient{
url: config.URL,
accessToken: config.AccessToken,
eventChan: make(chan []byte, config.BufferSize),
logger: logger.Named("sse-client"),
reconnectDelay: config.ReconnectDelay,
maxReconnect: config.MaxReconnect,
eventFilter: config.EventFilter,
ctx: ctx,
cancel: cancel,
}
}
// Connect 连接到 SSE 服务器
func (c *SSEClient) Connect(ctx context.Context) error {
c.logger.Info("Starting SSE client", zap.String("url", c.url))
go c.connectLoop(ctx)
return nil
}
// connectLoop 连接循环
func (c *SSEClient) connectLoop(ctx context.Context) {
reconnectCount := 0
for {
select {
case <-ctx.Done():
c.logger.Info("SSE client stopped")
return
case <-c.ctx.Done():
c.logger.Info("SSE client stopped")
return
default:
}
c.logger.Info("Connecting to SSE server",
zap.String("url", c.url),
zap.Int("reconnect_count", reconnectCount))
err := c.connect(ctx)
if err != nil {
c.logger.Error("SSE connection failed", zap.Error(err))
}
if c.maxReconnect >= 0 && reconnectCount >= c.maxReconnect {
c.logger.Error("Max reconnect attempts reached", zap.Int("count", reconnectCount))
return
}
reconnectCount++
c.logger.Info("Reconnecting after delay",
zap.Duration("delay", c.reconnectDelay),
zap.Int("attempt", reconnectCount))
select {
case <-time.After(c.reconnectDelay):
case <-ctx.Done():
return
case <-c.ctx.Done():
return
}
}
}
// connect 建立单次连接
func (c *SSEClient) connect(ctx context.Context) error {
req, err := http.NewRequestWithContext(ctx, "GET", c.url, nil)
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
if c.accessToken != "" {
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.accessToken))
}
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("Cache-Control", "no-cache")
req.Header.Set("Connection", "keep-alive")
client := &http.Client{
Timeout: 0,
Transport: &http.Transport{
DialContext: (&net.Dialer{
Timeout: 30 * time.Second,
KeepAlive: 30 * time.Second,
}).DialContext,
MaxIdleConns: 100,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
ExpectContinueTimeout: 1 * time.Second,
},
}
resp, err := client.Do(req)
if err != nil {
return fmt.Errorf("failed to connect: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
}
contentType := resp.Header.Get("Content-Type")
if !strings.HasPrefix(contentType, "text/event-stream") {
return fmt.Errorf("unexpected content type: %s", contentType)
}
c.logger.Info("SSE connection established")
return c.readEventStream(ctx, resp)
}
// readEventStream 读取事件流
func (c *SSEClient) readEventStream(ctx context.Context, resp *http.Response) error {
scanner := bufio.NewScanner(resp.Body)
scanner.Split(bufio.ScanLines)
var eventType string
var dataLines []string
for scanner.Scan() {
select {
case <-ctx.Done():
return ctx.Err()
case <-c.ctx.Done():
return c.ctx.Err()
default:
}
line := scanner.Text()
if line == "" {
if len(dataLines) > 0 {
c.processEvent(eventType, dataLines)
eventType = ""
dataLines = nil
}
continue
}
if strings.HasPrefix(line, ":") {
continue
}
if strings.HasPrefix(line, "event:") {
eventType = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
} else if strings.HasPrefix(line, "data:") {
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
dataLines = append(dataLines, data)
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("scanner error: %w", err)
}
return fmt.Errorf("connection closed")
}
// processEvent 处理事件
func (c *SSEClient) processEvent(eventType string, dataLines []string) {
if c.eventFilter != "" && eventType != c.eventFilter && eventType != "" {
c.logger.Debug("Ignoring filtered event", zap.String("event_type", eventType))
return
}
data := strings.Join(dataLines, "\n")
c.logger.Debug("Received SSE event",
zap.String("event_type", eventType),
zap.Int("data_length", len(data)))
select {
case c.eventChan <- []byte(data):
default:
c.logger.Warn("Event channel full, dropping event")
}
}
// Events 获取事件通道
func (c *SSEClient) Events() <-chan []byte {
return c.eventChan
}
// Close 关闭客户端
func (c *SSEClient) Close() error {
c.cancel()
close(c.eventChan)
c.logger.Info("SSE client closed")
return nil
}

View File

@@ -10,6 +10,7 @@ import (
"cellbot/internal/engine"
"cellbot/internal/protocol"
"github.com/bytedance/sonic"
"github.com/fasthttp/websocket"
"github.com/valyala/fasthttp"
"go.uber.org/zap"
@@ -40,27 +41,45 @@ func NewWebSocketManager(logger *zap.Logger, eventBus *engine.EventBus) *WebSock
}
}
// ConnectionType 连接类型
type ConnectionType string
const (
ConnectionTypeReverse ConnectionType = "reverse" // 反向连接(被动接受)
ConnectionTypeForward ConnectionType = "forward" // 正向连接(主动发起)
)
// WebSocketConnection WebSocket连接
type WebSocketConnection struct {
ID string
Conn *websocket.Conn
BotID string
Logger *zap.Logger
ctx context.Context
cancel context.CancelFunc
ID string
Conn *websocket.Conn
BotID string
Logger *zap.Logger
ctx context.Context
cancel context.CancelFunc
Type ConnectionType
RemoteAddr string
reconnectURL string // 用于正向连接重连
maxReconnect int // 最大重连次数
reconnectCount int // 当前重连次数
heartbeatTick time.Duration // 心跳间隔
}
// NewWebSocketConnection 创建WebSocket连接
func NewWebSocketConnection(conn *websocket.Conn, botID string, logger *zap.Logger) *WebSocketConnection {
func NewWebSocketConnection(conn *websocket.Conn, botID string, connType ConnectionType, logger *zap.Logger) *WebSocketConnection {
ctx, cancel := context.WithCancel(context.Background())
connID := generateConnID()
return &WebSocketConnection{
ID: connID,
Conn: conn,
BotID: botID,
Logger: logger.With(zap.String("conn_id", connID)),
ctx: ctx,
cancel: cancel,
ID: connID,
Conn: conn,
BotID: botID,
Logger: logger.With(zap.String("conn_id", connID)),
ctx: ctx,
cancel: cancel,
Type: connType,
RemoteAddr: conn.RemoteAddr().String(),
maxReconnect: 5,
heartbeatTick: 30 * time.Second,
}
}
@@ -85,21 +104,23 @@ func (wsm *WebSocketManager) UpgradeWebSocket(ctx *fasthttp.RequestCtx) (*WebSoc
// 等待连接在回调中建立
conn := <-connChan
// 创建连接对象
wsConn := NewWebSocketConnection(conn, botID, wsm.logger)
// 创建连接对象(反向连接)
wsConn := NewWebSocketConnection(conn, botID, ConnectionTypeReverse, wsm.logger)
// 存储连接
wsm.mu.Lock()
wsm.connections[wsConn.ID] = wsConn
wsm.mu.Unlock()
wsm.logger.Info("WebSocket connection established",
wsm.logger.Info("WebSocket reverse connection established",
zap.String("conn_id", wsConn.ID),
zap.String("bot_id", botID))
zap.String("bot_id", botID),
zap.String("remote_addr", wsConn.RemoteAddr))
// 启动读取循环
// 启动读取循环和心跳
go wsConn.readLoop(wsm.eventBus)
go wsConn.heartbeatLoop()
return wsConn, nil
}
@@ -119,9 +140,15 @@ func (wsc *WebSocketConnection) readLoop(eventBus *engine.EventBus) {
return
}
// 只处理文本消息,忽略其他类型
if messageType != websocket.TextMessage {
wsc.Logger.Warn("Received non-text message, ignoring",
zap.Int("message_type", messageType))
continue
}
// 处理消息
wsc.handleMessage(message, eventBus)
// messageType 可用于区分文本或二进制消息
}
}
}
@@ -130,18 +157,41 @@ func (wsc *WebSocketConnection) readLoop(eventBus *engine.EventBus) {
func (wsc *WebSocketConnection) handleMessage(data []byte, eventBus *engine.EventBus) {
wsc.Logger.Debug("Received message", zap.ByteString("data", data))
// TODO: 解析消息为Event对象
// 这里简化实现,实际应该根据协议解析
event := &protocol.BaseEvent{
Type: protocol.EventTypeMessage,
DetailType: "private",
Timestamp: time.Now().Unix(),
SelfID: wsc.BotID,
Data: make(map[string]interface{}),
// 解析JSON消息为BaseEvent
var event protocol.BaseEvent
if err := sonic.Unmarshal(data, &event); err != nil {
wsc.Logger.Error("Failed to parse message", zap.Error(err), zap.ByteString("data", data))
return
}
// 验证必需字段
if event.Type == "" {
wsc.Logger.Warn("Event type is empty", zap.ByteString("data", data))
return
}
// 如果没有时间戳,使用当前时间
if event.Timestamp == 0 {
event.Timestamp = time.Now().Unix()
}
// 如果没有SelfID使用连接的BotID
if event.SelfID == "" {
event.SelfID = wsc.BotID
}
// 确保Data字段不为nil
if event.Data == nil {
event.Data = make(map[string]interface{})
}
wsc.Logger.Info("Event received",
zap.String("type", string(event.Type)),
zap.String("detail_type", event.DetailType),
zap.String("self_id", event.SelfID))
// 发布到事件总线
eventBus.Publish(event)
eventBus.Publish(&event)
}
// SendMessage 发送消息
@@ -155,6 +205,80 @@ func (wsc *WebSocketConnection) SendMessage(data []byte) error {
return nil
}
// heartbeatLoop 心跳循环
func (wsc *WebSocketConnection) heartbeatLoop() {
ticker := time.NewTicker(wsc.heartbeatTick)
defer ticker.Stop()
for {
select {
case <-ticker.C:
// 发送ping消息
if err := wsc.Conn.WriteControl(websocket.PingMessage, []byte{}, time.Now().Add(10*time.Second)); err != nil {
wsc.Logger.Warn("Failed to send ping", zap.Error(err))
return
}
wsc.Logger.Debug("Heartbeat ping sent")
case <-wsc.ctx.Done():
return
}
}
}
// reconnectLoop 重连循环(仅用于正向连接)
func (wsc *WebSocketConnection) reconnectLoop(wsm *WebSocketManager) {
<-wsc.ctx.Done() // 等待连接断开
if wsc.Type != ConnectionTypeForward || wsc.reconnectURL == "" {
return
}
wsc.Logger.Info("Connection closed, attempting to reconnect",
zap.Int("max_reconnect", wsc.maxReconnect))
for wsc.reconnectCount < wsc.maxReconnect {
wsc.reconnectCount++
backoff := time.Duration(wsc.reconnectCount) * 5 * time.Second
wsc.Logger.Info("Reconnecting",
zap.Int("attempt", wsc.reconnectCount),
zap.Int("max", wsc.maxReconnect),
zap.Duration("backoff", backoff))
time.Sleep(backoff)
// 尝试重新连接
conn, _, err := websocket.DefaultDialer.Dial(wsc.reconnectURL, nil)
if err != nil {
wsc.Logger.Error("Reconnect failed", zap.Error(err))
continue
}
// 更新连接
wsc.Conn = conn
wsc.RemoteAddr = conn.RemoteAddr().String()
wsc.ctx, wsc.cancel = context.WithCancel(context.Background())
wsc.reconnectCount = 0 // 重置重连计数
wsc.Logger.Info("Reconnected successfully",
zap.String("remote_addr", wsc.RemoteAddr))
// 重新启动读取循环和心跳
go wsc.readLoop(wsm.eventBus)
go wsc.heartbeatLoop()
go wsc.reconnectLoop(wsm)
return
}
wsc.Logger.Error("Max reconnect attempts reached, giving up",
zap.Int("attempts", wsc.reconnectCount))
// 从管理器中移除连接
wsm.RemoveConnection(wsc.ID)
}
// close 关闭连接
func (wsc *WebSocketConnection) close() {
wsc.cancel()
@@ -209,31 +333,64 @@ func (wsm *WebSocketManager) BroadcastToBot(botID string, data []byte) {
}
}
// Dial 建立WebSocket客户端连接
// DialConfig WebSocket客户端连接配置
type DialConfig struct {
URL string
BotID string
MaxReconnect int
HeartbeatTick time.Duration
}
// Dial 建立WebSocket客户端连接正向连接
func (wsm *WebSocketManager) Dial(addr string, botID string) (*WebSocketConnection, error) {
u, err := url.Parse(addr)
return wsm.DialWithConfig(DialConfig{
URL: addr,
BotID: botID,
MaxReconnect: 5,
HeartbeatTick: 30 * time.Second,
})
}
// DialWithConfig 使用配置建立WebSocket客户端连接
func (wsm *WebSocketManager) DialWithConfig(config DialConfig) (*WebSocketConnection, error) {
u, err := url.Parse(config.URL)
if err != nil {
return nil, fmt.Errorf("invalid URL: %w", err)
}
conn, _, err := websocket.DefaultDialer.Dial(addr, nil)
// 验证URL scheme必须是ws或wss
if u.Scheme != "ws" && u.Scheme != "wss" {
return nil, fmt.Errorf("invalid URL scheme: %s, expected ws or wss", u.Scheme)
}
conn, _, err := websocket.DefaultDialer.Dial(config.URL, nil)
if err != nil {
return nil, fmt.Errorf("failed to dial: %w", err)
}
wsConn := NewWebSocketConnection(conn, botID, wsm.logger)
wsConn := NewWebSocketConnection(conn, config.BotID, ConnectionTypeForward, wsm.logger)
wsConn.reconnectURL = config.URL
wsConn.maxReconnect = config.MaxReconnect
wsConn.heartbeatTick = config.HeartbeatTick
wsm.mu.Lock()
wsm.connections[wsConn.ID] = wsConn
wsm.mu.Unlock()
wsm.logger.Info("WebSocket client connected",
wsm.logger.Info("WebSocket forward connection established",
zap.String("conn_id", wsConn.ID),
zap.String("bot_id", botID),
zap.String("addr", addr))
zap.String("bot_id", config.BotID),
zap.String("addr", config.URL),
zap.String("remote_addr", wsConn.RemoteAddr))
// 启动读取循环
// 启动读取循环和心跳
go wsConn.readLoop(wsm.eventBus)
go wsConn.heartbeatLoop()
// 如果是正向连接,启动重连监控
if wsConn.Type == ConnectionTypeForward {
go wsConn.reconnectLoop(wsm)
}
return wsConn, nil
}