diff --git a/pkg/channels/wecom/aibot_ws.go b/pkg/channels/wecom/aibot_ws.go index e6a405b81..cbdb18ac3 100644 --- a/pkg/channels/wecom/aibot_ws.go +++ b/pkg/channels/wecom/aibot_ws.go @@ -705,7 +705,7 @@ func (c *WeComAIBotWSChannel) wsHandleMediaMessage( ctx, cancel := context.WithTimeout(c.ctx, wsImageDownloadTimeout) defer cancel() - ref, err := c.storeWSImage(ctx, chatID, msg.MsgID, resourceURL, aesKey) + ref, err := c.storeWSMedia(ctx, chatID, msg.MsgID, resourceURL, aesKey, wsLabelToDefaultExt(label)) if err != nil { logger.WarnCF("wecom_aibot", "Failed to download/store WS "+label, map[string]any{"error": err.Error(), "url": resourceURL}) @@ -742,8 +742,8 @@ func (c *WeComAIBotWSChannel) handleWSMixedMessage(reqID string, msg WeComAIBotW } case "image": if item.Image != nil { - ref, err := c.storeWSImage(ctx, chatID, - msg.MsgID+"-"+wsGenerateID(), item.Image.URL, item.Image.AESKey) + ref, err := c.storeWSMedia(ctx, chatID, + msg.MsgID+"-"+wsGenerateID(), item.Image.URL, item.Image.AESKey, ".jpg") if err != nil { logger.WarnCF("wecom_aibot", "Failed to download/store mixed image", map[string]any{"error": err.Error()}) @@ -751,6 +751,9 @@ func (c *WeComAIBotWSChannel) handleWSMixedMessage(reqID string, msg WeComAIBotW mediaRefs = append(mediaRefs, ref) } } + default: + logger.WarnCF("wecom_aibot", "Unsupported item type in mixed message", + map[string]any{"msgtype": item.MsgType}) } } @@ -1116,12 +1119,15 @@ func wsGenerateID() string { return generateRandomID(10) } -// ---- Inbound image download helpers ---- +// ---- Inbound media download helpers ---- -// storeWSImage downloads the image at imageURL (with optional AES-CBC decryption) and stores it in the MediaStore. -func (c *WeComAIBotWSChannel) storeWSImage( +// storeWSMedia downloads the resource at resourceURL (with optional AES-CBC +// decryption) and stores it in the MediaStore. The file extension is inferred +// from the HTTP Content-Type response header; defaultExt is used as a fallback +// when the content type is absent or unrecognized. +func (c *WeComAIBotWSChannel) storeWSMedia( ctx context.Context, - chatID, msgID, imageURL, aesKey string, + chatID, msgID, resourceURL, aesKey, defaultExt string, ) (string, error) { store := c.GetMediaStore() if store == nil { @@ -1130,7 +1136,7 @@ func (c *WeComAIBotWSChannel) storeWSImage( const maxSize = 20 << 20 // 20 MB - req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, resourceURL, nil) if err != nil { return "", fmt.Errorf("create request: %w", err) } @@ -1143,13 +1149,19 @@ func (c *WeComAIBotWSChannel) storeWSImage( return "", fmt.Errorf("download HTTP %d", resp.StatusCode) } - // Buffer the image in memory, bounded to maxSize. + // Infer file extension from the Content-Type response header. + ext := wsMediaExtFromContentType(resp.Header.Get("Content-Type")) + if ext == "" { + ext = defaultExt + } + + // Buffer the media in memory, bounded to maxSize. data, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxSize)+1)) if err != nil { - return "", fmt.Errorf("read image: %w", err) + return "", fmt.Errorf("read media: %w", err) } if len(data) > maxSize { - return "", fmt.Errorf("image too large (> %d MB)", maxSize>>20) + return "", fmt.Errorf("media too large (> %d MB)", maxSize>>20) } // AES-CBC decryption if a key is present. @@ -1158,12 +1170,12 @@ func (c *WeComAIBotWSChannel) storeWSImage( if decErr != nil || len(key) != 32 { key, decErr = decodeWeComAESKey(aesKey) if decErr != nil { - return "", fmt.Errorf("decode image AES key: %w", decErr) + return "", fmt.Errorf("decode media AES key: %w", decErr) } } data, err = decryptAESCBC(key, data) if err != nil { - return "", fmt.Errorf("decrypt image: %w", err) + return "", fmt.Errorf("decrypt media: %w", err) } } @@ -1173,7 +1185,7 @@ func (c *WeComAIBotWSChannel) storeWSImage( if err = os.MkdirAll(mediaDir, 0o700); err != nil { return "", fmt.Errorf("mkdir: %w", err) } - tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*.jpg") + tmpFile, err := os.CreateTemp(mediaDir, msgID+"-*"+ext) if err != nil { return "", fmt.Errorf("create temp file: %w", err) } @@ -1182,16 +1194,16 @@ func (c *WeComAIBotWSChannel) storeWSImage( closeErr := tmpFile.Close() if writeErr != nil { os.Remove(tmpPath) - return "", fmt.Errorf("write image: %w", writeErr) + return "", fmt.Errorf("write media: %w", writeErr) } if closeErr != nil { os.Remove(tmpPath) - return "", fmt.Errorf("close image: %w", closeErr) + return "", fmt.Errorf("close media: %w", closeErr) } scope := channels.BuildMediaScope("wecom_aibot", chatID, msgID) ref, err := store.Store(tmpPath, media.MediaMeta{ - Filename: msgID + ".jpg", + Filename: msgID + ext, Source: "wecom_aibot", }, scope) if err != nil { @@ -1200,3 +1212,71 @@ func (c *WeComAIBotWSChannel) storeWSImage( } return ref, nil } + +// wsMediaExtFromContentType returns the lowercase file extension (with leading +// dot) for the given Content-Type value, or "" when the type is unrecognized. +func wsMediaExtFromContentType(contentType string) string { + if contentType == "" { + return "" + } + // Strip parameters (e.g. "image/jpeg; charset=utf-8" → "image/jpeg"). + mt := strings.ToLower(strings.TrimSpace(strings.SplitN(contentType, ";", 2)[0])) + switch mt { + case "image/jpeg", "image/jpg": + return ".jpg" + case "image/png": + return ".png" + case "image/gif": + return ".gif" + case "image/webp": + return ".webp" + case "video/mp4": + return ".mp4" + case "video/mpeg", "video/x-mpeg": + return ".mpeg" + case "video/quicktime": + return ".mov" + case "video/webm": + return ".webm" + case "audio/mpeg", "audio/mp3": + return ".mp3" + case "audio/ogg": + return ".ogg" + case "audio/wav": + return ".wav" + case "application/pdf": + return ".pdf" + case "application/zip": + return ".zip" + case "application/x-rar-compressed", "application/vnd.rar": + return ".rar" + case "text/plain": + return ".txt" + case "application/msword": + return ".doc" + case "application/vnd.openxmlformats-officedocument.wordprocessingml.document": + return ".docx" + case "application/vnd.ms-excel": + return ".xls" + case "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": + return ".xlsx" + case "application/vnd.ms-powerpoint": + return ".ppt" + case "application/vnd.openxmlformats-officedocument.presentationml.presentation": + return ".pptx" + } + return "" +} + +// wsLabelToDefaultExt returns the default file extension for the given media label +// used in wsHandleMediaMessage. It is the fallback when Content-Type detection fails. +func wsLabelToDefaultExt(label string) string { + switch label { + case "image": + return ".jpg" + case "video": + return ".mp4" + default: // "file" and any future labels + return ".bin" + } +} diff --git a/pkg/channels/wecom/aibot_ws_test.go b/pkg/channels/wecom/aibot_ws_test.go index b59d99ef0..4e3314df4 100644 --- a/pkg/channels/wecom/aibot_ws_test.go +++ b/pkg/channels/wecom/aibot_ws_test.go @@ -29,19 +29,19 @@ func newTestWSChannel(t *testing.T) *WeComAIBotWSChannel { return ch } -// TestStoreWSImage_NilStore verifies that storeWSImage returns an error when no +// TestStoreWSMedia_NilStore verifies that storeWSMedia returns an error when no // MediaStore has been injected. -func TestStoreWSImage_NilStore(t *testing.T) { +func TestStoreWSMedia_NilStore(t *testing.T) { ch := newTestWSChannel(t) - _, err := ch.storeWSImage(context.Background(), "chat1", "msg1", "http://any", "") + _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", "http://any", "", ".jpg") 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) { +// TestStoreWSMedia_HTTPError verifies that storeWSMedia propagates HTTP errors +// from the media server. +func TestStoreWSMedia_HTTPError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "not found", http.StatusNotFound) })) @@ -50,29 +50,29 @@ func TestStoreWSImage_HTTPError(t *testing.T) { ch := newTestWSChannel(t) ch.SetMediaStore(media.NewFileMediaStore()) - _, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "") + _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", srv.URL, "", ".jpg") 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) { +// TestStoreWSMedia_ServerUnavailable verifies that storeWSMedia returns a clear +// error when the media server cannot be reached. +func TestStoreWSMedia_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", "") + _, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", "http://127.0.0.1:1", "", ".jpg") if err == nil { t.Fatal("expected error for unreachable server") } } -// TestStoreWSImage_Success_NoAES verifies the happy path: the image is downloaded, +// TestStoreWSMedia_Success_NoAES verifies the happy path: the media 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) { +// ReleaseAll is called. The server returns no Content-Type, so the defaultExt is used. +func TestStoreWSMedia_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) @@ -84,7 +84,7 @@ func TestStoreWSImage_Success_NoAES(t *testing.T) { store := media.NewFileMediaStore() ch.SetMediaStore(store) - ref, err := ch.storeWSImage(context.Background(), "chat1", "msg1", srv.URL, "") + ref, err := ch.storeWSMedia(context.Background(), "chat1", "msg1", srv.URL, "", ".jpg") if err != nil { t.Fatalf("expected no error, got %v", err) } @@ -92,7 +92,7 @@ func TestStoreWSImage_Success_NoAES(t *testing.T) { t.Fatal("expected non-empty ref") } - // File must be accessible after storeWSImage returns (no premature deletion). + // File must be accessible after storeWSMedia returns (no premature deletion). path, err := store.Resolve(ref) if err != nil { t.Fatalf("ref should resolve: %v", err) @@ -115,9 +115,9 @@ func TestStoreWSImage_Success_NoAES(t *testing.T) { } } -// TestStoreWSImage_MultipleMessages verifies that concurrent image messages with +// TestStoreWSMedia_MultipleMessages verifies that concurrent media messages with // different msgIDs do not collide and each resolve to distinct files. -func TestStoreWSImage_MultipleMessages(t *testing.T) { +func TestStoreWSMedia_MultipleMessages(t *testing.T) { imageA := bytes.Repeat([]byte("a"), 64) imageB := bytes.Repeat([]byte("b"), 64) @@ -136,13 +136,13 @@ func TestStoreWSImage_MultipleMessages(t *testing.T) { store := media.NewFileMediaStore() ch.SetMediaStore(store) - refA, err := ch.storeWSImage(context.Background(), "chat1", "msgA", srvA.URL, "") + refA, err := ch.storeWSMedia(context.Background(), "chat1", "msgA", srvA.URL, "", ".jpg") if err != nil { - t.Fatalf("storeWSImage A: %v", err) + t.Fatalf("storeWSMedia A: %v", err) } - refB, err := ch.storeWSImage(context.Background(), "chat1", "msgB", srvB.URL, "") + refB, err := ch.storeWSMedia(context.Background(), "chat1", "msgB", srvB.URL, "", ".jpg") if err != nil { - t.Fatalf("storeWSImage B: %v", err) + t.Fatalf("storeWSMedia B: %v", err) } if refA == refB { t.Fatal("distinct messages must produce distinct refs") @@ -163,3 +163,56 @@ func TestStoreWSImage_MultipleMessages(t *testing.T) { t.Errorf("content mismatch for message B") } } + +// TestStoreWSMedia_ContentTypeExt verifies that the file extension is inferred +// from the HTTP Content-Type header and the defaultExt fallback is used when the +// type is absent or unrecognized. +func TestStoreWSMedia_ContentTypeExt(t *testing.T) { + tests := []struct { + contentType string + wantExt string + }{ + {"image/jpeg", ".jpg"}, + {"image/png", ".png"}, + {"video/mp4", ".mp4"}, + {"application/pdf", ".pdf"}, + {"application/zip", ".zip"}, + // With parameters stripped. + {"video/mp4; codecs=avc1", ".mp4"}, + // Unknown type → falls back to defaultExt. + {"", ""}, + {"application/octet-stream", ""}, + } + for _, tc := range tests { + got := wsMediaExtFromContentType(tc.contentType) + if got != tc.wantExt { + t.Errorf("wsMediaExtFromContentType(%q) = %q, want %q", tc.contentType, got, tc.wantExt) + } + } + + // End-to-end: server returns Content-Type: video/mp4, defaultExt is .bin. + // The stored file should carry the .mp4 extension, not .bin. + payload := bytes.Repeat([]byte("v"), 128) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "video/mp4") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(payload) + })) + defer srv.Close() + + ch := newTestWSChannel(t) + store := media.NewFileMediaStore() + ch.SetMediaStore(store) + + ref, err := ch.storeWSMedia(context.Background(), "chat1", "vid1", srv.URL, "", ".bin") + if err != nil { + t.Fatalf("storeWSMedia: %v", err) + } + path, err := store.Resolve(ref) + if err != nil { + t.Fatalf("resolve: %v", err) + } + if ext := path[len(path)-4:]; ext != ".mp4" { + t.Errorf("expected .mp4 extension from Content-Type, got %q", ext) + } +}