From d69ace8c94620cea1ce05b6e68feaa9c45f5346d Mon Sep 17 00:00:00 2001 From: Hz_ Date: Fri, 31 Jul 2026 10:49:54 +0800 Subject: [PATCH] feat(go-channels): add Telegram and WeCom channels (#17564) ## Summary - Add Telegram long-polling channel support. - Add WeCom webhook and WebSocket channel support with tests. Related to #17520 --- internal/channels/bootstrap.go | 4 + internal/channels/telegram.go | 395 ++++++++++++++++ internal/channels/telegram_test.go | 254 ++++++++++ internal/channels/wecom.go | 737 +++++++++++++++++++++++++++++ internal/channels/wecom_test.go | 524 ++++++++++++++++++++ internal/channels/wecom_webhook.go | 315 ++++++++++++ 6 files changed, 2229 insertions(+) create mode 100644 internal/channels/telegram_test.go create mode 100644 internal/channels/wecom_test.go create mode 100644 internal/channels/wecom_webhook.go diff --git a/internal/channels/bootstrap.go b/internal/channels/bootstrap.go index 6a7677946e..13146cf654 100644 --- a/internal/channels/bootstrap.go +++ b/internal/channels/bootstrap.go @@ -281,6 +281,10 @@ func buildChannel(accountID string, wanted desiredChannel) (core.Channel, error) return newDiscordChannelFromConfig(accountID, wanted.credential) case "whatsapp": return newWhatsAppChannelFromConfig(accountID, wanted.credential) + case "telegram": + return newTelegramChannelFromConfig(accountID, wanted.credential) + case "wecom": + return newWeComChannelFromConfig(accountID, wanted.credential) case "dingtalk": return newDingTalkChannelFromConfig(accountID, wanted.credential) default: diff --git a/internal/channels/telegram.go b/internal/channels/telegram.go index a4ce67278c..b40e485319 100644 --- a/internal/channels/telegram.go +++ b/internal/channels/telegram.go @@ -15,3 +15,398 @@ // package channels + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "net/url" + "strconv" + "strings" + "sync" + "time" + + "ragflow/internal/channels/core" +) + +const ( + defaultTelegramAPIBaseURL = "https://api.telegram.org" + telegramPollTimeout = 30 * time.Second + telegramHTTPTimeout = 40 * time.Second + telegramReconnectDelay = 3 * time.Second +) + +type telegramAccount struct { + AccountID string + Token string + APIBaseURL string +} + +type telegramChannel struct { + account telegramAccount + + mu sync.Mutex + cancel context.CancelFunc + done chan struct{} + handler core.MessageHandler + client *http.Client + offset int64 +} + +type telegramAPIResponse struct { + OK bool `json:"ok"` + Result json.RawMessage `json:"result"` + Description string `json:"description"` +} + +type telegramTransportError struct { + method string + err error +} + +func (e *telegramTransportError) Error() string { + return fmt.Sprintf("Telegram API %s transport request failed", e.method) +} + +func (e *telegramTransportError) Unwrap() error { + return e.err +} + +type telegramUpdate struct { + UpdateID int64 `json:"update_id"` + Message *telegramMessage `json:"message"` + EditedMessage *telegramMessage `json:"edited_message"` + ChannelPost *telegramMessage `json:"channel_post"` + EditedChannelPost *telegramMessage `json:"edited_channel_post"` + Raw map[string]any +} + +type telegramMessage struct { + MessageID int64 `json:"message_id"` + From *telegramUser `json:"from"` + Chat telegramChat `json:"chat"` + Text string `json:"text"` + Caption string `json:"caption"` +} + +type telegramUser struct { + ID int64 `json:"id"` + IsBot bool `json:"is_bot"` +} + +type telegramChat struct { + ID int64 `json:"id"` + Type string `json:"type"` +} + +// newTelegramChannel creates a Telegram channel with default API settings. +func newTelegramChannel(account telegramAccount) *telegramChannel { + account.Token = strings.TrimSpace(account.Token) + account.APIBaseURL = strings.TrimSpace(account.APIBaseURL) + if account.APIBaseURL == "" { + account.APIBaseURL = defaultTelegramAPIBaseURL + } + return &telegramChannel{ + account: account, + client: &http.Client{Timeout: telegramHTTPTimeout}, + } +} + +// newTelegramChannelFromConfig builds a Telegram channel from chat_channel.config.credential. +func newTelegramChannelFromConfig(accountID string, cfg map[string]any) (*telegramChannel, error) { + token := strings.TrimSpace(fmt.Sprint(cfg["token"])) + if token == "" || token == "" { + return nil, fmt.Errorf("telegram account %q is missing token", accountID) + } + + // APIBaseURL is intentionally optional and is useful for Bot API-compatible + // proxies and tests. The Python implementation defaults to api.telegram.org. + baseURL := firstString(cfg, "api_base_url", "base_url") + return newTelegramChannel(telegramAccount{ + AccountID: accountID, + Token: token, + APIBaseURL: baseURL, + }), nil +} + +// ChannelID returns the platform identifier used by the chat-channel runtime. +func (c *telegramChannel) ChannelID() string { + return "telegram" +} + +// AccountID returns the chat_channel.id bound to this Telegram runtime instance. +func (c *telegramChannel) AccountID() string { + return c.account.AccountID +} + +// SetMessageHandler installs the RAGFlow bridge invoked for inbound messages. +func (c *telegramChannel) SetMessageHandler(handler core.MessageHandler) { + c.mu.Lock() + defer c.mu.Unlock() + c.handler = handler +} + +// Start starts Telegram long polling in the background. +func (c *telegramChannel) Start(ctx context.Context) error { + c.mu.Lock() + if c.cancel != nil { + c.mu.Unlock() + return nil + } + runCtx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + c.cancel = cancel + c.done = done + c.offset = 0 + c.mu.Unlock() + + go func() { + defer close(done) + c.run(runCtx) + + c.mu.Lock() + if c.done == done { + c.cancel = nil + } + c.mu.Unlock() + }() + return nil +} + +// Stop cancels Telegram long polling and waits for the poll request to finish. +func (c *telegramChannel) Stop(ctx context.Context) error { + c.mu.Lock() + cancel := c.cancel + done := c.done + c.mu.Unlock() + if cancel == nil { + return nil + } + + cancel() + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// Send sends a Telegram message and optionally replies to an existing message. +// Ai -> user +func (c *telegramChannel) Send(ctx context.Context, msg core.OutgoingMessage) error { + chatID, err := strconv.ParseInt(strings.TrimSpace(msg.ChatID), 10, 64) + if err != nil { + log.Printf("[telegram:%s] invalid chat_id %q: %v", c.account.AccountID, msg.ChatID, err) + return nil + } + + payload := map[string]any{ + "chat_id": chatID, + "text": msg.Text, + } + if replyID := strings.TrimSpace(msg.ReplyToMessageID); replyID != "" { + if messageID, parseErr := strconv.ParseInt(replyID, 10, 64); parseErr == nil { + payload["reply_parameters"] = map[string]any{ + "message_id": messageID, + "allow_sending_without_reply": true, + } + } + } + if err := c.callAPI(ctx, "sendMessage", payload, nil); err != nil { + log.Printf("[telegram:%s] send failed: %v", c.account.AccountID, err) + } + return nil +} + +// run drops queued updates and consumes Telegram updates until stopped. +func (c *telegramChannel) run(ctx context.Context) { + for ctx.Err() == nil { + err := c.callAPI(ctx, "deleteWebhook", map[string]any{ + "drop_pending_updates": true, + }, nil) + if err == nil { + break + } + log.Printf("[telegram:%s] failed to start polling: %v", c.account.AccountID, err) + if !waitForTelegramRetry(ctx) { + return + } + } + if ctx.Err() != nil { + return + } + + for ctx.Err() == nil { + updates, err := c.getUpdates(ctx) + if err != nil { + if ctx.Err() != nil { + return + } + log.Printf("[telegram:%s] polling failed: %v", c.account.AccountID, err) + if !waitForTelegramRetry(ctx) { + return + } + continue + } + for _, update := range updates { + if update.UpdateID >= c.offset { + c.offset = update.UpdateID + 1 + } + c.handleUpdate(ctx, update) + } + } +} + +// getUpdates performs one Telegram long-poll request and decodes raw updates. +func (c *telegramChannel) getUpdates(ctx context.Context) ([]telegramUpdate, error) { + params := map[string]any{ + "offset": c.offset, + "timeout": int(telegramPollTimeout / time.Second), + } + var rawUpdates []json.RawMessage + if err := c.callAPI(ctx, "getUpdates", params, &rawUpdates); err != nil { + return nil, err + } + + updates := make([]telegramUpdate, 0, len(rawUpdates)) + for _, raw := range rawUpdates { + var update telegramUpdate + if err := json.Unmarshal(raw, &update); err != nil { + log.Printf("[telegram:%s] ignored invalid update: %v", c.account.AccountID, err) + continue + } + _ = json.Unmarshal(raw, &update.Raw) + updates = append(updates, update) + } + return updates, nil +} + +// handleUpdate maps Telegram's effective message into the common channel model. +// Telegram -> AI +func (c *telegramChannel) handleUpdate(ctx context.Context, update telegramUpdate) { + msg := update.effectiveMessage() + if msg == nil { + return + } + if msg.From == nil && msg != update.ChannelPost && msg != update.EditedChannelPost { + return + } + if msg.From != nil && msg.From.IsBot { + return + } + + text := msg.Text + if text == "" { + text = msg.Caption + } + senderID := strconv.FormatInt(msg.Chat.ID, 10) + if msg.From != nil { + senderID = strconv.FormatInt(msg.From.ID, 10) + } + incoming := core.IncomingMessage{ + Channel: c.ChannelID(), + AccountID: c.account.AccountID, + ChatID: strconv.FormatInt(msg.Chat.ID, 10), + ChatType: telegramChatType(msg.Chat.Type), + MessageID: strconv.FormatInt(msg.MessageID, 10), + SenderID: senderID, + Text: text, + Raw: update.Raw, + } + + c.mu.Lock() + handler := c.handler + c.mu.Unlock() + if handler == nil { + return + } + if err := handler(ctx, incoming); err != nil { + log.Printf("[telegram:%s] message handler error: %v", c.account.AccountID, err) + } +} + +func (u telegramUpdate) effectiveMessage() *telegramMessage { + if u.Message != nil { + return u.Message + } + if u.EditedMessage != nil { + return u.EditedMessage + } + if u.ChannelPost != nil { + return u.ChannelPost + } + return u.EditedChannelPost +} + +func telegramChatType(value string) string { + switch value { + case "private": + return "p2p" + case "group", "supergroup": + return "group" + case "channel": + return "channel" + default: + if value == "" { + return "unknown" + } + return value + } +} + +// callAPI invokes a Telegram Bot API method using the configured token. +func (c *telegramChannel) callAPI(ctx context.Context, method string, params map[string]any, out any) error { + payload, err := json.Marshal(params) + if err != nil { + return err + } + baseURL := strings.TrimRight(c.account.APIBaseURL, "/") + endpoint := baseURL + "/bot" + url.PathEscape(c.account.Token) + "/" + method + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(payload)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + resp, err := c.client.Do(req) + if err != nil { + return &telegramTransportError{method: method, err: err} + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1024*1024)) + if err != nil { + return err + } + if resp.StatusCode >= http.StatusBadRequest { + return fmt.Errorf("Telegram API %s returned status %d: %s", method, resp.StatusCode, string(body)) + } + + var response telegramAPIResponse + if err := json.Unmarshal(body, &response); err != nil { + return fmt.Errorf("Telegram API %s returned invalid JSON: %w", method, err) + } + if !response.OK { + return fmt.Errorf("Telegram API %s failed: %s", method, response.Description) + } + if out != nil && len(response.Result) > 0 { + if err := json.Unmarshal(response.Result, out); err != nil { + return fmt.Errorf("Telegram API %s returned invalid result: %w", method, err) + } + } + return nil +} + +func waitForTelegramRetry(ctx context.Context) bool { + timer := time.NewTimer(telegramReconnectDelay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/internal/channels/telegram_test.go b/internal/channels/telegram_test.go new file mode 100644 index 0000000000..f1b04ea18d --- /dev/null +++ b/internal/channels/telegram_test.go @@ -0,0 +1,254 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package channels + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "ragflow/internal/channels/core" +) + +type telegramRoundTripperFunc func(*http.Request) (*http.Response, error) + +func (f telegramRoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +func TestNewTelegramChannelFromConfigRequiresToken(t *testing.T) { + if _, err := newTelegramChannelFromConfig("account-1", map[string]any{}); err == nil { + t.Fatal("newTelegramChannelFromConfig accepted a missing token") + } +} + +func TestTelegramChatType(t *testing.T) { + tests := map[string]string{ + "private": "p2p", + "group": "group", + "supergroup": "group", + "channel": "channel", + "": "unknown", + "custom": "custom", + } + for input, want := range tests { + if got := telegramChatType(input); got != want { + t.Fatalf("telegramChatType(%q) = %q, want %q", input, got, want) + } + } +} + +func TestTelegramHandleUpdateMapsCaptionAndIgnoresBots(t *testing.T) { + ch := newTelegramChannel(telegramAccount{AccountID: "account-1", Token: "token"}) + messages := make(chan core.IncomingMessage, 1) + ch.SetMessageHandler(func(_ context.Context, msg core.IncomingMessage) error { + messages <- msg + return nil + }) + + ch.handleUpdate(context.Background(), telegramUpdate{ + UpdateID: 1, + Message: &telegramMessage{ + MessageID: 42, + From: &telegramUser{ID: 7}, + Chat: telegramChat{ID: -10, Type: "supergroup"}, + Caption: "caption text", + }, + }) + ch.handleUpdate(context.Background(), telegramUpdate{ + UpdateID: 2, + Message: &telegramMessage{ + MessageID: 43, + From: &telegramUser{ID: 8, IsBot: true}, + }, + }) + + select { + case msg := <-messages: + if msg.ChatID != "-10" || msg.ChatType != "group" || msg.MessageID != "42" || msg.SenderID != "7" || msg.Text != "caption text" { + t.Fatalf("unexpected incoming message: %+v", msg) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for incoming message") + } + + select { + case msg := <-messages: + t.Fatalf("bot message was dispatched: %+v", msg) + case <-time.After(20 * time.Millisecond): + } +} + +func TestTelegramHandleUpdateMapsChannelPostWithoutSender(t *testing.T) { + ch := newTelegramChannel(telegramAccount{AccountID: "account-1", Token: "token"}) + messages := make(chan core.IncomingMessage, 1) + ch.SetMessageHandler(func(_ context.Context, msg core.IncomingMessage) error { + messages <- msg + return nil + }) + + ch.handleUpdate(context.Background(), telegramUpdate{ + ChannelPost: &telegramMessage{ + MessageID: 42, + Chat: telegramChat{ID: -100123, Type: "channel"}, + Text: "channel post", + }, + }) + + select { + case msg := <-messages: + if msg.ChatID != "-100123" || msg.ChatType != "channel" || msg.MessageID != "42" || msg.SenderID != "-100123" || msg.Text != "channel post" { + t.Fatalf("unexpected incoming message: %+v", msg) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for channel post") + } +} + +func TestTelegramCallAPISanitizesTransportError(t *testing.T) { + const token = "sensitive-bot-token" + transportErr := &url.Error{ + Op: "Post", + URL: "https://api.telegram.org/bot" + token + "/sendMessage", + Err: errors.New("connection refused"), + } + ch := newTelegramChannel(telegramAccount{AccountID: "account-1", Token: token}) + ch.client = &http.Client{Transport: telegramRoundTripperFunc(func(*http.Request) (*http.Response, error) { + return nil, transportErr + })} + + err := ch.callAPI(context.Background(), "sendMessage", map[string]any{}, nil) + if err == nil { + t.Fatal("callAPI() error = nil") + } + if strings.Contains(err.Error(), token) { + t.Fatalf("callAPI() exposed token in error: %v", err) + } + if !strings.Contains(err.Error(), "sendMessage") { + t.Fatalf("callAPI() error = %q, want method context", err) + } + if !errors.Is(err, transportErr) { + t.Fatalf("callAPI() error does not unwrap transport error: %v", err) + } +} + +func TestTelegramSendUsesReplyParameters(t *testing.T) { + request := make(chan map[string]any, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/sendMessage") { + t.Errorf("unexpected Telegram API path: %s", r.URL.Path) + http.Error(w, "unexpected path", http.StatusNotFound) + return + } + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode request: %v", err) + http.Error(w, "invalid payload", http.StatusBadRequest) + return + } + request <- payload + _, _ = w.Write([]byte(`{"ok":true,"result":{"message_id":1}}`)) + })) + defer server.Close() + + ch := newTelegramChannel(telegramAccount{ + AccountID: "account-1", + Token: "token", + APIBaseURL: server.URL, + }) + if err := ch.Send(context.Background(), core.OutgoingMessage{ + ChatID: "-100", + Text: "answer", + ReplyToMessageID: "12", + }); err != nil { + t.Fatalf("Send() error = %v", err) + } + + payload := <-request + if payload["chat_id"] != float64(-100) || payload["text"] != "answer" { + t.Fatalf("unexpected send payload: %+v", payload) + } + reply, ok := payload["reply_parameters"].(map[string]any) + if !ok || reply["message_id"] != float64(12) || reply["allow_sending_without_reply"] != true { + t.Fatalf("unexpected reply parameters: %+v", payload["reply_parameters"]) + } +} + +func TestTelegramSendIgnoresInvalidChatID(t *testing.T) { + ch := newTelegramChannel(telegramAccount{AccountID: "account-1", Token: "token"}) + if err := ch.Send(context.Background(), core.OutgoingMessage{ChatID: "invalid", Text: "answer"}); err != nil { + t.Fatalf("Send() error = %v, want nil", err) + } +} + +func TestTelegramStartAndStopPollsUpdates(t *testing.T) { + updateReturned := false + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.HasSuffix(r.URL.Path, "/deleteWebhook"): + _, _ = w.Write([]byte(`{"ok":true,"result":true}`)) + case strings.HasSuffix(r.URL.Path, "/getUpdates"): + if !updateReturned { + updateReturned = true + _, _ = w.Write([]byte(`{"ok":true,"result":[{"update_id":9,"message":{"message_id":3,"from":{"id":4,"is_bot":false},"chat":{"id":5,"type":"private"},"text":"hello"}}]}`)) + return + } + _, _ = w.Write([]byte(`{"ok":true,"result":[]}`)) + case strings.HasSuffix(r.URL.Path, "/sendMessage"): + _, _ = w.Write([]byte(`{"ok":true,"result":{"message_id":1}}`)) + default: + t.Errorf("unexpected Telegram API path: %s", r.URL.Path) + http.Error(w, "unexpected path", http.StatusNotFound) + } + })) + defer server.Close() + + ch := newTelegramChannel(telegramAccount{ + AccountID: "account-1", + Token: "token", + APIBaseURL: server.URL, + }) + message := make(chan core.IncomingMessage, 1) + ch.SetMessageHandler(func(_ context.Context, msg core.IncomingMessage) error { + message <- msg + return nil + }) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := ch.Start(ctx); err != nil { + t.Fatalf("Start() error = %v", err) + } + select { + case msg := <-message: + if msg.Text != "hello" || msg.MessageID != "3" || ch.offset != 10 { + t.Fatalf("unexpected polled message: %+v, offset=%d", msg, ch.offset) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for polled message") + } + + if err := ch.Stop(context.Background()); err != nil { + t.Fatalf("Stop() error = %v", err) + } +} diff --git a/internal/channels/wecom.go b/internal/channels/wecom.go index a4ce67278c..678c62a6c3 100644 --- a/internal/channels/wecom.go +++ b/internal/channels/wecom.go @@ -15,3 +15,740 @@ // package channels + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "net/url" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + + "ragflow/internal/channels/core" +) + +const ( + defaultWeComAPIBaseURL = "https://qyapi.weixin.qq.com/cgi-bin" + defaultWeComWSURL = "wss://openws.work.weixin.qq.com" + defaultWeComWebhookHost = "0.0.0.0" + defaultWeComWebhookPort = 3002 + weComHTTPTimeout = 30 * time.Second + weComWSHandshakeTimeout = 10 * time.Second + weComSubscribeTimeout = 10 * time.Second + weComWSWriteTimeout = 10 * time.Second + weComReconnectDelay = 3 * time.Second + weComHeartbeatInterval = 25 * time.Second + weComTokenSafetyMargin = 60 * time.Second +) + +var ( + errWeComAuthentication = errors.New("WeCom websocket authentication failed") + errWeComDisconnected = errors.New("WeCom websocket disconnected by server event") + weComRequestSequence atomic.Uint64 +) + +type wecomAccount struct { + AccountID string + ConnectionType string + CorpID string + AgentID int64 + Secret string + Token string + AESKey string + BotID string + WebhookHost string + WebhookPort int + APIBaseURL string + WSURL string +} + +type wecomChannel struct { + account wecomAccount + crypto *wecomCrypto + client *http.Client + + lifecycleMu sync.Mutex + mu sync.Mutex + handler core.MessageHandler + cancel context.CancelFunc + done chan struct{} + ws *websocket.Conn + server *wecomWebhookServer + wsWriteMu sync.Mutex + tokenMu sync.Mutex + accessToken string + tokenExpiry time.Time +} + +type wecomAPIResponse struct { + ErrCode int `json:"errcode"` + ErrMsg string `json:"errmsg"` +} + +type wecomAccessTokenResponse struct { + wecomAPIResponse + AccessToken string `json:"access_token"` + ExpiresIn int `json:"expires_in"` +} + +// newWeComChannel creates a WeCom channel with default endpoints filled in. +func newWeComChannel(account wecomAccount) (*wecomChannel, error) { + account.ConnectionType = strings.ToLower(strings.TrimSpace(account.ConnectionType)) + if account.ConnectionType == "" { + account.ConnectionType = "webhook" + } + account.APIBaseURL = strings.TrimRight(strings.TrimSpace(account.APIBaseURL), "/") + if account.APIBaseURL == "" { + account.APIBaseURL = defaultWeComAPIBaseURL + } + account.WSURL = strings.TrimSpace(account.WSURL) + if account.WSURL == "" { + account.WSURL = defaultWeComWSURL + } + account.WebhookHost = strings.TrimSpace(account.WebhookHost) + if account.WebhookHost == "" { + account.WebhookHost = defaultWeComWebhookHost + } + if account.WebhookPort == 0 { + account.WebhookPort = defaultWeComWebhookPort + } + + channel := &wecomChannel{ + account: account, + client: &http.Client{Timeout: weComHTTPTimeout}, + } + if account.ConnectionType == "webhook" { + crypto, err := newWeComCrypto(account.Token, account.AESKey, account.CorpID) + if err != nil { + return nil, err + } + channel.crypto = crypto + } + return channel, nil +} + +// newWeComChannelFromConfig builds a WeCom channel from chat_channel.config.credential. +func newWeComChannelFromConfig(accountID string, cfg map[string]any) (*wecomChannel, error) { + connectionType := strings.ToLower(firstString(cfg, "connection_type")) + if connectionType == "" { + connectionType = "webhook" + } + if connectionType != "websocket" { + // Python treats every non-websocket value as webhook. Keep the same + // effective runtime path while still applying webhook validation. + connectionType = "webhook" + } + + account := wecomAccount{ + AccountID: accountID, + ConnectionType: connectionType, + Secret: firstString(cfg, "secret"), + WebhookHost: firstString(cfg, "webhook_host"), + APIBaseURL: firstString(cfg, "api_base_url"), + WSURL: firstString(cfg, "ws_url"), + } + if account.Secret == "" { + return nil, fmt.Errorf("wecom account %q missing required field: secret", accountID) + } + + if connectionType == "websocket" { + account.BotID = firstString(cfg, "bot_id") + if account.BotID == "" { + return nil, fmt.Errorf("wecom account %q missing required field: bot_id", accountID) + } + return newWeComChannel(account) + } + + account.CorpID = firstString(cfg, "corp_id") + account.Token = firstString(cfg, "token") + account.AESKey = firstString(cfg, "aes_key") + missing := make([]string, 0, 3) + if account.CorpID == "" { + missing = append(missing, "corp_id") + } + if account.Token == "" { + missing = append(missing, "token") + } + if account.AESKey == "" { + missing = append(missing, "aes_key") + } + if len(missing) > 0 { + return nil, fmt.Errorf("wecom account %q missing required fields: %s", accountID, strings.Join(missing, ", ")) + } + if len(account.AESKey) != 43 { + return nil, fmt.Errorf("wecom account %q aes_key (EncodingAESKey) must be 43 characters, got %d", accountID, len(account.AESKey)) + } + + agentID, err := parseWeComInt(cfg["agent_id"], "agent_id") + if err != nil { + return nil, fmt.Errorf("wecom account %q %w", accountID, err) + } + account.AgentID = agentID + var configuredWebhookPort *int + if raw, ok := cfg["webhook_port"]; ok { + port, err := parseWeComInt(raw, "webhook_port") + if err != nil { + return nil, fmt.Errorf("wecom account %q %w", accountID, err) + } + account.WebhookPort = int(port) + configuredWebhookPort = &account.WebhookPort + } + channel, err := newWeComChannel(account) + if err != nil { + return nil, err + } + if configuredWebhookPort != nil { + channel.account.WebhookPort = *configuredWebhookPort + } + return channel, nil +} + +// ChannelID returns the platform identifier used by the chat-channel runtime. +func (c *wecomChannel) ChannelID() string { + return "wecom" +} + +// AccountID returns the chat_channel.id bound to this WeCom runtime instance. +func (c *wecomChannel) AccountID() string { + return c.account.AccountID +} + +// SetMessageHandler installs the RAGFlow bridge invoked for inbound messages. +func (c *wecomChannel) SetMessageHandler(handler core.MessageHandler) { + c.mu.Lock() + defer c.mu.Unlock() + c.handler = handler +} + +// Start starts either the shared webhook listener or the WebSocket client. +func (c *wecomChannel) Start(ctx context.Context) error { + c.lifecycleMu.Lock() + defer c.lifecycleMu.Unlock() + + if c.account.ConnectionType == "webhook" { + c.mu.Lock() + started := c.server != nil + c.mu.Unlock() + if started { + return nil + } + server, err := acquireWeComWebhookServer(c) + if err != nil { + return err + } + c.mu.Lock() + c.server = server + c.mu.Unlock() + log.Printf("[wecom:%s] registered webhook at /wecom/%s/callback", c.account.AccountID, c.account.AccountID) + return nil + } + + c.mu.Lock() + if c.cancel != nil { + c.mu.Unlock() + return nil + } + runCtx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + c.cancel = cancel + c.done = done + c.mu.Unlock() + + go func() { + defer close(done) + c.runWebSocket(runCtx) + c.mu.Lock() + if c.done == done { + c.cancel = nil + c.ws = nil + } + c.mu.Unlock() + }() + return nil +} + +// Stop releases the webhook listener or stops the WebSocket reconnect loop. +func (c *wecomChannel) Stop(ctx context.Context) error { + c.lifecycleMu.Lock() + defer c.lifecycleMu.Unlock() + + if c.account.ConnectionType == "webhook" { + c.mu.Lock() + server := c.server + c.server = nil + c.mu.Unlock() + if server != nil { + return releaseWeComWebhookServer(ctx, server, c.account.AccountID) + } + return nil + } + + c.mu.Lock() + cancel := c.cancel + done := c.done + conn := c.ws + c.mu.Unlock() + if cancel == nil { + return nil + } + cancel() + if conn != nil { + _ = conn.Close() + } + select { + case <-done: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// Send sends an application message or a WebSocket bot message. +func (c *wecomChannel) Send(ctx context.Context, msg core.OutgoingMessage) error { + if c.account.ConnectionType == "websocket" { + if err := c.sendWebSocketMessage(msg); err != nil { + log.Printf("[wecom:%s] websocket send failed: %v", c.account.AccountID, err) + } + return nil + } + if err := c.sendApplicationMessage(ctx, msg); err != nil { + log.Printf("[wecom:%s] application send failed: %v", c.account.AccountID, err) + } + return nil +} + +func (c *wecomChannel) runWebSocket(ctx context.Context) { + for ctx.Err() == nil { + err := c.runWebSocketOnce(ctx) + if ctx.Err() != nil { + return + } + if errors.Is(err, errWeComAuthentication) || errors.Is(err, errWeComDisconnected) { + log.Printf("[wecom:%s] websocket stopped: %v", c.account.AccountID, err) + return + } + if err != nil { + log.Printf("[wecom:%s] websocket loop error: %v", c.account.AccountID, err) + } + if !waitForWeComRetry(ctx) { + return + } + } +} + +func (c *wecomChannel) runWebSocketOnce(ctx context.Context) error { + dialer := *websocket.DefaultDialer + dialer.HandshakeTimeout = weComWSHandshakeTimeout + conn, _, err := dialer.DialContext(ctx, c.account.WSURL, nil) + if err != nil { + return err + } + defer conn.Close() + + c.mu.Lock() + c.ws = conn + c.mu.Unlock() + defer func() { + c.mu.Lock() + if c.ws == conn { + c.ws = nil + } + c.mu.Unlock() + }() + + connectionDone := make(chan struct{}) + defer close(connectionDone) + go func() { + select { + case <-ctx.Done(): + _ = conn.Close() + case <-connectionDone: + } + }() + + if err := c.subscribeWebSocket(conn); err != nil { + return err + } + heartbeatDone := make(chan struct{}) + defer close(heartbeatDone) + go c.runWebSocketHeartbeat(ctx, conn, heartbeatDone) + + for ctx.Err() == nil { + if err := conn.SetReadDeadline(time.Now().Add(2 * weComHeartbeatInterval)); err != nil { + return err + } + _, payload, err := conn.ReadMessage() + if err != nil { + return err + } + disconnect, err := c.handleWebSocketPayload(ctx, payload) + if err != nil { + log.Printf("[wecom:%s] invalid websocket payload: %v", c.account.AccountID, err) + continue + } + if disconnect { + return errWeComDisconnected + } + } + return ctx.Err() +} + +func (c *wecomChannel) subscribeWebSocket(conn *websocket.Conn) error { + payload := map[string]any{ + "cmd": "aibot_subscribe", + "headers": map[string]any{"req_id": newWeComRequestID()}, + "body": map[string]any{ + "bot_id": c.account.BotID, + "secret": c.account.Secret, + }, + } + if err := c.writeWebSocketJSON(conn, payload); err != nil { + return err + } + _ = conn.SetReadDeadline(time.Now().Add(weComSubscribeTimeout)) + defer conn.SetReadDeadline(time.Time{}) + var response map[string]any + if err := conn.ReadJSON(&response); err != nil { + return fmt.Errorf("WeCom websocket subscribe acknowledgement: %w", err) + } + if cmd := weComRawStringValue(response["cmd"]); cmd != "aibot_subscribe" { + log.Printf("[wecom:%s] unexpected subscribe response command %q", c.account.AccountID, cmd) + } + errCode := weComIntValue(response["errcode"]) + if errCode == 853000 { + return fmt.Errorf("%w: invalid bot_id or secret", errWeComAuthentication) + } + if errCode != 0 { + return fmt.Errorf("WeCom websocket subscribe failed: errcode=%d errmsg=%s", errCode, weComStringValue(response["errmsg"])) + } + log.Printf("[wecom:%s] websocket subscribed", c.account.AccountID) + return nil +} + +func (c *wecomChannel) runWebSocketHeartbeat(ctx context.Context, conn *websocket.Conn, done <-chan struct{}) { + ticker := time.NewTicker(weComHeartbeatInterval) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-done: + return + case <-ticker.C: + if err := c.writeWebSocketJSON(conn, map[string]any{ + "cmd": "ping", + "headers": map[string]any{"req_id": newWeComRequestID()}, + }); err != nil { + _ = conn.Close() + return + } + } + } +} + +func (c *wecomChannel) handleWebSocketPayload(ctx context.Context, payload []byte) (bool, error) { + var object map[string]any + if err := json.Unmarshal(payload, &object); err != nil { + return false, err + } + cmd := weComRawStringValue(object["cmd"]) + switch cmd { + case "aibot_msg_callback": + c.handleWebSocketMessage(ctx, object) + case "aibot_event_callback": + body := weComMapValue(object["body"]) + event := weComMapValue(body["event"]) + if weComRawStringValue(event["eventtype"]) == "disconnected_event" { + return true, nil + } + case "aibot_subscribe", "aibot_respond_msg", "aibot_respond_welcome_msg", "aibot_send_msg", "ping": + if code := weComIntValue(object["errcode"]); code != 0 { + log.Printf("[wecom:%s] websocket response error: %s", c.account.AccountID, string(payload)) + } + } + return false, nil +} + +func (c *wecomChannel) handleWebSocketMessage(ctx context.Context, raw map[string]any) { + headers := weComMapValue(raw["headers"]) + body := weComMapValue(raw["body"]) + if weComRawStringValue(body["msgtype"]) != "text" { + return + } + sender := weComMapValue(body["from"]) + senderID := weComRawStringValue(sender["userid"]) + if senderID == "" { + return + } + chatID := weComRawStringValue(body["chatid"]) + if chatID == "" { + chatID = senderID + } + messageID := weComRawStringValue(headers["req_id"]) + if messageID == "" { + messageID = weComRawStringValue(body["msgid"]) + } + chatType := "p2p" + if weComRawStringValue(body["chattype"]) == "group" { + chatType = "group" + } + text := weComRawStringValue(weComMapValue(body["text"])["content"]) + c.handleTextMessage(ctx, chatID, senderID, messageID, text, chatType, raw) +} + +func (c *wecomChannel) sendWebSocketMessage(msg core.OutgoingMessage) error { + if msg.Text == "" { + return errors.New("WeCom websocket reply text is required") + } + c.mu.Lock() + conn := c.ws + c.mu.Unlock() + if conn == nil { + return errors.New("WeCom websocket is not connected") + } + return c.writeWebSocketJSON(conn, map[string]any{ + "cmd": "aibot_send_msg", + "headers": map[string]any{"req_id": newWeComRequestID()}, + "body": map[string]any{ + "chatid": msg.ChatID, + "msgtype": "markdown", + "markdown": map[string]any{ + "content": msg.Text, + }, + }, + }) +} + +func (c *wecomChannel) writeWebSocketJSON(conn *websocket.Conn, payload any) error { + c.wsWriteMu.Lock() + defer c.wsWriteMu.Unlock() + if err := conn.SetWriteDeadline(time.Now().Add(weComWSWriteTimeout)); err != nil { + return err + } + return conn.WriteJSON(payload) +} + +func (c *wecomChannel) sendApplicationMessage(ctx context.Context, msg core.OutgoingMessage) error { + if msg.ChatID == "" { + return errors.New("WeCom chat_id is required") + } + accessToken, err := c.getAccessToken(ctx) + if err != nil { + return err + } + endpoint, err := url.Parse(c.account.APIBaseURL + "/message/send") + if err != nil { + return err + } + query := endpoint.Query() + query.Set("access_token", accessToken) + endpoint.RawQuery = query.Encode() + + var response wecomAPIResponse + if err := c.requestJSON(ctx, http.MethodPost, endpoint.String(), map[string]any{ + "touser": msg.ChatID, + "msgtype": "text", + "agentid": c.account.AgentID, + "text": map[string]any{"content": msg.Text}, + "safe": 0, + }, &response); err != nil { + return err + } + if response.ErrCode != 0 { + if response.ErrCode == 40014 || response.ErrCode == 42001 { + c.clearAccessToken() + } + return fmt.Errorf("WeCom message/send failed: errcode=%d errmsg=%s", response.ErrCode, response.ErrMsg) + } + return nil +} + +func (c *wecomChannel) getAccessToken(ctx context.Context) (string, error) { + c.tokenMu.Lock() + defer c.tokenMu.Unlock() + if c.accessToken != "" && time.Now().Before(c.tokenExpiry) { + return c.accessToken, nil + } + + endpoint, err := url.Parse(c.account.APIBaseURL + "/gettoken") + if err != nil { + return "", err + } + query := endpoint.Query() + query.Set("corpid", c.account.CorpID) + query.Set("corpsecret", c.account.Secret) + endpoint.RawQuery = query.Encode() + var response wecomAccessTokenResponse + if err := c.requestJSON(ctx, http.MethodGet, endpoint.String(), nil, &response); err != nil { + return "", err + } + if response.ErrCode != 0 || response.AccessToken == "" { + return "", fmt.Errorf("WeCom gettoken failed: errcode=%d errmsg=%s", response.ErrCode, response.ErrMsg) + } + expiresIn := response.ExpiresIn + if expiresIn == 0 { + expiresIn = 7200 + } + c.accessToken = response.AccessToken + c.tokenExpiry = time.Now().Add(time.Duration(expiresIn)*time.Second - weComTokenSafetyMargin) + return c.accessToken, nil +} + +func (c *wecomChannel) clearAccessToken() { + c.tokenMu.Lock() + defer c.tokenMu.Unlock() + c.accessToken = "" + c.tokenExpiry = time.Time{} +} + +func (c *wecomChannel) requestJSON(ctx context.Context, method, endpoint string, body map[string]any, out any) error { + var reader io.Reader + if body != nil { + payload, err := json.Marshal(body) + if err != nil { + return err + } + reader = bytes.NewReader(payload) + } + req, err := http.NewRequestWithContext(ctx, method, endpoint, reader) + if err != nil { + return err + } + if body != nil { + req.Header.Set("Content-Type", "application/json") + } + resp, err := c.client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + payload, err := io.ReadAll(io.LimitReader(resp.Body, 1024*1024)) + if err != nil { + return err + } + if out == nil { + return nil + } + if err := json.Unmarshal(payload, out); err != nil { + return fmt.Errorf("WeCom API returned invalid JSON: %w", err) + } + return nil +} + +func (c *wecomChannel) handleTextMessage(ctx context.Context, chatID, senderID, messageID, text, chatType string, raw map[string]any) { + if strings.TrimSpace(text) == "" { + return + } + c.mu.Lock() + handler := c.handler + c.mu.Unlock() + if handler == nil { + return + } + if err := handler(ctx, core.IncomingMessage{ + Channel: c.ChannelID(), + AccountID: c.account.AccountID, + ChatID: chatID, + ChatType: chatType, + MessageID: messageID, + SenderID: senderID, + Text: text, + Raw: raw, + }); err != nil { + log.Printf("[wecom:%s] message handler error: %v", c.account.AccountID, err) + } +} + +func parseWeComInt(value any, field string) (int64, error) { + switch v := value.(type) { + case int: + return int64(v), nil + case int64: + return v, nil + case float64: + if v != float64(int64(v)) { + return 0, fmt.Errorf("%s must be an integer", field) + } + return int64(v), nil + case json.Number: + result, err := v.Int64() + if err != nil { + return 0, fmt.Errorf("%s must be an integer", field) + } + return result, nil + case string: + result, err := strconv.ParseInt(strings.TrimSpace(v), 10, 64) + if err != nil { + return 0, fmt.Errorf("%s must be an integer", field) + } + return result, nil + default: + return 0, fmt.Errorf("%s must be an integer", field) + } +} + +func weComMapValue(value any) map[string]any { + result, _ := value.(map[string]any) + return result +} + +func weComStringValue(value any) string { + if value == nil { + return "" + } + return strings.TrimSpace(fmt.Sprint(value)) +} + +// weComRawStringValue mirrors Python's str(value or "") conversion for +// message fields without changing user-visible whitespace. +func weComRawStringValue(value any) string { + if value == nil { + return "" + } + if text, ok := value.(string); ok { + return text + } + return fmt.Sprint(value) +} + +func weComIntValue(value any) int { + switch v := value.(type) { + case float64: + return int(v) + case int: + return v + case json.Number: + result, _ := v.Int64() + return int(result) + case string: + result, _ := strconv.Atoi(v) + return result + default: + return 0 + } +} + +func newWeComRequestID() string { + return fmt.Sprintf("req-%d-%d", time.Now().UnixNano(), weComRequestSequence.Add(1)) +} + +func waitForWeComRetry(ctx context.Context) bool { + timer := time.NewTimer(weComReconnectDelay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} diff --git a/internal/channels/wecom_test.go b/internal/channels/wecom_test.go new file mode 100644 index 0000000000..ddacb79305 --- /dev/null +++ b/internal/channels/wecom_test.go @@ -0,0 +1,524 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package channels + +import ( + "bytes" + "context" + "crypto/aes" + "crypto/cipher" + "crypto/sha1" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "encoding/json" + "encoding/xml" + "net" + "net/http" + "net/http/httptest" + "net/url" + "sort" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" + + "ragflow/internal/channels/core" +) + +func TestNewWeComChannelFromConfigValidatesConnectionFields(t *testing.T) { + if _, err := newWeComChannelFromConfig("account-1", map[string]any{ + "connection_type": "websocket", + "secret": "secret", + }); err == nil { + t.Fatal("websocket config accepted a missing bot_id") + } + + channel, err := newWeComChannelFromConfig("account-1", map[string]any{ + "connection_type": "websocket", + "bot_id": "bot-id", + "secret": "secret", + }) + if err != nil { + t.Fatalf("valid websocket config error = %v", err) + } + if channel.ChannelID() != "wecom" || channel.AccountID() != "account-1" { + t.Fatalf("unexpected channel identity: %s:%s", channel.ChannelID(), channel.AccountID()) + } + + if _, err := newWeComChannelFromConfig("account-2", map[string]any{ + "corp_id": "corp-id", + "agent_id": "1000001", + "secret": "secret", + "token": "token", + "aes_key": "too-short", + }); err == nil { + t.Fatal("webhook config accepted an invalid aes_key") + } +} + +func TestNewWeComChannelFromConfigTreatsUnknownConnectionTypeAsWebhook(t *testing.T) { + encodingAESKey, _ := testWeComAESKey() + channel, err := newWeComChannelFromConfig("account-1", map[string]any{ + "connection_type": "unexpected", + "corp_id": "corp-id", + "agent_id": "1000001", + "secret": "secret", + "token": "token", + "aes_key": encodingAESKey, + }) + if err != nil { + t.Fatalf("unknown connection type error = %v", err) + } + if channel.account.ConnectionType != "webhook" || channel.crypto == nil { + t.Fatalf("unknown connection type did not use webhook mode: %+v", channel.account) + } +} + +func TestNewWeComChannelFromConfigAcceptsIntegerAgentID(t *testing.T) { + encodingAESKey, _ := testWeComAESKey() + channel, err := newWeComChannelFromConfig("account-1", map[string]any{ + "corp_id": "corp-id", + "agent_id": 0, + "secret": "secret", + "token": "token", + "aes_key": encodingAESKey, + "webhook_port": 0, + }) + if err != nil { + t.Fatalf("agent_id=0 error = %v", err) + } + if channel.account.AgentID != 0 { + t.Fatalf("agent_id = %d, want 0", channel.account.AgentID) + } + if channel.account.WebhookPort != 0 { + t.Fatalf("webhook_port = %d, want 0", channel.account.WebhookPort) + } +} + +func TestWeComCryptoDecryptsAndValidatesReceiveID(t *testing.T) { + encodingAESKey, aesKey := testWeComAESKey() + crypto, err := newWeComCrypto("token", encodingAESKey, "corp-id") + if err != nil { + t.Fatalf("newWeComCrypto() error = %v", err) + } + encrypted := encryptWeComTestPayload(t, aesKey, []byte("hello"), "corp-id") + signature := testWeComSignature("token", "123", "nonce", encrypted) + + plaintext, err := crypto.decrypt(signature, "123", "nonce", encrypted) + if err != nil { + t.Fatalf("decrypt() error = %v", err) + } + if string(plaintext) != "hello" { + t.Fatalf("decrypt() = %q, want hello", plaintext) + } + if _, err := crypto.decrypt("bad-signature", "123", "nonce", encrypted); err == nil { + t.Fatal("decrypt() accepted an invalid signature") + } + + wrongReceiveID := encryptWeComTestPayload(t, aesKey, []byte("hello"), "other-corp") + wrongSignature := testWeComSignature("token", "123", "nonce", wrongReceiveID) + if _, err := crypto.decrypt(wrongSignature, "123", "nonce", wrongReceiveID); err == nil { + t.Fatal("decrypt() accepted a mismatched receive ID") + } +} + +func TestWeComWebhookHandlesVerificationAndTextMessage(t *testing.T) { + encodingAESKey, aesKey := testWeComAESKey() + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "webhook", + CorpID: "corp-id", + AgentID: 1000001, + Secret: "secret", + Token: "token", + AESKey: encodingAESKey, + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + messages := make(chan core.IncomingMessage, 1) + channel.SetMessageHandler(func(_ context.Context, msg core.IncomingMessage) error { + messages <- msg + return nil + }) + server := &wecomWebhookServer{channels: map[string]*wecomChannel{"account-1": channel}} + + echo := encryptWeComTestPayload(t, aesKey, []byte("verified"), "corp-id") + echoQuery := url.Values{ + "msg_signature": {testWeComSignature("token", "123", "nonce", echo)}, + "timestamp": {"123"}, + "nonce": {"nonce"}, + "echostr": {echo}, + } + echoRequest := httptest.NewRequest(http.MethodGet, "/wecom/account-1/callback?"+echoQuery.Encode(), nil) + echoResponse := httptest.NewRecorder() + server.handleRequest(echoResponse, echoRequest) + if echoResponse.Code != http.StatusOK || echoResponse.Body.String() != "verified" { + t.Fatalf("verification response = %d %q", echoResponse.Code, echoResponse.Body.String()) + } + + plaintext := []byte(`corp-iduser-1123texthello wecommessage-1`) + encrypted := encryptWeComTestPayload(t, aesKey, plaintext, "corp-id") + body, err := xml.Marshal(wecomEncryptedXML{Encrypt: encrypted}) + if err != nil { + t.Fatalf("marshal encrypted XML: %v", err) + } + query := url.Values{ + "msg_signature": {testWeComSignature("token", "124", "nonce-2", encrypted)}, + "timestamp": {"124"}, + "nonce": {"nonce-2"}, + } + request := httptest.NewRequest(http.MethodPost, "/wecom/account-1/callback?"+query.Encode(), bytes.NewReader(body)) + response := httptest.NewRecorder() + server.handleRequest(response, request) + if response.Code != http.StatusOK { + t.Fatalf("message response status = %d, body=%q", response.Code, response.Body.String()) + } + + select { + case msg := <-messages: + if msg.ChatID != "user-1" || msg.SenderID != "user-1" || msg.MessageID != "message-1" || msg.Text != "hello wecom" || msg.ChatType != "p2p" { + t.Fatalf("unexpected webhook message: %+v", msg) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for webhook message") + } +} + +func TestWeComWebhookAcceptsMalformedEncryptedMessage(t *testing.T) { + encodingAESKey, _ := testWeComAESKey() + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "webhook", + CorpID: "corp-id", + AgentID: 1000001, + Secret: "secret", + Token: "token", + AESKey: encodingAESKey, + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + server := &wecomWebhookServer{channels: map[string]*wecomChannel{"account-1": channel}} + request := httptest.NewRequest(http.MethodPost, "/wecom/account-1/callback", strings.NewReader("not xml")) + response := httptest.NewRecorder() + server.handleRequest(response, request) + if response.Code != http.StatusOK { + t.Fatalf("malformed webhook response = %d, want 200", response.Code) + } +} + +func TestWeComWebhookServerSharesListener(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("reserve webhook port: %v", err) + } + port := listener.Addr().(*net.TCPAddr).Port + if err := listener.Close(); err != nil { + t.Fatalf("release webhook port: %v", err) + } + + encodingAESKey, _ := testWeComAESKey() + newChannel := func(accountID string) *wecomChannel { + channel, err := newWeComChannel(wecomAccount{ + AccountID: accountID, + ConnectionType: "webhook", + CorpID: "corp-id", + AgentID: 1000001, + Secret: "secret", + Token: "token", + AESKey: encodingAESKey, + WebhookHost: "127.0.0.1", + WebhookPort: port, + }) + if err != nil { + t.Fatalf("newWeComChannel(%s) error = %v", accountID, err) + } + return channel + } + first := newChannel("account-1") + second := newChannel("account-2") + if err := first.Start(context.Background()); err != nil { + t.Fatalf("first Start() error = %v", err) + } + defer first.Stop(context.Background()) + if err := second.Start(context.Background()); err != nil { + t.Fatalf("second Start() error = %v", err) + } + defer second.Stop(context.Background()) + + first.mu.Lock() + firstServer := first.server + first.mu.Unlock() + second.mu.Lock() + secondServer := second.server + second.mu.Unlock() + if firstServer == nil || firstServer != secondServer { + t.Fatal("webhook channels did not share one listener") + } + firstServer.mu.RLock() + refs := firstServer.refs + registered := len(firstServer.channels) + firstServer.mu.RUnlock() + if refs != 2 || registered != 2 { + t.Fatalf("shared server refs=%d channels=%d, want 2 and 2", refs, registered) + } + + if err := first.Stop(context.Background()); err != nil { + t.Fatalf("first Stop() error = %v", err) + } + firstServer.mu.RLock() + refs = firstServer.refs + registered = len(firstServer.channels) + firstServer.mu.RUnlock() + if refs != 1 || registered != 1 { + t.Fatalf("shared server after first stop refs=%d channels=%d, want 1 and 1", refs, registered) + } + if err := second.Stop(context.Background()); err != nil { + t.Fatalf("second Stop() error = %v", err) + } +} + +func TestWeComApplicationSendCachesAccessToken(t *testing.T) { + getTokenCalls := 0 + sendCalls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/gettoken": + getTokenCalls++ + if r.URL.Query().Get("corpid") != "corp-id" || r.URL.Query().Get("corpsecret") != "secret" { + t.Errorf("unexpected gettoken query: %s", r.URL.RawQuery) + } + _, _ = w.Write([]byte(`{"errcode":0,"errmsg":"ok","access_token":"access-token","expires_in":7200}`)) + case "/message/send": + sendCalls++ + if r.URL.Query().Get("access_token") != "access-token" { + t.Errorf("unexpected access_token query: %s", r.URL.RawQuery) + } + var payload map[string]any + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Errorf("decode send payload: %v", err) + http.Error(w, "invalid payload", http.StatusBadRequest) + return + } + if payload["touser"] != "user-1" || payload["agentid"] != float64(1000001) || weComMapValue(payload["text"])["content"] != "answer" { + t.Errorf("unexpected send payload: %+v", payload) + } + _, _ = w.Write([]byte(`{"errcode":0,"errmsg":"ok"}`)) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + encodingAESKey, _ := testWeComAESKey() + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "webhook", + CorpID: "corp-id", + AgentID: 1000001, + Secret: "secret", + Token: "token", + AESKey: encodingAESKey, + APIBaseURL: server.URL, + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + message := core.OutgoingMessage{ChatID: "user-1", Text: "answer"} + if err := channel.Send(context.Background(), message); err != nil { + t.Fatalf("first Send() error = %v", err) + } + if err := channel.Send(context.Background(), message); err != nil { + t.Fatalf("second Send() error = %v", err) + } + if getTokenCalls != 1 || sendCalls != 2 { + t.Fatalf("gettoken calls = %d, send calls = %d", getTokenCalls, sendCalls) + } +} + +func TestWeComWebSocketSubscribesHandlesAndSends(t *testing.T) { + upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} + outgoing := make(chan map[string]any, 1) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + t.Errorf("upgrade websocket: %v", err) + return + } + defer conn.Close() + + var subscribe map[string]any + if err := conn.ReadJSON(&subscribe); err != nil { + t.Errorf("read subscribe: %v", err) + return + } + if subscribe["cmd"] != "aibot_subscribe" || weComMapValue(subscribe["body"])["bot_id"] != "bot-id" { + t.Errorf("unexpected subscribe payload: %+v", subscribe) + } + if err := conn.WriteJSON(map[string]any{"cmd": "aibot_subscribe", "errcode": 0, "errmsg": "ok"}); err != nil { + t.Errorf("write subscribe response: %v", err) + return + } + if err := conn.WriteJSON(map[string]any{ + "cmd": "aibot_msg_callback", + "headers": map[string]any{"req_id": "request-1"}, + "body": map[string]any{ + "msgid": "message-1", + "msgtype": "text", + "chattype": "group", + "chatid": "chat-1", + "from": map[string]any{"userid": "user-1"}, + "text": map[string]any{"content": "question"}, + }, + }); err != nil { + t.Errorf("write callback: %v", err) + return + } + var response map[string]any + if err := conn.ReadJSON(&response); err != nil { + t.Errorf("read outgoing response: %v", err) + return + } + outgoing <- response + })) + defer server.Close() + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "websocket", + BotID: "bot-id", + Secret: "secret", + WSURL: wsURL, + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + incoming := make(chan core.IncomingMessage, 1) + channel.SetMessageHandler(func(ctx context.Context, msg core.IncomingMessage) error { + incoming <- msg + return channel.Send(ctx, core.OutgoingMessage{ChatID: msg.ChatID, Text: "answer", ReplyToMessageID: msg.MessageID}) + }) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if err := channel.Start(ctx); err != nil { + t.Fatalf("Start() error = %v", err) + } + + select { + case msg := <-incoming: + if msg.ChatID != "chat-1" || msg.ChatType != "group" || msg.SenderID != "user-1" || msg.MessageID != "request-1" || msg.Text != "question" { + t.Fatalf("unexpected websocket message: %+v", msg) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for websocket message") + } + select { + case response := <-outgoing: + body := weComMapValue(response["body"]) + if response["cmd"] != "aibot_send_msg" || body["chatid"] != "chat-1" || weComMapValue(body["markdown"])["content"] != "answer" { + t.Fatalf("unexpected websocket response: %+v", response) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for websocket response") + } + if err := channel.Stop(context.Background()); err != nil { + t.Fatalf("Stop() error = %v", err) + } +} + +func TestWeComWebSocketMessagePreservesWhitespaceAndFallsBackToSender(t *testing.T) { + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "websocket", + BotID: "bot-id", + Secret: "secret", + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + messages := make(chan core.IncomingMessage, 1) + channel.SetMessageHandler(func(_ context.Context, msg core.IncomingMessage) error { + messages <- msg + return nil + }) + channel.handleWebSocketMessage(context.Background(), map[string]any{ + "headers": map[string]any{"req_id": "request-1"}, + "body": map[string]any{ + "msgtype": "text", + "from": map[string]any{"userid": "user-1"}, + "text": map[string]any{"content": " question\n"}, + }, + }) + select { + case message := <-messages: + if message.ChatID != "user-1" || message.Text != " question\n" { + t.Fatalf("unexpected websocket message: %+v", message) + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for websocket message") + } +} + +func TestWeComSendIgnoresTransportErrors(t *testing.T) { + channel, err := newWeComChannel(wecomAccount{ + AccountID: "account-1", + ConnectionType: "websocket", + BotID: "bot-id", + Secret: "secret", + }) + if err != nil { + t.Fatalf("newWeComChannel() error = %v", err) + } + if err := channel.Send(context.Background(), core.OutgoingMessage{ChatID: "chat-1", Text: "answer"}); err != nil { + t.Fatalf("Send() error = %v, want nil", err) + } +} + +func testWeComAESKey() (string, []byte) { + key := []byte("0123456789abcdef0123456789abcdef") + return strings.TrimSuffix(base64.StdEncoding.EncodeToString(key), "="), key +} + +func encryptWeComTestPayload(t *testing.T, key, message []byte, receiveID string) string { + t.Helper() + payload := make([]byte, 20+len(message)+len(receiveID)) + copy(payload[:16], []byte("0123456789abcdef")) + binary.BigEndian.PutUint32(payload[16:20], uint32(len(message))) + copy(payload[20:], message) + copy(payload[20+len(message):], receiveID) + padding := weComPKCS7BlockSize - len(payload)%weComPKCS7BlockSize + payload = append(payload, bytes.Repeat([]byte{byte(padding)}, padding)...) + + block, err := aes.NewCipher(key) + if err != nil { + t.Fatalf("create test cipher: %v", err) + } + ciphertext := make([]byte, len(payload)) + cipher.NewCBCEncrypter(block, key[:block.BlockSize()]).CryptBlocks(ciphertext, payload) + return base64.StdEncoding.EncodeToString(ciphertext) +} + +func testWeComSignature(token, timestamp, nonce, encrypted string) string { + parts := []string{token, timestamp, nonce, encrypted} + sort.Strings(parts) + hash := sha1.Sum([]byte(strings.Join(parts, ""))) + return hex.EncodeToString(hash[:]) +} diff --git a/internal/channels/wecom_webhook.go b/internal/channels/wecom_webhook.go new file mode 100644 index 0000000000..76d61e66af --- /dev/null +++ b/internal/channels/wecom_webhook.go @@ -0,0 +1,315 @@ +// +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// + +package channels + +import ( + "context" + "crypto/aes" + "crypto/cipher" + "crypto/sha1" + "crypto/subtle" + "encoding/base64" + "encoding/binary" + "encoding/hex" + "encoding/xml" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "net/url" + "sort" + "strconv" + "strings" + "sync" + "time" +) + +const ( + weComWebhookBodyLimit = 1024 * 1024 + weComWebhookShutdownTimeout = 5 * time.Second + weComPKCS7BlockSize = 32 +) + +type wecomCrypto struct { + token string + aesKey []byte + receiveID string +} + +type wecomEncryptedXML struct { + Encrypt string `xml:"Encrypt"` +} + +type wecomXMLMessage struct { + ToUserName string `xml:"ToUserName"` + FromUserName string `xml:"FromUserName"` + CreateTime int64 `xml:"CreateTime"` + MsgType string `xml:"MsgType"` + Content string `xml:"Content"` + MsgID string `xml:"MsgId"` +} + +type wecomWebhookServerKey struct { + host string + port int +} + +type wecomWebhookServer struct { + key wecomWebhookServerKey + server *http.Server + listener net.Listener + + mu sync.RWMutex + channels map[string]*wecomChannel + refs int +} + +var sharedWeComWebhookServers = struct { + sync.Mutex + servers map[wecomWebhookServerKey]*wecomWebhookServer +}{ + servers: map[wecomWebhookServerKey]*wecomWebhookServer{}, +} + +func newWeComCrypto(token, encodingAESKey, receiveID string) (*wecomCrypto, error) { + if len(encodingAESKey) != 43 { + return nil, fmt.Errorf("WeCom EncodingAESKey must be 43 characters, got %d", len(encodingAESKey)) + } + aesKey, err := base64.StdEncoding.DecodeString(encodingAESKey + "=") + if err != nil { + return nil, fmt.Errorf("decode WeCom EncodingAESKey: %w", err) + } + if len(aesKey) != 32 { + return nil, fmt.Errorf("decoded WeCom AES key must be 32 bytes, got %d", len(aesKey)) + } + return &wecomCrypto{ + token: token, + aesKey: aesKey, + receiveID: receiveID, + }, nil +} + +func (c *wecomCrypto) decrypt(signature, timestamp, nonce, encrypted string) ([]byte, error) { + if signature == "" || timestamp == "" || nonce == "" || encrypted == "" { + return nil, errors.New("missing WeCom signature parameters") + } + if !c.hasValidSignature(signature, timestamp, nonce, encrypted) { + return nil, errors.New("invalid WeCom message signature") + } + ciphertext, err := base64.StdEncoding.DecodeString(encrypted) + if err != nil { + return nil, fmt.Errorf("decode WeCom ciphertext: %w", err) + } + block, err := aes.NewCipher(c.aesKey) + if err != nil { + return nil, err + } + if len(ciphertext) == 0 || len(ciphertext)%block.BlockSize() != 0 { + return nil, errors.New("invalid WeCom ciphertext length") + } + plaintext := make([]byte, len(ciphertext)) + cipher.NewCBCDecrypter(block, c.aesKey[:block.BlockSize()]).CryptBlocks(plaintext, ciphertext) + plaintext, err = removeWeComPKCS7Padding(plaintext) + if err != nil { + return nil, err + } + if len(plaintext) < 20 { + return nil, errors.New("invalid WeCom decrypted payload") + } + messageLength := int(binary.BigEndian.Uint32(plaintext[16:20])) + messageEnd := 20 + messageLength + if messageLength < 0 || messageEnd > len(plaintext) { + return nil, errors.New("invalid WeCom decrypted message length") + } + if receiveID := string(plaintext[messageEnd:]); c.receiveID != "" && receiveID != c.receiveID { + return nil, fmt.Errorf("WeCom receive ID mismatch: got %q", receiveID) + } + return plaintext[20:messageEnd], nil +} + +func (c *wecomCrypto) hasValidSignature(signature, timestamp, nonce, encrypted string) bool { + parts := []string{c.token, timestamp, nonce, encrypted} + sort.Strings(parts) + hash := sha1.Sum([]byte(strings.Join(parts, ""))) + expected := hex.EncodeToString(hash[:]) + return subtle.ConstantTimeCompare([]byte(strings.ToLower(signature)), []byte(expected)) == 1 +} + +func removeWeComPKCS7Padding(payload []byte) ([]byte, error) { + if len(payload) == 0 { + return nil, errors.New("empty WeCom padded payload") + } + padding := int(payload[len(payload)-1]) + if padding < 1 || padding > weComPKCS7BlockSize || padding > len(payload) { + return nil, errors.New("invalid WeCom PKCS#7 padding") + } + for _, value := range payload[len(payload)-padding:] { + if int(value) != padding { + return nil, errors.New("invalid WeCom PKCS#7 padding bytes") + } + } + return payload[:len(payload)-padding], nil +} + +func acquireWeComWebhookServer(channel *wecomChannel) (*wecomWebhookServer, error) { + key := wecomWebhookServerKey{ + host: channel.account.WebhookHost, + port: channel.account.WebhookPort, + } + sharedWeComWebhookServers.Lock() + defer sharedWeComWebhookServers.Unlock() + + if server := sharedWeComWebhookServers.servers[key]; server != nil { + server.mu.Lock() + server.channels[channel.account.AccountID] = channel + server.refs++ + server.mu.Unlock() + return server, nil + } + + listener, err := net.Listen("tcp", net.JoinHostPort(key.host, strconv.Itoa(key.port))) + if err != nil { + return nil, fmt.Errorf("start WeCom webhook listener: %w", err) + } + server := &wecomWebhookServer{ + key: key, + listener: listener, + channels: map[string]*wecomChannel{ + channel.account.AccountID: channel, + }, + refs: 1, + } + mux := http.NewServeMux() + mux.HandleFunc("/wecom/", server.handleRequest) + server.server = &http.Server{ + Handler: mux, + ReadHeaderTimeout: 10 * time.Second, + } + sharedWeComWebhookServers.servers[key] = server + + go func() { + if err := server.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) { + log.Printf("[wecom] webhook server %s stopped: %v", listener.Addr(), err) + } + }() + log.Printf("[wecom] webhook listening on http://%s/wecom//callback", listener.Addr()) + return server, nil +} + +func releaseWeComWebhookServer(ctx context.Context, server *wecomWebhookServer, accountID string) error { + sharedWeComWebhookServers.Lock() + server.mu.Lock() + delete(server.channels, accountID) + if server.refs > 0 { + server.refs-- + } + remaining := server.refs + server.mu.Unlock() + if remaining > 0 { + sharedWeComWebhookServers.Unlock() + return nil + } + if sharedWeComWebhookServers.servers[server.key] == server { + delete(sharedWeComWebhookServers.servers, server.key) + } + sharedWeComWebhookServers.Unlock() + + shutdownCtx, cancel := context.WithTimeout(ctx, weComWebhookShutdownTimeout) + defer cancel() + return server.server.Shutdown(shutdownCtx) +} + +func (s *wecomWebhookServer) handleRequest(w http.ResponseWriter, r *http.Request) { + accountID, ok := weComAccountIDFromPath(r.URL.Path) + if !ok { + http.NotFound(w, r) + return + } + s.mu.RLock() + channel := s.channels[accountID] + s.mu.RUnlock() + if channel == nil { + http.Error(w, "unknown account", http.StatusNotFound) + return + } + + query := r.URL.Query() + signature := query.Get("msg_signature") + timestamp := query.Get("timestamp") + nonce := query.Get("nonce") + if r.Method == http.MethodGet { + plaintext, err := channel.crypto.decrypt(signature, timestamp, nonce, query.Get("echostr")) + if err != nil { + http.Error(w, "bad signature", http.StatusForbidden) + return + } + w.Header().Set("Content-Type", "text/plain; charset=utf-8") + _, _ = w.Write(plaintext) + return + } + if r.Method != http.MethodPost { + w.WriteHeader(http.StatusMethodNotAllowed) + return + } + + body, err := io.ReadAll(io.LimitReader(r.Body, weComWebhookBodyLimit)) + if err != nil { + log.Printf("[wecom:%s] failed to read request body: %v", accountID, err) + w.WriteHeader(http.StatusOK) + return + } + var envelope wecomEncryptedXML + if err := xml.Unmarshal(body, &envelope); err != nil || strings.TrimSpace(envelope.Encrypt) == "" { + log.Printf("[wecom:%s] invalid encrypted message: %v", accountID, err) + w.WriteHeader(http.StatusOK) + return + } + plaintext, err := channel.crypto.decrypt(signature, timestamp, nonce, envelope.Encrypt) + if err != nil { + http.Error(w, "bad signature", http.StatusForbidden) + return + } + var message wecomXMLMessage + if err := xml.Unmarshal(plaintext, &message); err != nil { + log.Printf("[wecom:%s] failed to parse decrypted message: %v", accountID, err) + w.WriteHeader(http.StatusOK) + return + } + if message.MsgType == "text" && message.FromUserName != "" { + channel.handleTextMessage(r.Context(), message.FromUserName, message.FromUserName, message.MsgID, message.Content, "p2p", map[string]any{ + "to_user_name": message.ToUserName, + "from_user_name": message.FromUserName, + "create_time": message.CreateTime, + "msg_type": message.MsgType, + "content": message.Content, + "msg_id": message.MsgID, + "xml": string(plaintext), + }) + } + w.WriteHeader(http.StatusOK) +} + +func weComAccountIDFromPath(path string) (string, bool) { + parts := strings.Split(strings.Trim(path, "/"), "/") + if len(parts) != 3 || parts[0] != "wecom" || parts[2] != "callback" { + return "", false + } + accountID, err := url.PathUnescape(parts[1]) + return accountID, err == nil && accountID != "" +}