diff --git a/config/config.example.json b/config/config.example.json index 36783d0ea..872358bd4 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -159,7 +159,8 @@ "webhook_port": 18791, "webhook_path": "/webhook/wecom-aibot", "max_steps": 10, - "welcome_message": "Hello! I'm your AI assistant. How can I help you today?" + "welcome_message": "Hello! I'm your AI assistant. How can I help you today?", + "reasoning_channel_id": "" } }, "providers": { diff --git a/pkg/channels/wecom/aibot.go b/pkg/channels/wecom/aibot.go index 115fde9f7..c4970d8fb 100644 --- a/pkg/channels/wecom/aibot.go +++ b/pkg/channels/wecom/aibot.go @@ -130,6 +130,7 @@ func NewWeComAIBotChannel( base := channels.NewBaseChannel("wecom_aibot", cfg, messageBus, cfg.AllowFrom, channels.WithMaxMessageLength(2048), + channels.WithReasoningChannelID(cfg.ReasoningChannelID), ) return &WeComAIBotChannel{ @@ -336,8 +337,9 @@ func (c *WeComAIBotChannel) handleMessageCallback( timestamp := r.URL.Query().Get("timestamp") nonce := r.URL.Query().Get("nonce") - // Read request body - body, err := io.ReadAll(r.Body) + // Read request body (limit to 4 MB to prevent memory exhaustion) + const maxBodySize = 4 << 20 // 4 MB + body, err := io.ReadAll(io.LimitReader(r.Body, maxBodySize+1)) if err != nil { logger.ErrorCF("wecom_aibot", "Failed to read request body", map[string]any{ "error": err, @@ -345,6 +347,10 @@ func (c *WeComAIBotChannel) handleMessageCallback( http.Error(w, "Failed to read body", http.StatusBadRequest) return } + if len(body) > maxBodySize { + http.Error(w, "Request body too large", http.StatusRequestEntityTooLarge) + return + } // Parse JSON body to get encrypted message // Format: {"encrypt": "base64_encrypted_string"} @@ -1024,10 +1030,15 @@ func (c *WeComAIBotChannel) downloadAndDecryptImage( return nil, fmt.Errorf("download failed with status: %d", resp.StatusCode) } - encryptedData, err := io.ReadAll(resp.Body) + // Limit image download to 20 MB to prevent memory exhaustion + const maxImageSize = 20 << 20 // 20 MB + encryptedData, err := io.ReadAll(io.LimitReader(resp.Body, maxImageSize+1)) if err != nil { return nil, fmt.Errorf("failed to read image data: %w", err) } + if len(encryptedData) > maxImageSize { + return nil, fmt.Errorf("image too large (exceeds %d MB)", maxImageSize>>20) + } logger.DebugCF("wecom_aibot", "Image downloaded", map[string]any{ "size": len(encryptedData), diff --git a/pkg/config/config.go b/pkg/config/config.go index 66f3945ed..439c2b995 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -362,16 +362,17 @@ type WeComAppConfig struct { } type WeComAIBotConfig struct { - Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"` - Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"` - EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"` - WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_HOST"` - WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PORT"` - WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"` - AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"` - ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"` - MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` // Maximum streaming steps - WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` // Sent on enter_chat event; empty = no welcome + Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENABLED"` + Token string `json:"token" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_TOKEN"` + EncodingAESKey string `json:"encoding_aes_key" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ENCODING_AES_KEY"` + WebhookHost string `json:"webhook_host" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_HOST"` + WebhookPort int `json:"webhook_port" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PORT"` + WebhookPath string `json:"webhook_path" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WEBHOOK_PATH"` + AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_ALLOW_FROM"` + ReplyTimeout int `json:"reply_timeout" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REPLY_TIMEOUT"` + MaxSteps int `json:"max_steps" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_MAX_STEPS"` // Maximum streaming steps + WelcomeMessage string `json:"welcome_message" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_WELCOME_MESSAGE"` // Sent on enter_chat event; empty = no welcome + ReasoningChannelID string `json:"reasoning_channel_id" env:"PICOCLAW_CHANNELS_WECOM_AIBOT_REASONING_CHANNEL_ID"` } type PicoConfig struct {