feat: add rate limiting and improve event handling

- Introduced rate limiting configuration in config.toml with options for enabling, requests per second (RPS), and burst capacity.
- Enhanced event handling in the OneBot11 adapter to ignore messages sent by the bot itself.
- Updated the dispatcher to register rate limit middleware based on configuration settings.
- Refactored WebSocket message handling to support flexible JSON parsing and improved event type detection.
- Removed deprecated echo plugin and associated tests to streamline the codebase.
This commit is contained in:
lafay
2026-01-05 01:00:38 +08:00
parent 44fe05ff62
commit d16261e6bd
11 changed files with 1001 additions and 427 deletions

View File

@@ -157,32 +157,122 @@ 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))
// 解析JSON消息为BaseEvent
var event protocol.BaseEvent
if err := sonic.Unmarshal(data, &event); err != nil {
// 解析为 map 以支持灵活的字段类型(如 self_id 可能是数字或字符串)
// 使用 sonic.Config 配置更宽松的解析,允许数字和字符串之间的转换
cfg := sonic.Config{
UseInt64: true, // 使用 int64 而不是 float64 来解析数字
NoValidateJSONSkip: true, // 跳过类型验证,允许更灵活的类型转换
}.Froze()
var rawMap map[string]interface{}
if err := cfg.Unmarshal(data, &rawMap); err != nil {
wsc.Logger.Error("Failed to parse message", zap.Error(err), zap.ByteString("data", data))
return
}
// 检查是否是 API 响应(有 echo 字段且没有 post_type
// 如果是响应,不在这里处理,让 adapter 的 handleWebSocketMessages 处理
if echo, hasEcho := rawMap["echo"].(string); hasEcho && echo != "" {
if _, hasPostType := rawMap["post_type"]; !hasPostType {
// 这是 API 响应,不在这里处理
// 正向 WebSocket 时adapter 的 handleWebSocketMessages 会处理
// 反向 WebSocket 时,响应应该通过 adapter 处理
wsc.Logger.Debug("Skipping API response in handleMessage, will be handled by adapter",
zap.String("echo", echo))
return
}
}
// 构建 BaseEvent
event := &protocol.BaseEvent{
Data: make(map[string]interface{}),
}
// 处理 self_id可能是数字或字符串
if selfIDVal, ok := rawMap["self_id"]; ok {
switch v := selfIDVal.(type) {
case string:
event.SelfID = v
case float64:
event.SelfID = fmt.Sprintf("%.0f", v)
case int64:
event.SelfID = fmt.Sprintf("%d", v)
case int:
event.SelfID = fmt.Sprintf("%d", v)
default:
event.SelfID = fmt.Sprintf("%v", v)
}
}
// 如果没有SelfID使用连接的BotID
if event.SelfID == "" {
event.SelfID = wsc.BotID
}
// 处理时间戳
if timeVal, ok := rawMap["time"]; ok {
switch v := timeVal.(type) {
case float64:
event.Timestamp = int64(v)
case int64:
event.Timestamp = v
case int:
event.Timestamp = int64(v)
}
}
if event.Timestamp == 0 {
event.Timestamp = time.Now().Unix()
}
// 处理类型字段
if typeVal, ok := rawMap["post_type"]; ok {
if typeStr, ok := typeVal.(string); ok {
// OneBot11 格式post_type -> EventType 映射
switch typeStr {
case "message":
event.Type = protocol.EventTypeMessage
case "notice":
event.Type = protocol.EventTypeNotice
case "request":
event.Type = protocol.EventTypeRequest
case "meta_event":
event.Type = protocol.EventTypeMeta
case "message_sent":
// 忽略机器人自己发送的消息
wsc.Logger.Debug("Ignoring message_sent event")
return
default:
event.Type = protocol.EventType(typeStr)
}
}
} else if typeVal, ok := rawMap["type"]; ok {
if typeStr, ok := typeVal.(string); ok {
event.Type = protocol.EventType(typeStr)
}
}
// 验证必需字段
if event.Type == "" {
wsc.Logger.Warn("Event type is empty", zap.ByteString("data", data))
return
}
// 如果没有时间戳,使用当前时间
if event.Timestamp == 0 {
event.Timestamp = time.Now().Unix()
// 处理 detail_type
if detailTypeVal, ok := rawMap["message_type"]; ok {
if detailTypeStr, ok := detailTypeVal.(string); ok {
event.DetailType = detailTypeStr
}
} else if detailTypeVal, ok := rawMap["detail_type"]; ok {
if detailTypeStr, ok := detailTypeVal.(string); ok {
event.DetailType = detailTypeStr
}
}
// 如果没有SelfID使用连接的BotID
if event.SelfID == "" {
event.SelfID = wsc.BotID
}
// 确保Data字段不为nil
if event.Data == nil {
event.Data = make(map[string]interface{})
// 将所有其他字段放入 Data
for k, v := range rawMap {
if k != "self_id" && k != "time" && k != "post_type" && k != "type" && k != "message_type" && k != "detail_type" {
event.Data[k] = v
}
}
wsc.Logger.Info("Event received",
@@ -191,7 +281,7 @@ func (wsc *WebSocketConnection) handleMessage(data []byte, eventBus *engine.Even
zap.String("self_id", event.SelfID))
// 发布到事件总线
eventBus.Publish(&event)
eventBus.Publish(event)
}
// SendMessage 发送消息
@@ -339,6 +429,7 @@ type DialConfig struct {
BotID string
MaxReconnect int
HeartbeatTick time.Duration
AutoReadLoop bool // 是否自动启动 readLoopadapter 自己处理消息时设为 false
}
// Dial 建立WebSocket客户端连接正向连接
@@ -348,6 +439,7 @@ func (wsm *WebSocketManager) Dial(addr string, botID string) (*WebSocketConnecti
BotID: botID,
MaxReconnect: 5,
HeartbeatTick: 30 * time.Second,
AutoReadLoop: false, // adapter 自己处理消息,不自动启动 readLoop
})
}
@@ -383,8 +475,10 @@ func (wsm *WebSocketManager) DialWithConfig(config DialConfig) (*WebSocketConnec
zap.String("addr", config.URL),
zap.String("remote_addr", wsConn.RemoteAddr))
// 启动读取循环和心跳
go wsConn.readLoop(wsm.eventBus)
// 启动读取循环和心跳(如果启用)
if config.AutoReadLoop {
go wsConn.readLoop(wsm.eventBus)
}
go wsConn.heartbeatLoop()
// 如果是正向连接,启动重连监控