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
This commit is contained in:
Hz_
2026-07-31 10:49:54 +08:00
committed by GitHub
parent 39d8609d6c
commit d69ace8c94
6 changed files with 2229 additions and 0 deletions

View File

@@ -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:

View File

@@ -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 == "<nil>" {
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
}
}

View File

@@ -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)
}
}

View File

@@ -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
}
}

View File

@@ -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(`<xml><ToUserName>corp-id</ToUserName><FromUserName>user-1</FromUserName><CreateTime>123</CreateTime><MsgType>text</MsgType><Content>hello wecom</Content><MsgId>message-1</MsgId></xml>`)
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[:])
}

View File

@@ -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/<account_id>/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 != ""
}