From 9a19a0c19b9114d184ae1b8901237216b8794ef8 Mon Sep 17 00:00:00 2001 From: Zhang Rui Date: Wed, 11 Mar 2026 17:17:58 +0800 Subject: [PATCH] feat(wecom): implement shared HTTP clients for WeCom image handling and response URL posting --- pkg/channels/wecom/aibot.go | 10 +- pkg/channels/wecom/aibot_test.go | 118 +++++++++++++++++++ pkg/channels/wecom/aibot_ws.go | 171 +++++++++++++++++++--------- pkg/channels/wecom/aibot_ws_test.go | 165 +++++++++++++++++++++++++++ pkg/media/store.go | 60 ++++++++++ 5 files changed, 466 insertions(+), 58 deletions(-) create mode 100644 pkg/channels/wecom/aibot_ws_test.go diff --git a/pkg/channels/wecom/aibot.go b/pkg/channels/wecom/aibot.go index 45385ead1..999f4f13b 100644 --- a/pkg/channels/wecom/aibot.go +++ b/pkg/channels/wecom/aibot.go @@ -22,6 +22,10 @@ import ( "github.com/sipeed/picoclaw/pkg/utils" ) +// responseURLHTTPClient is a shared HTTP client for posting to WeCom response_url. +// Reusing it enables connection pooling across replies. +var responseURLHTTPClient = &http.Client{Timeout: 15 * time.Second} + // WeComAIBotChannel implements the Channel interface for WeCom AI Bot (企业微信智能机器人) type WeComAIBotChannel struct { *channels.BaseChannel @@ -794,8 +798,7 @@ func (c *WeComAIBotChannel) sendViaResponseURL(responseURL, content string) erro } req.Header.Set("Content-Type", "application/json; charset=utf-8") - client := &http.Client{Timeout: 15 * time.Second} - resp, err := client.Do(req) + resp, err := responseURLHTTPClient.Do(req) if err != nil { return fmt.Errorf("post to response_url failed: %w: %w", channels.ErrTemporary, err) } @@ -805,7 +808,8 @@ func (c *WeComAIBotChannel) sendViaResponseURL(responseURL, content string) erro return nil } - respBody, err := io.ReadAll(resp.Body) + const maxErrBody = 64 << 10 // 64 KB is more than enough for any error response + respBody, err := io.ReadAll(io.LimitReader(resp.Body, maxErrBody)) if err != nil { return fmt.Errorf("reading response_url body: %w: %w", channels.ErrTemporary, err) } diff --git a/pkg/channels/wecom/aibot_test.go b/pkg/channels/wecom/aibot_test.go index 6926f7c6a..7c5ae67b1 100644 --- a/pkg/channels/wecom/aibot_test.go +++ b/pkg/channels/wecom/aibot_test.go @@ -3,6 +3,7 @@ package wecom import ( "context" "testing" + "time" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" @@ -335,3 +336,120 @@ func TestWSGenerateID(t *testing.T) { ids[id] = true } } + +// ---- Webhook streaming fallback tests ---- + +// makeWebhookChannel creates a started WeComAIBotChannel for testing. +func makeWebhookChannel(t *testing.T) *WeComAIBotChannel { + t.Helper() + cfg := config.WeComAIBotConfig{ + Enabled: true, + Token: "test_token", + EncodingAESKey: "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFG", + } + ch, err := NewWeComAIBotChannel(cfg, bus.NewMessageBus()) + if err != nil { + t.Fatalf("create channel: %v", err) + } + wc := ch.(*WeComAIBotChannel) + wc.ctx, wc.cancel = context.WithCancel(context.Background()) + return wc +} + +// makeStreamTask creates and registers a streamTask for testing. +func makeStreamTask(t *testing.T, ch *WeComAIBotChannel, streamID, chatID string, deadline time.Time) *streamTask { + t.Helper() + task := &streamTask{ + StreamID: streamID, + ChatID: chatID, + Deadline: deadline, + answerCh: make(chan string, 1), + } + task.ctx, task.cancel = context.WithCancel(ch.ctx) + ch.taskMu.Lock() + ch.streamTasks[streamID] = task + ch.chatTasks[chatID] = append(ch.chatTasks[chatID], task) + ch.taskMu.Unlock() + return task +} + +// TestGetStreamResponse_ImmediateAnswer verifies that when the agent has already +// placed its answer in answerCh, getStreamResponse returns a finish=true response +// and fully removes the task. +func TestGetStreamResponse_ImmediateAnswer(t *testing.T) { + ch := makeWebhookChannel(t) + defer ch.cancel() + + task := makeStreamTask(t, ch, "stream-1", "chat-1", time.Now().Add(30*time.Second)) + task.answerCh <- "hello from agent" + + result := ch.getStreamResponse(task, "ts123", "nonce123") + if result == "" { + t.Fatal("expected non-empty encrypted response") + } + + ch.taskMu.RLock() + _, exists := ch.streamTasks["stream-1"] + ch.taskMu.RUnlock() + if exists { + t.Error("task should have been removed from streamTasks after normal finish") + } + if !task.Finished { + t.Error("task.Finished should be true after normal finish") + } +} + +// TestGetStreamResponse_DeadlinePassed verifies that when the stream deadline has +// elapsed (no agent reply yet), getStreamResponse closes the stream but keeps the +// task alive so the response_url fallback can still deliver the answer. +func TestGetStreamResponse_DeadlinePassed(t *testing.T) { + ch := makeWebhookChannel(t) + defer ch.cancel() + + task := makeStreamTask(t, ch, "stream-2", "chat-2", time.Now().Add(-time.Millisecond)) + + result := ch.getStreamResponse(task, "ts456", "nonce456") + if result == "" { + t.Fatal("expected non-empty encrypted response") + } + + ch.taskMu.RLock() + _, stillStreaming := ch.streamTasks["stream-2"] + ch.taskMu.RUnlock() + if stillStreaming { + t.Error("task should have been removed from streamTasks after deadline") + } + if !task.StreamClosed { + t.Error("task.StreamClosed should be true after deadline") + } + if task.Finished { + t.Error("task.Finished must remain false: agent reply still expected via response_url") + } +} + +// TestGetStreamResponse_StillPending verifies that when neither the agent has +// replied nor the deadline has passed, getStreamResponse returns without altering +// task state (client should poll again). +func TestGetStreamResponse_StillPending(t *testing.T) { + ch := makeWebhookChannel(t) + defer ch.cancel() + + task := makeStreamTask(t, ch, "stream-3", "chat-3", time.Now().Add(30*time.Second)) + + result := ch.getStreamResponse(task, "ts789", "nonce789") + if result == "" { + t.Fatal("expected non-empty encrypted response") + } + + ch.taskMu.RLock() + _, exists := ch.streamTasks["stream-3"] + ch.taskMu.RUnlock() + if !exists { + t.Error("pending task should still be in streamTasks") + } + if task.Finished || task.StreamClosed { + t.Error("pending task should not be finished or stream-closed") + } + // Cleanup. + ch.removeTask(task) +} diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go index 1bea2fdf1..eddea1cb4 100644 --- a/pkg/channels/wecom/aibot_ws.go +++ b/pkg/channels/wecom/aibot_ws.go @@ -1,6 +1,7 @@ package wecom import ( + "bytes" "context" "encoding/base64" "encoding/json" @@ -49,6 +50,10 @@ const ( wsLateReplyRouteTTL = 30 * time.Minute ) +// wsImageHTTPClient is a shared HTTP client for downloading inbound images. +// Reusing it enables connection pooling across multiple image downloads. +var wsImageHTTPClient = &http.Client{Timeout: wsImageDownloadTimeout} + // WeComAIBotWSChannel implements channels.Channel for WeCom AI Bot using the // WebSocket long-connection API. // Unlike the webhook counterpart it does NOT implement WebhookHandler, so the @@ -1040,49 +1045,10 @@ func wsGenerateID() string { // ---- Inbound image download helpers ---- -// downloadWSImage fetches and optionally decrypts a WeCom WS image resource. -// aesKey is the per-resource AES key provided in the callback (may be empty). -func (c *WeComAIBotWSChannel) downloadWSImage(ctx context.Context, imageURL, aesKey string) ([]byte, error) { - const maxSize = 20 << 20 // 20 MB - req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil) - if err != nil { - return nil, fmt.Errorf("create request: %w", err) - } - client := &http.Client{Timeout: wsImageDownloadTimeout} - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("download: %w", err) - } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("download HTTP %d", resp.StatusCode) - } - data, err := io.ReadAll(io.LimitReader(resp.Body, maxSize+1)) - if err != nil { - return nil, fmt.Errorf("read body: %w", err) - } - if len(data) > maxSize { - return nil, fmt.Errorf("image too large (> %d MB)", maxSize>>20) - } - if aesKey == "" { - return data, nil - } - // WeCom per-image AES key: try standard base64 first, then WeCom 43-char format. - key, decErr := base64.StdEncoding.DecodeString(aesKey) - if decErr != nil || len(key) != 32 { - key, decErr = decodeWeComAESKey(aesKey) - if decErr != nil { - return nil, fmt.Errorf("decode image AES key: %w", decErr) - } - } - decrypted, err := decryptAESCBC(key, data) - if err != nil { - return nil, fmt.Errorf("decrypt image: %w", err) - } - return decrypted, nil -} - // storeWSImage downloads, optionally decrypts, and stores an inbound image. +// It streams the HTTP response body directly into the MediaStore, avoiding any +// caller-managed temp files. File lifecycle (creation and deletion) is owned +// entirely by the store and cleaned up via store.ReleaseAll. func (c *WeComAIBotWSChannel) storeWSImage( ctx context.Context, chatID, msgID, imageURL, aesKey string, @@ -1091,26 +1057,121 @@ func (c *WeComAIBotWSChannel) storeWSImage( if store == nil { return "", fmt.Errorf("no media store available") } - data, err := c.downloadWSImage(ctx, imageURL, aesKey) + + const maxSize = 20 << 20 // 20 MB + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil) if err != nil { - return "", err + return "", fmt.Errorf("create request: %w", err) } + resp, err := wsImageHTTPClient.Do(req) + if err != nil { + return "", fmt.Errorf("download: %w", err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("download HTTP %d", resp.StatusCode) + } + mediaDir := filepath.Join(os.TempDir(), "picoclaw_media") + scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) + meta := media.MediaMeta{Filename: msgID + ".jpg", Source: "wecom_aibot"} + + // lr wraps the response body with a +1 limit: if lr.N reaches 0 after the + // transfer, the body was at least maxSize+1 bytes and must be rejected. + lr := &io.LimitedReader{R: resp.Body, N: int64(maxSize) + 1} + var r io.Reader = lr + + if aesKey != "" { + // AES-CBC decryption requires full ciphertext. Buffer the bounded encrypted + // bytes, decrypt, then pass the plaintext as a reader. + encrypted, readErr := io.ReadAll(lr) + if readErr != nil { + return "", fmt.Errorf("read for decrypt: %w", readErr) + } + if lr.N == 0 { + return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) + } + key, decErr := base64.StdEncoding.DecodeString(aesKey) + if decErr != nil || len(key) != 32 { + key, decErr = decodeWeComAESKey(aesKey) + if decErr != nil { + return "", fmt.Errorf("decode image AES key: %w", decErr) + } + } + decrypted, decErr := decryptAESCBC(key, encrypted) + if decErr != nil { + return "", fmt.Errorf("decrypt image: %w", decErr) + } + r = bytes.NewReader(decrypted) + } + + // Fast path: FileMediaStore supports StoreFromReader which manages the file + // lifecycle entirely (temp-rename-fsync pattern, deleted by ReleaseAll). + if fsStore, ok := store.(*media.FileMediaStore); ok { + ref, storeErr := fsStore.StoreFromReader(r, meta, scope, mediaDir) + if storeErr != nil { + return "", storeErr + } + // For the no-AES path, check whether the response body hit the size limit. + if aesKey == "" && lr.N == 0 { + _ = store.ReleaseAll(scope) // remove the oversized file + return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) + } + return ref, nil + } + + // Fallback path for non-FileMediaStore implementations: stream into a temp + // file using the temp-rename-fsync pattern, then register the path. + // The file lifecycle is managed by store.ReleaseAll(). if mkErr := os.MkdirAll(mediaDir, 0o700); mkErr != nil { return "", fmt.Errorf("mkdir: %w", mkErr) } - filename := msgID + ".jpg" - localPath := filepath.Join(mediaDir, utils.SanitizeFilename(filename)) - if writeErr := os.WriteFile(localPath, data, 0o600); writeErr != nil { - return "", fmt.Errorf("write: %w", writeErr) - } - scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) - ref, err := store.Store(localPath, media.MediaMeta{ - Filename: filename, - Source: "wecom_aibot", - }, scope) + tmpPath := filepath.Join(mediaDir, ".tmp-"+wsGenerateID()) + tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) if err != nil { - os.Remove(localPath) + return "", fmt.Errorf("create temp file: %w", err) + } + cleanup := true + defer func() { + if cleanup { + tmpFile.Close() + os.Remove(tmpPath) + } + }() + n, cpErr := io.Copy(tmpFile, r) + syncErr := tmpFile.Sync() + closeErr := tmpFile.Close() + cleanup = false + if cpErr != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("write image: %w", cpErr) + } + if syncErr != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("sync image: %w", syncErr) + } + if closeErr != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("close image: %w", closeErr) + } + if aesKey == "" && lr.N == 0 { + os.Remove(tmpPath) + return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) + } + _ = n + finalPath := filepath.Join(mediaDir, wsGenerateID()+".jpg") + if err = os.Rename(tmpPath, finalPath); err != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("rename: %w", err) + } + if d, openErr := os.Open(mediaDir); openErr == nil { + _ = d.Sync() + d.Close() + } + ref, err := store.Store(finalPath, meta, scope) + if err != nil { + os.Remove(finalPath) return "", fmt.Errorf("store: %w", err) } return ref, nil diff --git a/pkg/channels/wecom/aibot_ws_test.go b/pkg/channels/wecom/aibot_ws_test.go new file mode 100644 index 000000000..b59d99ef0 --- /dev/null +++ b/pkg/channels/wecom/aibot_ws_test.go @@ -0,0 +1,165 @@ +package wecom + +import ( + "bytes" + "context" + "net/http" + "net/http/httptest" + "os" + "testing" + + "github.com/sipeed/picoclaw/pkg/bus" + "github.com/sipeed/picoclaw/pkg/channels" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/media" +) + +// newTestWSChannel creates a WeComAIBotWSChannel ready for unit testing. +func newTestWSChannel(t *testing.T) *WeComAIBotWSChannel { + t.Helper() + cfg := config.WeComAIBotConfig{ + Enabled: true, + BotID: "test_bot_id", + Secret: "test_secret", + } + ch, err := newWeComAIBotWSChannel(cfg, bus.NewMessageBus()) + if err != nil { + t.Fatalf("create WS channel: %v", err) + } + return ch +} + +// TestStoreWSImage_NilStore verifies that storeWSImage returns an error when no +// MediaStore has been injected. +func TestStoreWSImage_NilStore(t *testing.T) { + ch := newTestWSChannel(t) + _, err := ch.storeWSImage(context.Background(), "chat1", "msg1", "http://any", "") + if err == nil { + t.Fatal("expected error when no MediaStore is set") + } +} + +// TestStoreWSImage_HTTPError verifies that storeWSImage propagates HTTP errors +// from the image server. +func TestStoreWSImage_HTTPError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + http.Error(w, "not found", http.StatusNotFound) + })) + defer srv.Close() + + ch := newTestWSChannel(t) + ch.SetMediaStore(media.NewFileMediaStore()) + + _, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "") + if err == nil { + t.Fatal("expected error for HTTP 404") + } +} + +// TestStoreWSImage_ServerUnavailable verifies that storeWSImage returns a clear +// error when the image server cannot be reached. +func TestStoreWSImage_ServerUnavailable(t *testing.T) { + ch := newTestWSChannel(t) + ch.SetMediaStore(media.NewFileMediaStore()) + + // Port 1 is reserved and will refuse the connection immediately. + _, err := ch.storeWSImage(context.Background(), "chat1", "msg1", "http://127.0.0.1:1", "") + if err == nil { + t.Fatal("expected error for unreachable server") + } +} + +// TestStoreWSImage_Success_NoAES verifies the happy path: the image is downloaded, +// a media ref is returned, and the file persists and is readable via Resolve until +// ReleaseAll is called. +func TestStoreWSImage_Success_NoAES(t *testing.T) { + imageData := bytes.Repeat([]byte("x"), 256) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write(imageData) + })) + defer srv.Close() + + ch := newTestWSChannel(t) + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + ref, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if ref == "" { + t.Fatal("expected non-empty ref") + } + + // File must be accessible after storeWSImage returns (no premature deletion). + path, err := store.Resolve(ref) + if err != nil { + t.Fatalf("ref should resolve: %v", err) + } + got, err := os.ReadFile(path) + if err != nil { + t.Fatalf("file should exist at %s: %v", path, err) + } + if !bytes.Equal(got, imageData) { + t.Errorf("content mismatch: got len=%d, want len=%d", len(got), len(imageData)) + } + + // ReleaseAll must delete the file (store owns lifecycle). + scope := channels.BuildMediaScope("wecom_aibot", "chat1", "msg1") + if err := store.ReleaseAll(scope); err != nil { + t.Fatalf("ReleaseAll failed: %v", err) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Errorf("file should have been deleted by ReleaseAll, stat err: %v", err) + } +} + +// TestStoreWSImage_MultipleMessages verifies that concurrent image messages with +// different msgIDs do not collide and each resolve to distinct files. +func TestStoreWSImage_MultipleMessages(t *testing.T) { + imageA := bytes.Repeat([]byte("a"), 64) + imageB := bytes.Repeat([]byte("b"), 64) + + srvA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write(imageA) + })) + defer srvA.Close() + srvB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write(imageB) + })) + defer srvB.Close() + + ch := newTestWSChannel(t) + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + refA, err := ch.storeWSImage(context.Background(), "chat1", "msgA", srvA.URL, "") + if err != nil { + t.Fatalf("storeWSImage A: %v", err) + } + refB, err := ch.storeWSImage(context.Background(), "chat1", "msgB", srvB.URL, "") + if err != nil { + t.Fatalf("storeWSImage B: %v", err) + } + if refA == refB { + t.Fatal("distinct messages must produce distinct refs") + } + + pathA, _ := store.Resolve(refA) + pathB, _ := store.Resolve(refB) + if pathA == pathB { + t.Fatal("distinct messages must be stored at distinct paths") + } + + gotA, _ := os.ReadFile(pathA) + gotB, _ := os.ReadFile(pathB) + if !bytes.Equal(gotA, imageA) { + t.Errorf("content mismatch for message A") + } + if !bytes.Equal(gotB, imageB) { + t.Errorf("content mismatch for message B") + } +} diff --git a/pkg/media/store.go b/pkg/media/store.go index 30220986c..7eb3cbf72 100644 --- a/pkg/media/store.go +++ b/pkg/media/store.go @@ -2,7 +2,9 @@ package media import ( "fmt" + "io" "os" + "path/filepath" "sync" "time" @@ -108,6 +110,64 @@ func (s *FileMediaStore) Store(localPath string, meta MediaMeta, scope string) ( return ref, nil } +// StoreFromReader creates a temp file in dir by streaming r into it, then +// atomically renames it to a UUID-based final filename and registers it under +// scope. The file is owned by the store and deleted by ReleaseAll. +// This uses the same temp-write + sync + rename pattern as fileutil.WriteFileAtomic +// to guarantee flash-storage durability. +func (s *FileMediaStore) StoreFromReader(r io.Reader, meta MediaMeta, scope, dir string) (string, error) { + if err := os.MkdirAll(dir, 0o755); err != nil { + return "", fmt.Errorf("media store: mkdir: %w", err) + } + tmpPath := filepath.Join(dir, ".tmp-picoclaw-"+uuid.New().String()) + tmpFile, err := os.OpenFile(tmpPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600) + if err != nil { + return "", fmt.Errorf("media store: create temp: %w", err) + } + cleanup := true + defer func() { + if cleanup { + tmpFile.Close() + os.Remove(tmpPath) + } + }() + + if _, err = io.Copy(tmpFile, r); err != nil { + return "", fmt.Errorf("media store: write: %w", err) + } + if err = tmpFile.Sync(); err != nil { + return "", fmt.Errorf("media store: sync: %w", err) + } + if err = tmpFile.Close(); err != nil { + return "", fmt.Errorf("media store: close: %w", err) + } + cleanup = false + + // Atomic rename to a UUID-based final name (preserving the original extension). + ext := filepath.Ext(meta.Filename) + finalPath := filepath.Join(dir, uuid.New().String()+ext) + if err = os.Rename(tmpPath, finalPath); err != nil { + os.Remove(tmpPath) + return "", fmt.Errorf("media store: rename: %w", err) + } + // Sync the directory so the rename is durable. + if d, openErr := os.Open(dir); openErr == nil { + _ = d.Sync() + d.Close() + } + + ref := "media://" + uuid.New().String() + s.mu.Lock() + defer s.mu.Unlock() + s.refs[ref] = mediaEntry{path: finalPath, meta: meta, storedAt: s.nowFunc()} + if s.scopeToRefs[scope] == nil { + s.scopeToRefs[scope] = make(map[string]struct{}) + } + s.scopeToRefs[scope][ref] = struct{}{} + s.refToScope[ref] = scope + return ref, nil +} + // Resolve returns the local path for the given ref. func (s *FileMediaStore) Resolve(ref string) (string, error) { s.mu.RLock()