mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-01 13:33:48 +08:00
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:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
254
internal/channels/telegram_test.go
Normal file
254
internal/channels/telegram_test.go
Normal 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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
524
internal/channels/wecom_test.go
Normal file
524
internal/channels/wecom_test.go
Normal 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[:])
|
||||
}
|
||||
315
internal/channels/wecom_webhook.go
Normal file
315
internal/channels/wecom_webhook.go
Normal 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 != ""
|
||||
}
|
||||
Reference in New Issue
Block a user