d16261e6bd
- 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.
422 lines
10 KiB
Go
422 lines
10 KiB
Go
package engine
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"sync"
|
||
"sync/atomic"
|
||
|
||
"cellbot/internal/protocol"
|
||
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
// PluginFactory 插件工厂函数类型
|
||
// 用于在运行时创建插件实例，支持依赖注入
|
||
type PluginFactory func(botManager *protocol.BotManager, logger *zap.Logger) protocol.EventHandler
|
||
|
||
// globalPluginRegistry 全局插件注册表
|
||
var (
|
||
globalPluginRegistry = make(map[string]PluginFactory)
|
||
globalPluginMu sync.RWMutex
|
||
)
|
||
|
||
// RegisterPlugin 注册插件（供插件在 init 函数中调用）
|
||
func RegisterPlugin(name string, factory PluginFactory) {
|
||
globalPluginMu.Lock()
|
||
defer globalPluginMu.Unlock()
|
||
|
||
if _, exists := globalPluginRegistry[name]; exists {
|
||
panic("plugin already registered: " + name)
|
||
}
|
||
|
||
globalPluginRegistry[name] = factory
|
||
}
|
||
|
||
// GetRegisteredPlugins 获取所有已注册的插件名称
|
||
func GetRegisteredPlugins() []string {
|
||
globalPluginMu.RLock()
|
||
defer globalPluginMu.RUnlock()
|
||
|
||
names := make([]string, 0, len(globalPluginRegistry))
|
||
for name := range globalPluginRegistry {
|
||
names = append(names, name)
|
||
}
|
||
return names
|
||
}
|
||
|
||
// LoadPlugin 加载指定名称的插件
|
||
func LoadPlugin(name string, botManager *protocol.BotManager, logger *zap.Logger) (protocol.EventHandler, error) {
|
||
globalPluginMu.RLock()
|
||
factory, exists := globalPluginRegistry[name]
|
||
globalPluginMu.RUnlock()
|
||
|
||
if !exists {
|
||
return nil, nil // 插件不存在，返回 nil 而不是错误
|
||
}
|
||
|
||
return factory(botManager, logger), nil
|
||
}
|
||
|
||
// LoadAllPlugins 加载所有已注册的插件
|
||
func LoadAllPlugins(botManager *protocol.BotManager, logger *zap.Logger) []protocol.EventHandler {
|
||
globalPluginMu.RLock()
|
||
defer globalPluginMu.RUnlock()
|
||
|
||
plugins := make([]protocol.EventHandler, 0, len(globalPluginRegistry))
|
||
for name, factory := range globalPluginRegistry {
|
||
plugin := factory(botManager, logger.Named("plugin."+name))
|
||
plugins = append(plugins, plugin)
|
||
}
|
||
|
||
return plugins
|
||
}
|
||
|
||
// PluginRegistry 插件注册表
|
||
type PluginRegistry struct {
|
||
plugins []protocol.EventHandler
|
||
dispatcher *Dispatcher
|
||
logger *zap.Logger
|
||
mu sync.RWMutex
|
||
}
|
||
|
||
// NewPluginRegistry 创建插件注册表
|
||
func NewPluginRegistry(dispatcher *Dispatcher, logger *zap.Logger) *PluginRegistry {
|
||
return &PluginRegistry{
|
||
plugins: make([]protocol.EventHandler, 0),
|
||
dispatcher: dispatcher,
|
||
logger: logger.Named("plugin-registry"),
|
||
}
|
||
}
|
||
|
||
// Register 注册插件
|
||
func (r *PluginRegistry) Register(plugin protocol.EventHandler) {
|
||
r.mu.Lock()
|
||
defer r.mu.Unlock()
|
||
|
||
r.plugins = append(r.plugins, plugin)
|
||
r.dispatcher.RegisterHandler(plugin)
|
||
r.logger.Info("Plugin registered",
|
||
zap.String("name", plugin.Name()),
|
||
zap.String("description", plugin.Description()))
|
||
}
|
||
|
||
// GetPlugins 获取所有插件
|
||
func (r *PluginRegistry) GetPlugins() []protocol.EventHandler {
|
||
r.mu.RLock()
|
||
defer r.mu.RUnlock()
|
||
return r.plugins
|
||
}
|
||
|
||
// PluginBuilder 插件构建器
|
||
type PluginBuilder struct {
|
||
name string
|
||
description string
|
||
priority int
|
||
matchFunc func(protocol.Event) bool
|
||
handleFunc func(context.Context, protocol.Event) error
|
||
}
|
||
|
||
// NewPlugin 创建插件构建器
|
||
func NewPlugin(name string) *PluginBuilder {
|
||
return &PluginBuilder{
|
||
name: name,
|
||
priority: 100, // 默认优先级
|
||
}
|
||
}
|
||
|
||
// Description 设置插件描述
|
||
func (b *PluginBuilder) Description(desc string) *PluginBuilder {
|
||
b.description = desc
|
||
return b
|
||
}
|
||
|
||
// Priority 设置优先级
|
||
func (b *PluginBuilder) Priority(priority int) *PluginBuilder {
|
||
b.priority = priority
|
||
return b
|
||
}
|
||
|
||
// Match 设置匹配函数
|
||
func (b *PluginBuilder) Match(fn func(protocol.Event) bool) *PluginBuilder {
|
||
b.matchFunc = fn
|
||
return b
|
||
}
|
||
|
||
// Handle 设置处理函数
|
||
func (b *PluginBuilder) Handle(fn func(context.Context, protocol.Event) error) *PluginBuilder {
|
||
b.handleFunc = fn
|
||
return b
|
||
}
|
||
|
||
// Build 构建插件
|
||
func (b *PluginBuilder) Build() protocol.EventHandler {
|
||
return &simplePlugin{
|
||
name: b.name,
|
||
description: b.description,
|
||
priority: b.priority,
|
||
matchFunc: b.matchFunc,
|
||
handleFunc: b.handleFunc,
|
||
}
|
||
}
|
||
|
||
// simplePlugin 简单插件实现
|
||
type simplePlugin struct {
|
||
name string
|
||
description string
|
||
priority int
|
||
matchFunc func(protocol.Event) bool
|
||
handleFunc func(context.Context, protocol.Event) error
|
||
}
|
||
|
||
func (p *simplePlugin) Name() string {
|
||
return p.name
|
||
}
|
||
|
||
func (p *simplePlugin) Description() string {
|
||
return p.description
|
||
}
|
||
|
||
func (p *simplePlugin) Priority() int {
|
||
return p.priority
|
||
}
|
||
|
||
func (p *simplePlugin) Match(event protocol.Event) bool {
|
||
if p.matchFunc == nil {
|
||
return true
|
||
}
|
||
return p.matchFunc(event)
|
||
}
|
||
|
||
func (p *simplePlugin) Handle(ctx context.Context, event protocol.Event) error {
|
||
if p.handleFunc == nil {
|
||
return nil
|
||
}
|
||
return p.handleFunc(ctx, event)
|
||
}
|
||
|
||
// ============================================================================
|
||
// ZeroBot 风格的全局注册 API
|
||
// ============================================================================
|
||
|
||
// globalHandlerRegistry 全局处理器注册表
|
||
var (
|
||
globalHandlerRegistry = make([]*HandlerBuilder, 0)
|
||
globalHandlerMu sync.RWMutex
|
||
)
|
||
|
||
// HandlerFunc 处理函数类型（支持依赖注入）
|
||
type HandlerFunc func(ctx context.Context, event protocol.Event, botManager *protocol.BotManager, logger *zap.Logger) error
|
||
|
||
// HandlerBuilder 处理器构建器（类似 ZeroBot 的 API）
|
||
type HandlerBuilder struct {
|
||
matchFunc func(protocol.Event) bool
|
||
priority int
|
||
handleFunc HandlerFunc
|
||
}
|
||
|
||
// OnPrivateMessage 匹配私聊消息
|
||
func OnPrivateMessage() *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
return event.GetType() == protocol.EventTypeMessage && event.GetDetailType() == "private"
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnGroupMessage 匹配群消息
|
||
func OnGroupMessage() *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
return event.GetType() == protocol.EventTypeMessage && event.GetDetailType() == "group"
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnMessage 匹配所有消息
|
||
func OnMessage() *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
return event.GetType() == protocol.EventTypeMessage
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnNotice 匹配通知事件
|
||
func OnNotice() *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
return event.GetType() == protocol.EventTypeNotice
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnRequest 匹配请求事件
|
||
func OnRequest() *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
return event.GetType() == protocol.EventTypeRequest
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// On 自定义匹配器
|
||
func On(matchFunc func(protocol.Event) bool) *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: matchFunc,
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnCommand 匹配命令（以指定前缀开头的消息）
|
||
func OnCommand(prefix string) *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
if event.GetType() != protocol.EventTypeMessage {
|
||
return false
|
||
}
|
||
data := event.GetData()
|
||
rawMessage, ok := data["raw_message"].(string)
|
||
if !ok {
|
||
return false
|
||
}
|
||
// 检查是否以命令前缀开头
|
||
if len(rawMessage) > 0 && len(prefix) > 0 {
|
||
return len(rawMessage) >= len(prefix) && rawMessage[:len(prefix)] == prefix
|
||
}
|
||
return false
|
||
},
|
||
priority: 50, // 命令通常优先级较高
|
||
}
|
||
}
|
||
|
||
// OnPrefix 匹配以指定前缀开头的消息
|
||
func OnPrefix(prefix string) *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
if event.GetType() != protocol.EventTypeMessage {
|
||
return false
|
||
}
|
||
data := event.GetData()
|
||
rawMessage, ok := data["raw_message"].(string)
|
||
if !ok {
|
||
return false
|
||
}
|
||
return len(rawMessage) >= len(prefix) && rawMessage[:len(prefix)] == prefix
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnSuffix 匹配以指定后缀结尾的消息
|
||
func OnSuffix(suffix string) *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
if event.GetType() != protocol.EventTypeMessage {
|
||
return false
|
||
}
|
||
data := event.GetData()
|
||
rawMessage, ok := data["raw_message"].(string)
|
||
if !ok {
|
||
return false
|
||
}
|
||
return len(rawMessage) >= len(suffix) && rawMessage[len(rawMessage)-len(suffix):] == suffix
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// OnKeyword 匹配包含指定关键词的消息
|
||
func OnKeyword(keyword string) *HandlerBuilder {
|
||
return &HandlerBuilder{
|
||
matchFunc: func(event protocol.Event) bool {
|
||
if event.GetType() != protocol.EventTypeMessage {
|
||
return false
|
||
}
|
||
data := event.GetData()
|
||
rawMessage, ok := data["raw_message"].(string)
|
||
if !ok {
|
||
return false
|
||
}
|
||
// 简单的字符串包含检查
|
||
return len(rawMessage) >= len(keyword) &&
|
||
(len(rawMessage) == len(keyword) ||
|
||
(rawMessage[:len(keyword)] == keyword ||
|
||
rawMessage[len(rawMessage)-len(keyword):] == keyword ||
|
||
contains(rawMessage, keyword)))
|
||
},
|
||
priority: 100,
|
||
}
|
||
}
|
||
|
||
// contains 检查字符串是否包含子串（简单实现）
|
||
func contains(s, substr string) bool {
|
||
for i := 0; i <= len(s)-len(substr); i++ {
|
||
if s[i:i+len(substr)] == substr {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// Priority 设置优先级
|
||
func (b *HandlerBuilder) Priority(priority int) *HandlerBuilder {
|
||
b.priority = priority
|
||
return b
|
||
}
|
||
|
||
// Handle 注册处理函数（在 init 中调用）
|
||
func (b *HandlerBuilder) Handle(handleFunc HandlerFunc) {
|
||
globalHandlerMu.Lock()
|
||
defer globalHandlerMu.Unlock()
|
||
|
||
b.handleFunc = handleFunc
|
||
globalHandlerRegistry = append(globalHandlerRegistry, b)
|
||
}
|
||
|
||
// generateHandlerName 生成处理器名称
|
||
var handlerCounter int64
|
||
|
||
func generateHandlerName() string {
|
||
counter := atomic.AddInt64(&handlerCounter, 1)
|
||
return fmt.Sprintf("handler_%d", counter)
|
||
}
|
||
|
||
// LoadAllHandlers 加载所有已注册的处理器（注入依赖）
|
||
func LoadAllHandlers(botManager *protocol.BotManager, logger *zap.Logger) []protocol.EventHandler {
|
||
globalHandlerMu.RLock()
|
||
defer globalHandlerMu.RUnlock()
|
||
|
||
handlers := make([]protocol.EventHandler, 0, len(globalHandlerRegistry))
|
||
for i, builder := range globalHandlerRegistry {
|
||
if builder.handleFunc == nil {
|
||
continue
|
||
}
|
||
|
||
// 生成唯一的插件名称
|
||
pluginName := fmt.Sprintf("handler_%d", i+1)
|
||
|
||
// 创建包装的处理函数，注入依赖
|
||
handleFunc := func(ctx context.Context, event protocol.Event) error {
|
||
return builder.handleFunc(ctx, event, botManager, logger)
|
||
}
|
||
|
||
handler := &simplePlugin{
|
||
name: pluginName,
|
||
description: "Handler registered via OnXXX().Handle()",
|
||
priority: builder.priority,
|
||
matchFunc: builder.matchFunc,
|
||
handleFunc: handleFunc,
|
||
}
|
||
|
||
handlers = append(handlers, handler)
|
||
}
|
||
|
||
return handlers
|
||
}
|