diff --git a/cmd/dobotai/main.go b/cmd/dobotai/main.go index 6b33635..3b3a84d 100644 --- a/cmd/dobotai/main.go +++ b/cmd/dobotai/main.go @@ -68,10 +68,15 @@ func main() { if err != nil { log.Fatal(err) } + downloader, err := dingtalk.NewImageDownloader(dingCfg.ClientID, dingCfg.ClientSecret, nil) + if err != nil { + log.Fatal(err) + } handler, err := dingtalk.NewHandler( dingtalk.TenantConfig{TenantKey: dingCfg.TenantKey, Channel: "dingtalk"}, qaService, sender, + downloader, ) if err != nil { log.Fatal(err) diff --git a/db/migrations/000002_turn_media.down.sql b/db/migrations/000002_turn_media.down.sql new file mode 100644 index 0000000..22adb4c --- /dev/null +++ b/db/migrations/000002_turn_media.down.sql @@ -0,0 +1,2 @@ +ALTER TABLE turns + DROP COLUMN IF EXISTS media_json; diff --git a/db/migrations/000002_turn_media.up.sql b/db/migrations/000002_turn_media.up.sql new file mode 100644 index 0000000..6fd892d --- /dev/null +++ b/db/migrations/000002_turn_media.up.sql @@ -0,0 +1,2 @@ +ALTER TABLE turns + ADD COLUMN IF NOT EXISTS media_json JSONB NOT NULL DEFAULT '[]'::JSONB; diff --git a/db/migrations/000002_turn_media_contract_test.go b/db/migrations/000002_turn_media_contract_test.go new file mode 100644 index 0000000..50eda2e --- /dev/null +++ b/db/migrations/000002_turn_media_contract_test.go @@ -0,0 +1,80 @@ +// Package migrations_test contains source-level contract tests for the +// 000002_turn_media migration. No real database connection is used. +package migrations_test + +import ( + "os" + "path/filepath" + "regexp" + "strings" + "testing" +) + +func readMediaUpMigration(t *testing.T) string { + t.Helper() + data, err := os.ReadFile(filepath.Join(".", "000002_turn_media.up.sql")) + if err != nil { + t.Fatalf("read up migration: %v", err) + } + return string(data) +} + +func readMediaDownMigration(t *testing.T) string { + t.Helper() + data, err := os.ReadFile(filepath.Join(".", "000002_turn_media.down.sql")) + if err != nil { + t.Fatalf("read down migration: %v", err) + } + return string(data) +} + +// TestUpMediaMigrationAddsMediaJSONColumn enforces that the up migration +// adds a media_json JSONB column to the turns table with appropriate defaults. +func TestUpMediaMigrationAddsMediaJSONColumn(t *testing.T) { + up := readMediaUpMigration(t) + upper := strings.ToUpper(up) + + if !strings.Contains(upper, "ADD COLUMN IF NOT EXISTS MEDIA_JSON") { + t.Error("up migration must use ADD COLUMN IF NOT EXISTS media_json") + } + if !strings.Contains(upper, "JSONB NOT NULL DEFAULT '[]'::JSONB") { + t.Error("up migration must declare media_json JSONB NOT NULL DEFAULT '[]'::JSONB") + } +} + +// TestUpMediaMigrationIsIdempotent enforces that the up migration uses IF NOT +// EXISTS so it can be safely re-run. +func TestUpMediaMigrationIsIdempotent(t *testing.T) { + up := readMediaUpMigration(t) + upper := strings.ToUpper(up) + + if !strings.Contains(upper, "IF NOT EXISTS") { + t.Error("up migration must use IF NOT EXISTS for idempotency") + } +} + +// TestDownMediaMigrationDropsColumn enforces the down migration drops the +// media_json column with IF EXISTS. +func TestDownMediaMigrationDropsColumn(t *testing.T) { + down := readMediaDownMigration(t) + upper := strings.ToUpper(down) + + if !strings.Contains(upper, "DROP COLUMN IF EXISTS MEDIA_JSON") { + t.Error("down migration must use DROP COLUMN IF EXISTS media_json") + } +} + +// TestUpMigrationHasNoReservedWordMedia ensures the column is named media_json +// not a bare "media" which could be a reserved word in future PostgreSQL. +func TestUpMigrationHasNoReservedWordMedia(t *testing.T) { + up := readMediaUpMigration(t) + + // The bare word "media" as a column name (not media_json) should be avoided + bare := regexp.MustCompile(`(?i)\bMEDIA\s+JSONB\b`) + if bare.MatchString(up) { + t.Error("up migration must not use bare 'media' column name; use media_json") + } + if !regexp.MustCompile(`(?i)\bMEDIA_JSON\s+JSONB\b`).MatchString(up) { + t.Error("up migration must use media_json column name") + } +} diff --git a/internal/channels/dingtalk/handler.go b/internal/channels/dingtalk/handler.go index f681f74..c92d193 100644 --- a/internal/channels/dingtalk/handler.go +++ b/internal/channels/dingtalk/handler.go @@ -4,37 +4,30 @@ import ( "context" "errors" "fmt" + "log" "strings" "gitea.zheromiku.top/Weknora/dobotai-platform/internal/core" ) const ( - // msgTypeText is the DingTalk robot callback message type for text. - msgTypeText = "text" - // conversationTypeDirect is a 1:1 (direct) conversation callback. + msgTypeText = "text" + msgTypePicture = "picture" conversationTypeDirect = "1" - // conversationTypeGroup is a group conversation callback. - conversationTypeGroup = "2" + conversationTypeGroup = "2" + replyTitle = "DobotAI" - // replyTitle is the markdown title used for session webhook replies. - replyTitle = "DobotAI" - - // unsupportedMessageReply is sent for non-text messages. unsupportedMessageReply = "当前版本暂仅支持文字消息。" - // qaErrorReply is sent when the QA backend fails. It must never leak - // internal error details to end users. - qaErrorReply = "抱歉,当前暂时无法完成回答,请稍后重试。" + qaErrorReply = "抱歉,当前暂时无法完成回答,请稍后重试。" + imageErrorReply = "图片识别暂时不可用,请稍后重试或发送文字。" + imageQuery = "请描述并分析这张图片。" ) -// TenantConfig identifies the tenant + channel this handler serves. type TenantConfig struct { TenantKey string Channel string } -// BotCallback is the decoded DingTalk robot callback payload for the -// fields the text ingress path needs. type BotCallback struct { ConversationType string ConversationID string @@ -46,28 +39,31 @@ type BotCallback struct { IsInAtList bool SessionWebhook string Text string + DownloadCode string } -// QAAsker answers questions (implemented by the QA service). type QAAsker interface { Ask(ctx context.Context, req core.QARequest) (core.QAResult, error) } -// MarkdownSender sends a markdown reply to a session webhook -// (implemented by WebhookSender). type MarkdownSender interface { SendMarkdown(ctx context.Context, webhookURL, title, text string) error } -// Handler processes DingTalk robot callbacks synchronously. -type Handler struct { - cfg TenantConfig - qa QAAsker - reply MarkdownSender +type ImageDownloader interface { + Download(ctx context.Context, downloadCode string) (core.ImageInput, error) } -// NewHandler validates its dependencies and returns a Handler. -func NewHandler(cfg TenantConfig, qa QAAsker, reply MarkdownSender) (*Handler, error) { +type Handler struct { + cfg TenantConfig + qa QAAsker + reply MarkdownSender + downloader ImageDownloader +} + +// NewHandler accepts at most one image downloader. A nil/omitted downloader +// keeps text ingress available while picture messages fail closed. +func NewHandler(cfg TenantConfig, qa QAAsker, reply MarkdownSender, downloaders ...ImageDownloader) (*Handler, error) { if strings.TrimSpace(cfg.TenantKey) == "" { return nil, errors.New("dingtalk handler: tenant key is required") } @@ -80,78 +76,113 @@ func NewHandler(cfg TenantConfig, qa QAAsker, reply MarkdownSender) (*Handler, e if reply == nil { return nil, errors.New("dingtalk handler: markdown sender is required") } - return &Handler{cfg: cfg, qa: qa, reply: reply}, nil + if len(downloaders) > 1 { + return nil, errors.New("dingtalk handler: at most one image downloader is allowed") + } + var downloader ImageDownloader + if len(downloaders) == 1 { + downloader = downloaders[0] + } + return &Handler{cfg: cfg, qa: qa, reply: reply, downloader: downloader}, nil } -// Handle processes a single bot callback synchronously and returns nil -// once any required reply has been sent (or the message was silently -// ignored). No goroutines are spawned. func (h *Handler) Handle(ctx context.Context, cb BotCallback) error { - // Group messages that do not @ the bot are silently ignored. + // Admission comes first: a group message without @ must have zero side effects. if cb.ConversationType == conversationTypeGroup && !cb.IsInAtList { return nil } - // Non-text messages never reach QA; reply with the unsupported notice. - if cb.Msgtype != msgTypeText { - if strings.TrimSpace(cb.SessionWebhook) == "" { - return errors.New("dingtalk handler: sessionWebhook is required to reply") - } - return h.reply.SendMarkdown(ctx, cb.SessionWebhook, replyTitle, unsupportedMessageReply) + // Validate the only reply destination before downloading media or spending a + // model call. Webhook errors never include its query string. + webhookURL := strings.TrimSpace(cb.SessionWebhook) + if webhookURL == "" { + return errors.New("dingtalk handler: sessionWebhook is required to reply") } - - query := strings.TrimSpace(cb.Text) - if query == "" { - return errors.New("dingtalk handler: text message content is empty") + if err := ValidateWebhookURL(webhookURL); err != nil { + return fmt.Errorf("dingtalk handler: invalid sessionWebhook: %w", err) } - if cb.MsgID == "" { + if strings.TrimSpace(cb.MsgID) == "" { return errors.New("dingtalk handler: msgId is required") } - externalUserID := cb.SenderStaffID + externalUserID := strings.TrimSpace(cb.SenderStaffID) if externalUserID == "" { - externalUserID = cb.SenderID + externalUserID = strings.TrimSpace(cb.SenderID) } if externalUserID == "" { return errors.New("dingtalk handler: sender staffId/senderId is required") } - // The session webhook is the only reply channel; fail fast before - // calling QA so a missing webhook never wastes a model call. - if strings.TrimSpace(cb.SessionWebhook) == "" { - return errors.New("dingtalk handler: sessionWebhook is required to reply") - } - req := core.QARequest{ TenantKey: h.cfg.TenantKey, Channel: h.cfg.Channel, ExternalUserID: externalUserID, - ExternalMessageID: cb.MsgID, + ExternalMessageID: strings.TrimSpace(cb.MsgID), SenderName: cb.SenderNick, - Query: query, } switch cb.ConversationType { case conversationTypeGroup: - if cb.ConversationID == "" { + if strings.TrimSpace(cb.ConversationID) == "" { return errors.New("dingtalk handler: conversationId is required for group messages") } req.ConversationType = core.ConversationGroup - req.ExternalChatID = cb.ConversationID + req.ExternalChatID = strings.TrimSpace(cb.ConversationID) default: req.ConversationType = core.ConversationDirect } + switch strings.ToLower(strings.TrimSpace(cb.Msgtype)) { + case msgTypeText: + req.Query = strings.TrimSpace(cb.Text) + if req.Query == "" { + return errors.New("dingtalk handler: text message content is empty") + } + case msgTypePicture: + if strings.TrimSpace(cb.DownloadCode) == "" { + log.Printf("DingTalk image preparation failed: stage=missing_download_code") + return h.sendReply(ctx, webhookURL, imageErrorReply) + } + if h.downloader == nil { + log.Printf("DingTalk image preparation failed: stage=downloader_unavailable") + return h.sendReply(ctx, webhookURL, imageErrorReply) + } + image, err := h.downloader.Download(ctx, strings.TrimSpace(cb.DownloadCode)) + if err != nil { + log.Printf("DingTalk image preparation failed: stage=%s", imageFailureStage(err)) + return h.sendReply(ctx, webhookURL, imageErrorReply) + } + req.Query = imageQuery + req.Images = []core.ImageInput{image} + default: + return h.sendReply(ctx, webhookURL, unsupportedMessageReply) + } + result, err := h.qa.Ask(ctx, req) if err != nil { - // Never leak internal error details to the user. If the friendly - // reply cannot be delivered, surface that failure instead. - if replyErr := h.sendReply(ctx, cb.SessionWebhook, qaErrorReply); replyErr != nil { + if replyErr := h.sendReply(ctx, webhookURL, qaErrorReply); replyErr != nil { return replyErr } return nil } + return h.sendReply(ctx, webhookURL, result.Answer) +} - return h.sendReply(ctx, cb.SessionWebhook, result.Answer) +func imageFailureStage(err error) string { + message := strings.ToLower(err.Error()) + switch { + case strings.Contains(message, "oauth"): + return "oauth" + case strings.Contains(message, "api"): + return "file_api" + case strings.Contains(message, "unsafe") || strings.Contains(message, "destination"): + return "unsafe_url" + case strings.Contains(message, "format") || strings.Contains(message, "image data") || strings.Contains(message, "dimensions") || strings.Contains(message, "size limit"): + return "image_validation" + case strings.Contains(message, "request") || strings.Contains(message, "http") || strings.Contains(message, "read"): + return "image_fetch" + default: + return "unknown" + } } func (h *Handler) sendReply(ctx context.Context, webhookURL, text string) error { diff --git a/internal/channels/dingtalk/handler_test.go b/internal/channels/dingtalk/handler_test.go index a30234c..05564e6 100644 --- a/internal/channels/dingtalk/handler_test.go +++ b/internal/channels/dingtalk/handler_test.go @@ -38,13 +38,30 @@ func (f *fakeReply) SendMarkdown(_ context.Context, url, title, text string) err return f.err } +type fakeImageDownloader struct { + called bool + code string + image core.ImageInput + err error +} + +func (f *fakeImageDownloader) Download(_ context.Context, code string) (core.ImageInput, error) { + f.called = true + f.code = code + return f.image, f.err +} + func newTestHandler(t *testing.T, qa *fakeQA, reply *fakeReply) *Handler { + return newTestHandlerWithDownloader(t, qa, reply, nil) +} + +func newTestHandlerWithDownloader(t *testing.T, qa *fakeQA, reply *fakeReply, downloader ImageDownloader) *Handler { t.Helper() cfg := TenantConfig{ TenantKey: "tenant-a", Channel: "dingtalk", } - h, err := NewHandler(cfg, qa, reply) + h, err := NewHandler(cfg, qa, reply, downloader) if err != nil { t.Fatalf("NewHandler: %v", err) } @@ -79,6 +96,14 @@ func groupTextCallback() BotCallback { } } +func directPictureCallback() BotCallback { + data := directTextCallback() + data.Msgtype = "picture" + data.Text = "" + data.DownloadCode = "temporary-download-code" + return data +} + // --- A: direct text --- func TestHandleDirectTextCallsQAAndRepliesAnswer(t *testing.T) { @@ -212,7 +237,7 @@ func TestHandleNonTextRepliesUnsupported(t *testing.T) { h := newTestHandler(t, qa, reply) data := directTextCallback() - data.Msgtype = "picture" + data.Msgtype = "audio" err := h.Handle(context.Background(), data) if err != nil { t.Fatalf("Handle: %v", err) @@ -228,6 +253,92 @@ func TestHandleNonTextRepliesUnsupported(t *testing.T) { } } +func TestHandlePictureDownloadsAndCallsQA(t *testing.T) { + qa := &fakeQA{result: core.QAResult{Answer: "图片回答"}} + reply := &fakeReply{} + image := core.ImageInput{DataURI: "data:image/png;base64,AAAA", MIMEType: "image/png", SizeBytes: 4, Width: 1, Height: 1} + downloader := &fakeImageDownloader{image: image} + h := newTestHandlerWithDownloader(t, qa, reply, downloader) + + if err := h.Handle(context.Background(), directPictureCallback()); err != nil { + t.Fatalf("Handle picture: %v", err) + } + if !downloader.called || downloader.code != "temporary-download-code" { + t.Fatalf("download called=%v code=%q", downloader.called, downloader.code) + } + if !qa.called || len(qa.lastReq.Images) != 1 || qa.lastReq.Images[0] != image { + t.Fatalf("QA request did not receive image: %+v", qa.lastReq) + } + if qa.lastReq.Query != "请描述并分析这张图片。" { + t.Fatalf("picture query=%q", qa.lastReq.Query) + } + if !reply.called || reply.lastText != "图片回答" { + t.Fatalf("reply=%+v", reply) + } +} + +func TestHandlePicturePreparationFailuresAreFriendlyAndDoNotCallQA(t *testing.T) { + tests := []struct { + name string + code string + downloader ImageDownloader + }{ + {name: "missing download code", code: "", downloader: &fakeImageDownloader{}}, + {name: "missing downloader", code: "code", downloader: nil}, + {name: "download error", code: "secret-code", downloader: &fakeImageDownloader{err: errors.New("signed URL and token must not leak")}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + qa := &fakeQA{} + reply := &fakeReply{} + h := newTestHandlerWithDownloader(t, qa, reply, tt.downloader) + data := directPictureCallback() + data.DownloadCode = tt.code + if err := h.Handle(context.Background(), data); err != nil { + t.Fatalf("Handle: %v", err) + } + if qa.called { + t.Fatal("QA must not be called") + } + if !reply.called || reply.lastText != "图片识别暂时不可用,请稍后重试或发送文字。" { + t.Fatalf("friendly reply=%q", reply.lastText) + } + }) + } +} + +func TestHandleGroupPictureWithoutMentionDoesNothing(t *testing.T) { + qa := &fakeQA{} + reply := &fakeReply{} + downloader := &fakeImageDownloader{} + h := newTestHandlerWithDownloader(t, qa, reply, downloader) + data := directPictureCallback() + data.ConversationType = conversationTypeGroup + data.IsInAtList = false + data.SessionWebhook = "not even validated" + if err := h.Handle(context.Background(), data); err != nil { + t.Fatalf("Handle: %v", err) + } + if downloader.called || qa.called || reply.called { + t.Fatalf("unexpected side effects: downloader=%v qa=%v reply=%v", downloader.called, qa.called, reply.called) + } +} + +func TestHandlePictureRejectsWebhookBeforeDownload(t *testing.T) { + qa := &fakeQA{} + reply := &fakeReply{} + downloader := &fakeImageDownloader{} + h := newTestHandlerWithDownloader(t, qa, reply, downloader) + data := directPictureCallback() + data.SessionWebhook = "http://127.0.0.1/steal" + if err := h.Handle(context.Background(), data); err == nil { + t.Fatal("expected invalid webhook error") + } + if downloader.called || qa.called || reply.called { + t.Fatalf("invalid webhook caused side effects: downloader=%v qa=%v reply=%v", downloader.called, qa.called, reply.called) + } +} + // --- D: QA error --- func TestHandleQAErrorRepliesFriendlyMessage(t *testing.T) { diff --git a/internal/channels/dingtalk/image.go b/internal/channels/dingtalk/image.go new file mode 100644 index 0000000..64deac3 --- /dev/null +++ b/internal/channels/dingtalk/image.go @@ -0,0 +1,342 @@ +package dingtalk + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "image" + _ "image/gif" + _ "image/jpeg" + _ "image/png" + "io" + "net" + "net/http" + "net/netip" + "net/url" + "strings" + "sync" + "time" + + "gitea.zheromiku.top/Weknora/dobotai-platform/internal/core" +) + +const ( + oauthTokenURL = "https://api.dingtalk.com/v1.0/oauth2/accessToken" + messageFileDownloadURL = "https://api.dingtalk.com/v1.0/robot/messageFiles/download" + maxAPIResponseBytes = 1 << 20 + maxImageBytes = 8 << 20 + maxImagePixels = 40_000_000 +) + +var blockedDownloadRanges = []netip.Prefix{ + netip.MustParsePrefix("100.64.0.0/10"), + netip.MustParsePrefix("192.0.0.0/24"), + netip.MustParsePrefix("192.0.2.0/24"), + netip.MustParsePrefix("198.18.0.0/15"), + netip.MustParsePrefix("198.51.100.0/24"), + netip.MustParsePrefix("203.0.113.0/24"), + netip.MustParsePrefix("2001:db8::/32"), +} + +type DingTalkImageDownloader struct { + clientID string + clientSecret string + tokenURL string + fileURL string + apiClient *http.Client + downloadHTTP *http.Client + allowTestURL bool + + mu sync.Mutex + token string + tokenTill time.Time +} + +func NewImageDownloader(clientID, clientSecret string, client *http.Client) (*DingTalkImageDownloader, error) { + return newImageDownloader(clientID, clientSecret, client, oauthTokenURL, messageFileDownloadURL, false) +} + +func newImageDownloaderForTest(clientID, clientSecret string, client *http.Client, tokenURL, fileURL string) (*DingTalkImageDownloader, error) { + return newImageDownloader(clientID, clientSecret, client, tokenURL, fileURL, true) +} + +func newImageDownloader(clientID, clientSecret string, client *http.Client, tokenURL, fileURL string, allowTest bool) (*DingTalkImageDownloader, error) { + clientID = strings.TrimSpace(clientID) + clientSecret = strings.TrimSpace(clientSecret) + if clientID == "" { + return nil, errors.New("dingtalk image downloader: clientID is required") + } + if clientSecret == "" { + return nil, errors.New("dingtalk image downloader: clientSecret is required") + } + base := &http.Client{Timeout: 30 * time.Second} + if client != nil { + copyClient := *client + base = ©Client + if base.Timeout <= 0 { + base.Timeout = 30 * time.Second + } + } + apiClient := *base + apiClient.CheckRedirect = func(*http.Request, []*http.Request) error { + return errors.New("dingtalk API redirect refused") + } + downloadClient := *base + if !allowTest { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.Proxy = nil + transport.DialContext = safeDialContext + apiClient.Transport = transport + downloadClient.Transport = transport.Clone() + } + d := &DingTalkImageDownloader{ + clientID: clientID, clientSecret: clientSecret, + tokenURL: tokenURL, fileURL: fileURL, + apiClient: &apiClient, downloadHTTP: &downloadClient, allowTestURL: allowTest, + } + d.downloadHTTP.CheckRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) >= 3 { + return errors.New("dingtalk image redirect limit exceeded") + } + req.Header = make(http.Header) + if !d.allowTestURL { + if err := validateDownloadURL(req.URL); err != nil { + return err + } + } + return nil + } + return d, nil +} + +func (d *DingTalkImageDownloader) Download(ctx context.Context, downloadCode string) (core.ImageInput, error) { + downloadCode = strings.TrimSpace(downloadCode) + if downloadCode == "" { + return core.ImageInput{}, errors.New("dingtalk image download: downloadCode is required") + } + token, err := d.accessToken(ctx) + if err != nil { + return core.ImageInput{}, err + } + payload, err := d.postJSON(ctx, d.fileURL, map[string]string{ + "downloadCode": downloadCode, + "robotCode": d.clientID, + }, token) + if err != nil { + return core.ImageInput{}, err + } + var fileResponse struct { + DownloadURL string `json:"downloadUrl"` + } + if json.Unmarshal(payload, &fileResponse) != nil || strings.TrimSpace(fileResponse.DownloadURL) == "" { + return core.ImageInput{}, errors.New("dingtalk image download: invalid message file response") + } + downloadURL := strings.TrimSpace(fileResponse.DownloadURL) + if !d.allowTestURL { + downloadURL, err = normalizeInitialDownloadURL(downloadURL) + if err != nil { + return core.ImageInput{}, err + } + parsed, parseErr := url.Parse(downloadURL) + if parseErr != nil || validateDownloadURL(parsed) != nil { + return core.ImageInput{}, errors.New("dingtalk image download: unsafe download URL") + } + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil) + if err != nil { + return core.ImageInput{}, errors.New("dingtalk image download: invalid download request") + } + resp, err := d.downloadHTTP.Do(req) + if err != nil { + if ctx.Err() != nil { + return core.ImageInput{}, ctx.Err() + } + return core.ImageInput{}, errors.New("dingtalk image download: request failed") + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return core.ImageInput{}, fmt.Errorf("dingtalk image download: HTTP %d", resp.StatusCode) + } + if resp.ContentLength > maxImageBytes { + return core.ImageInput{}, errors.New("dingtalk image download: image exceeds size limit") + } + data, err := io.ReadAll(io.LimitReader(resp.Body, maxImageBytes+1)) + if err != nil { + return core.ImageInput{}, errors.New("dingtalk image download: read failed") + } + if len(data) > maxImageBytes { + return core.ImageInput{}, errors.New("dingtalk image download: image exceeds size limit") + } + return validateImageBytes(data) +} + +func (d *DingTalkImageDownloader) accessToken(ctx context.Context) (string, error) { + d.mu.Lock() + defer d.mu.Unlock() + if d.token != "" && time.Now().Before(d.tokenTill) { + return d.token, nil + } + payload, err := d.postJSON(ctx, d.tokenURL, map[string]string{ + "appKey": d.clientID, "appSecret": d.clientSecret, + }, "") + if err != nil { + return "", err + } + var response struct { + AccessToken string `json:"accessToken"` + ExpireIn int64 `json:"expireIn"` + } + if json.Unmarshal(payload, &response) != nil || strings.TrimSpace(response.AccessToken) == "" { + return "", errors.New("dingtalk image download: invalid OAuth response") + } + ttl := time.Duration(response.ExpireIn) * time.Second + refreshAfter := ttl - 5*time.Minute + if refreshAfter <= 0 { + refreshAfter = ttl / 2 + } + if refreshAfter <= 0 { + refreshAfter = time.Minute + } + d.token = response.AccessToken + d.tokenTill = time.Now().Add(refreshAfter) + return d.token, nil +} + +func (d *DingTalkImageDownloader) postJSON(ctx context.Context, endpoint string, value any, token string) ([]byte, error) { + body, err := json.Marshal(value) + if err != nil { + return nil, errors.New("dingtalk image download: encode API request failed") + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body)) + if err != nil { + return nil, errors.New("dingtalk image download: create API request failed") + } + req.Header.Set("Content-Type", "application/json") + if token != "" { + req.Header.Set("x-acs-dingtalk-access-token", token) + } + resp, err := d.apiClient.Do(req) + if err != nil { + if ctx.Err() != nil { + return nil, ctx.Err() + } + return nil, errors.New("dingtalk image download: API request failed") + } + defer resp.Body.Close() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("dingtalk image download: API HTTP %d", resp.StatusCode) + } + return readBoundedAPIResponse(resp) +} + +func readBoundedAPIResponse(resp *http.Response) ([]byte, error) { + if resp.ContentLength > maxAPIResponseBytes { + return nil, errors.New("dingtalk image download: API response exceeds size limit") + } + payload, err := io.ReadAll(io.LimitReader(resp.Body, maxAPIResponseBytes+1)) + if err != nil { + return nil, errors.New("dingtalk image download: API response read failed") + } + if len(payload) > maxAPIResponseBytes { + return nil, errors.New("dingtalk image download: API response exceeds size limit") + } + return payload, nil +} + +func validateImageBytes(data []byte) (core.ImageInput, error) { + if len(data) == 0 || len(data) > maxImageBytes { + return core.ImageInput{}, errors.New("dingtalk image download: invalid image size") + } + mime := http.DetectContentType(data) + expectedFormat := map[string]string{ + "image/png": "png", "image/jpeg": "jpeg", "image/gif": "gif", + }[mime] + if expectedFormat == "" { + return core.ImageInput{}, errors.New("dingtalk image download: unsupported image format") + } + config, format, err := image.DecodeConfig(bytes.NewReader(data)) + if err != nil || format != expectedFormat || config.Width <= 0 || config.Height <= 0 { + return core.ImageInput{}, errors.New("dingtalk image download: invalid image data") + } + if int64(config.Width)*int64(config.Height) > maxImagePixels { + return core.ImageInput{}, errors.New("dingtalk image download: image dimensions exceed limit") + } + return core.ImageInput{ + DataURI: "data:" + mime + ";base64," + base64.StdEncoding.EncodeToString(data), + MIMEType: mime, SizeBytes: int64(len(data)), Width: config.Width, Height: config.Height, + }, nil +} + +func normalizeInitialDownloadURL(raw string) (string, error) { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || parsed.Hostname() == "" || parsed.User != nil { + return "", errors.New("dingtalk image download: unsafe download URL") + } + switch parsed.Scheme { + case "http": + parsed.Scheme = "https" + case "https": + default: + return "", errors.New("dingtalk image download: unsafe download URL") + } + return parsed.String(), nil +} + +func validateDownloadURL(parsed *url.URL) error { + if parsed == nil || parsed.Scheme != "https" || parsed.Hostname() == "" || parsed.User != nil { + return errors.New("dingtalk image download: unsafe download URL") + } + return nil +} + +func safeDialContext(ctx context.Context, network, address string) (net.Conn, error) { + host, port, err := net.SplitHostPort(address) + if err != nil { + return nil, errors.New("dingtalk image download: invalid destination") + } + var ips []net.IP + if literal := net.ParseIP(host); literal != nil { + ips = []net.IP{literal} + } else { + resolved, resolveErr := net.DefaultResolver.LookupIPAddr(ctx, host) + if resolveErr != nil || len(resolved) == 0 { + return nil, errors.New("dingtalk image download: destination resolution failed") + } + for _, item := range resolved { + ips = append(ips, item.IP) + } + } + for _, ip := range ips { + if !isSafeDownloadIP(ip) { + return nil, errors.New("dingtalk image download: destination is not public") + } + } + dialer := &net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second} + for _, ip := range ips { + conn, dialErr := dialer.DialContext(ctx, network, net.JoinHostPort(ip.String(), port)) + if dialErr == nil { + return conn, nil + } + } + return nil, errors.New("dingtalk image download: destination connection failed") +} + +func isSafeDownloadIP(ip net.IP) bool { + if ip == nil || !ip.IsGlobalUnicast() || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || ip.IsLinkLocalMulticast() || ip.IsMulticast() || ip.IsUnspecified() { + return false + } + addr, ok := netip.AddrFromSlice(ip) + if !ok { + return false + } + addr = addr.Unmap() + for _, prefix := range blockedDownloadRanges { + if prefix.Contains(addr) { + return false + } + } + return true +} diff --git a/internal/channels/dingtalk/image_test.go b/internal/channels/dingtalk/image_test.go new file mode 100644 index 0000000..4dd91a1 --- /dev/null +++ b/internal/channels/dingtalk/image_test.go @@ -0,0 +1,226 @@ +package dingtalk + +import ( + "bytes" + "context" + "encoding/json" + "image" + "image/color" + "image/png" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func tinyPNG(t *testing.T) []byte { + t.Helper() + img := image.NewRGBA(image.Rect(0, 0, 2, 3)) + for y := 0; y < 3; y++ { + for x := 0; x < 2; x++ { + img.Set(x, y, color.RGBA{R: 255, A: 255}) + } + } + var out bytes.Buffer + if err := png.Encode(&out, img); err != nil { + t.Fatal(err) + } + return out.Bytes() +} + +func TestNewImageDownloaderMissingCredentials(t *testing.T) { + for _, tc := range []struct{ id, secret string }{{"", "s"}, {"id", ""}, {" ", "s"}, {"id", " "}} { + if got, err := NewImageDownloader(tc.id, tc.secret, nil); err == nil || got != nil { + t.Fatalf("NewImageDownloader(%q,%q)=%v,%v", tc.id, tc.secret, got, err) + } + } +} + +func TestDownloadEmptyCodeReturnsError(t *testing.T) { + d, err := NewImageDownloader("id", "secret", nil) + if err != nil { + t.Fatal(err) + } + for _, code := range []string{"", " "} { + if _, err := d.Download(context.Background(), code); err == nil { + t.Fatalf("Download(%q) should fail", code) + } + } +} + +func TestImageDownloaderTwoHopAndTokenCache(t *testing.T) { + pngBytes := tinyPNG(t) + var tokenCalls, fileCalls int + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/token": + tokenCalls++ + var body map[string]string + _ = json.NewDecoder(r.Body).Decode(&body) + if body["appKey"] != "client-id" || body["appSecret"] != "client-secret" { + t.Errorf("oauth body=%v", body) + } + if body["clientId"] != "" || body["clientSecret"] != "" { + t.Errorf("legacy OAuth field names must not be sent: %v", body) + } + _ = json.NewEncoder(w).Encode(map[string]any{"accessToken": "access-token", "expireIn": 3600}) + case "/file": + fileCalls++ + if r.Header.Get("x-acs-dingtalk-access-token") != "access-token" { + t.Errorf("token header missing") + } + var body map[string]string + _ = json.NewDecoder(r.Body).Decode(&body) + if body["downloadCode"] != "download-code" || body["robotCode"] != "client-id" { + t.Errorf("file body=%v", body) + } + _ = json.NewEncoder(w).Encode(map[string]string{"downloadUrl": server.URL + "/image?signature=secret"}) + case "/image": + w.Header().Set("Content-Type", "text/plain") // deliberately untrusted + _, _ = w.Write(pngBytes) + default: + http.NotFound(w, r) + } + })) + defer server.Close() + + d, err := newImageDownloaderForTest("client-id", "client-secret", server.Client(), server.URL+"/token", server.URL+"/file") + if err != nil { + t.Fatal(err) + } + for i := 0; i < 2; i++ { + got, err := d.Download(context.Background(), "download-code") + if err != nil { + t.Fatalf("Download: %v", err) + } + if got.MIMEType != "image/png" || got.Width != 2 || got.Height != 3 || got.SizeBytes != int64(len(pngBytes)) || !strings.HasPrefix(got.DataURI, "data:image/png;base64,") { + t.Fatalf("image=%+v", got) + } + } + if tokenCalls != 1 || fileCalls != 2 { + t.Fatalf("tokenCalls=%d fileCalls=%d", tokenCalls, fileCalls) + } +} + +func TestImageDownloaderRejectsRedirectWithoutLeakingToken(t *testing.T) { + var leaked string + target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + leaked = r.Header.Get("x-acs-dingtalk-access-token") + w.WriteHeader(http.StatusNoContent) + })) + defer target.Close() + var source *httptest.Server + source = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/token" { + _ = json.NewEncoder(w).Encode(map[string]any{"accessToken": "top-secret-token", "expireIn": 3600}) + return + } + http.Redirect(w, r, target.URL, http.StatusFound) + })) + defer source.Close() + d, err := newImageDownloaderForTest("id", "secret", source.Client(), source.URL+"/token", source.URL+"/file") + if err != nil { + t.Fatal(err) + } + if _, err := d.Download(context.Background(), "code"); err == nil { + t.Fatal("API redirect should fail") + } + if leaked != "" { + t.Fatalf("token leaked across redirect: %q", leaked) + } +} + +func TestImageDownloaderRejectsInvalidImageAndOversize(t *testing.T) { + tests := []struct { + name string + contentType string + body []byte + length int64 + }{ + {name: "html disguised", contentType: "image/png", body: []byte("not image")}, + {name: "oversized content length", contentType: "image/png", body: []byte("x"), length: maxImageBytes + 1}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var server *httptest.Server + server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/token": + _ = json.NewEncoder(w).Encode(map[string]any{"accessToken": "token", "expireIn": 3600}) + case "/file": + _ = json.NewEncoder(w).Encode(map[string]string{"downloadUrl": server.URL + "/image?q=signed-secret"}) + case "/image": + w.Header().Set("Content-Type", tt.contentType) + if tt.length > 0 { + w.Header().Set("Content-Length", "8388609") + } + _, _ = w.Write(tt.body) + } + })) + defer server.Close() + d, _ := newImageDownloaderForTest("id", "secret", server.Client(), server.URL+"/token", server.URL+"/file") + _, err := d.Download(context.Background(), "download-secret") + if err == nil { + t.Fatal("expected rejection") + } + for _, secret := range []string{"download-secret", "signed-secret", "token"} { + if strings.Contains(err.Error(), secret) { + t.Fatalf("error leaked %q: %v", secret, err) + } + } + }) + } +} + +func TestNormalizeInitialDownloadURL(t *testing.T) { + got, err := normalizeInitialDownloadURL("http://example.com/media/path?signature=secret") + if err != nil { + t.Fatalf("normalize: %v", err) + } + if got != "https://example.com/media/path?signature=secret" { + t.Fatalf("normalized=%q", got) + } + for _, raw := range []string{"ftp://example.com/a", "http://user:pass@example.com/a", "not-a-url"} { + if _, err := normalizeInitialDownloadURL(raw); err == nil { + t.Fatalf("%q should fail", raw) + } + } +} + +func TestSafeDownloadIPRejectsNonPublicRanges(t *testing.T) { + for _, raw := range []string{"127.0.0.1", "10.0.0.1", "192.168.1.1", "169.254.1.1", "100.64.0.1", "192.0.2.1", "198.18.0.1", "203.0.113.1", "::1", "fc00::1", "2001:db8::1"} { + if isSafeDownloadIP(net.ParseIP(raw)) { + t.Errorf("%s must be rejected", raw) + } + } + for _, raw := range []string{"8.8.8.8", "1.1.1.1", "2606:4700:4700::1111"} { + if !isSafeDownloadIP(net.ParseIP(raw)) { + t.Errorf("%s should be public", raw) + } + } +} + +func TestReadImageRejectsPixelLimit(t *testing.T) { + // A PNG only needs a valid signature+IHDR for DecodeConfig. Patch a valid + // tiny PNG's dimensions and let CRC become invalid: DecodeConfig must reject + // it or the pixel policy must reject it, either outcome is fail closed. + data := tinyPNG(t) + if len(data) < 24 { + t.Fatal("short png fixture") + } + copy(data[16:20], []byte{0, 0, 0x27, 0x10}) // 10000 + copy(data[20:24], []byte{0, 0, 0x13, 0x88}) // 5000 = 50M + if _, err := validateImageBytes(data); err == nil { + t.Fatal("50M-pixel image must fail") + } +} + +func TestBoundedAPIResponse(t *testing.T) { + resp := &http.Response{Body: io.NopCloser(io.LimitReader(strings.NewReader(strings.Repeat("x", maxAPIResponseBytes+1)), maxAPIResponseBytes+1))} + if _, err := readBoundedAPIResponse(resp); err == nil { + t.Fatal("oversized API response must fail") + } +} diff --git a/internal/channels/dingtalk/stream.go b/internal/channels/dingtalk/stream.go index ddb47e1..501589a 100644 --- a/internal/channels/dingtalk/stream.go +++ b/internal/channels/dingtalk/stream.go @@ -2,6 +2,7 @@ package dingtalk import ( "context" + "encoding/json" "errors" "fmt" "strings" @@ -54,16 +55,19 @@ func (c StreamConfig) Validate() error { } // BotCallbackFromSDK converts the official SDK callback model into the -// channel-internal BotCallback. It is a pure one-to-one field mapping; no -// trimming, validation or business logic happens here (Handler owns that). +// channel-internal BotCallback. Msgtype is normalized (trimmed, lowercased). +// For picture messages the downloadCode is extracted from the SDK Content +// field (json marshal/unmarshal round-trip); malformed or missing codes +// produce an empty DownloadCode but never downgrade the Msgtype. func BotCallbackFromSDK(data *chatbot.BotCallbackDataModel) (BotCallback, error) { if data == nil { return BotCallback{}, errors.New("dingtalk stream: callback data is nil") } - return BotCallback{ + msgtype := strings.ToLower(strings.TrimSpace(data.Msgtype)) + cb := BotCallback{ ConversationType: data.ConversationType, ConversationID: data.ConversationId, - Msgtype: data.Msgtype, + Msgtype: msgtype, MsgID: data.MsgId, SenderNick: data.SenderNick, SenderStaffID: data.SenderStaffId, @@ -71,7 +75,41 @@ func BotCallbackFromSDK(data *chatbot.BotCallbackDataModel) (BotCallback, error) IsInAtList: data.IsInAtList, SessionWebhook: data.SessionWebhook, Text: data.Text.Content, - }, nil + } + if msgtype == "picture" { + cb.DownloadCode = extractDownloadCode(data.Content) + } + return cb, nil +} + +// extractDownloadCode attempts to pull content.downloadCode from the SDK +// Content field (interface{}) via a json marshal/unmarshal round-trip. It +// never panics on malformed input and returns an empty string for any +// extraction failure. +func extractDownloadCode(raw interface{}) string { + if raw == nil { + return "" + } + b, err := json.Marshal(raw) + if err != nil { + return "" + } + var direct struct { + DownloadCode string `json:"downloadCode"` + } + if json.Unmarshal(b, &direct) == nil && strings.TrimSpace(direct.DownloadCode) != "" { + return strings.TrimSpace(direct.DownloadCode) + } + // Keep compatibility with transports that preserve the callback envelope. + var nested struct { + Content struct { + DownloadCode string `json:"downloadCode"` + } `json:"content"` + } + if err := json.Unmarshal(b, &nested); err != nil { + return "" + } + return strings.TrimSpace(nested.Content.DownloadCode) } // StreamClient wraps the official DingTalk Stream SDK client. It wires diff --git a/internal/channels/dingtalk/stream_test.go b/internal/channels/dingtalk/stream_test.go index f8cba79..c7e3632 100644 --- a/internal/channels/dingtalk/stream_test.go +++ b/internal/channels/dingtalk/stream_test.go @@ -142,6 +142,117 @@ func TestBotCallbackFromSDK(t *testing.T) { }) } +func TestBotCallbackFromSDKPicture(t *testing.T) { + t.Parallel() + + t.Run("extracts downloadCode from picture content", func(t *testing.T) { + t.Parallel() + data := &chatbot.BotCallbackDataModel{ + ConversationType: "1", + Msgtype: "picture", + MsgId: "msg-pic-1", + Content: map[string]interface{}{ + "downloadCode": " abc123 ", + }, + } + cb, err := BotCallbackFromSDK(data) + if err != nil { + t.Fatalf("BotCallbackFromSDK() error = %v", err) + } + if cb.Msgtype != "picture" { + t.Fatalf("Msgtype = %q, want picture", cb.Msgtype) + } + if cb.DownloadCode != "abc123" { + t.Fatalf("DownloadCode = %q, want trimmed abc123", cb.DownloadCode) + } + if cb.Text != "" { + t.Fatalf("Text = %q, want empty for picture", cb.Text) + } + }) + + t.Run("malformed content keeps picture with empty downloadCode", func(t *testing.T) { + t.Parallel() + data := &chatbot.BotCallbackDataModel{ + ConversationType: "1", + Msgtype: "picture", + MsgId: "msg-pic-2", + Content: "not-a-map", + } + cb, err := BotCallbackFromSDK(data) + if err != nil { + t.Fatalf("BotCallbackFromSDK() error = %v", err) + } + if cb.Msgtype != "picture" { + t.Fatalf("Msgtype = %q, want picture (never downgrade)", cb.Msgtype) + } + if cb.DownloadCode != "" { + t.Fatalf("DownloadCode = %q, want empty", cb.DownloadCode) + } + }) + + t.Run("missing downloadCode keeps picture", func(t *testing.T) { + t.Parallel() + data := &chatbot.BotCallbackDataModel{ + ConversationType: "1", + Msgtype: "picture", + MsgId: "msg-pic-3", + Content: map[string]interface{}{ + "content": map[string]interface{}{}, + }, + } + cb, err := BotCallbackFromSDK(data) + if err != nil { + t.Fatalf("BotCallbackFromSDK() error = %v", err) + } + if cb.Msgtype != "picture" { + t.Fatalf("Msgtype = %q, want picture", cb.Msgtype) + } + if cb.DownloadCode != "" { + t.Fatalf("DownloadCode = %q, want empty", cb.DownloadCode) + } + }) + + t.Run("text regression: Msgtype normalized", func(t *testing.T) { + t.Parallel() + data := &chatbot.BotCallbackDataModel{ + ConversationType: "1", + Msgtype: " TEXT ", + MsgId: "msg-t-1", + Text: chatbot.BotCallbackDataTextModel{Content: "hello"}, + } + cb, err := BotCallbackFromSDK(data) + if err != nil { + t.Fatalf("BotCallbackFromSDK() error = %v", err) + } + if cb.Msgtype != "text" { + t.Fatalf("Msgtype = %q, want normalized text", cb.Msgtype) + } + if cb.Text != "hello" { + t.Fatalf("Text = %q, want hello", cb.Text) + } + }) + + t.Run("nil content keeps picture", func(t *testing.T) { + t.Parallel() + data := &chatbot.BotCallbackDataModel{ + ConversationType: "1", + Msgtype: "picture", + MsgId: "msg-pic-4", + Content: nil, + } + cb, err := BotCallbackFromSDK(data) + if err != nil { + t.Fatalf("BotCallbackFromSDK() error = %v", err) + } + if cb.Msgtype != "picture" { + t.Fatalf("Msgtype = %q, want picture", cb.Msgtype) + } + if cb.DownloadCode != "" { + t.Fatalf("DownloadCode = %q, want empty", cb.DownloadCode) + } + }) +} + func mustTestHandler(t *testing.T) *Handler { t.Helper() return newTestHandler(t, &fakeQA{}, &fakeReply{}) diff --git a/internal/core/types.go b/internal/core/types.go index 4653878..105f1a3 100644 --- a/internal/core/types.go +++ b/internal/core/types.go @@ -15,6 +15,27 @@ const ( var ErrTurnInProgress = errors.New("turn is already processing") +// ImageInput contains the image data for a QA request. The DataURI is a +// full data: URI (e.g. "data:image/png;base64,...") and must never be +// persisted to the audit log. +type ImageInput struct { + DataURI string + MIMEType string + SizeBytes int64 + Width int + Height int +} + +// MediaAudit records metadata about media in a turn for audit purposes. +// It must never contain DataURI, download codes, or URLs. +type MediaAudit struct { + Type string `json:"type"` // "image" + MIMEType string `json:"mime_type"` + SizeBytes int64 `json:"size_bytes"` + Width int `json:"width"` + Height int `json:"height"` +} + type QARequest struct { TenantKey string Channel string @@ -24,6 +45,7 @@ type QARequest struct { ExternalMessageID string SenderName string Query string + Images []ImageInput } type Reference struct { @@ -65,6 +87,7 @@ type Turn struct { Question string Answer string References []Reference + Media []MediaAudit Status string ErrorMessage string } diff --git a/internal/qa/service.go b/internal/qa/service.go index d482cf3..fd61fd8 100644 --- a/internal/qa/service.go +++ b/internal/qa/service.go @@ -56,9 +56,16 @@ func (s *Service) Ask(ctx context.Context, request core.QARequest) (core.QAResul return core.QAResult{}, err } + media := make([]core.MediaAudit, 0, len(request.Images)) + for _, image := range request.Images { + media = append(media, core.MediaAudit{ + Type: "image", MIMEType: image.MIMEType, SizeBytes: image.SizeBytes, + Width: image.Width, Height: image.Height, + }) + } turn := core.Turn{ID: newID(), TenantKey: request.TenantKey, Channel: request.Channel, ExternalMessageID: request.ExternalMessageID, SessionKey: sessionKey, - Question: request.Query, Status: "processing"} + Question: request.Query, Media: media, Status: "processing"} stored, created, err := s.repository.BeginTurn(ctx, turn) if err != nil { return core.QAResult{}, fmt.Errorf("begin turn: %w", err) @@ -91,8 +98,13 @@ func (s *Service) Ask(ctx context.Context, request core.QARequest) (core.QAResul if request.ConversationType == core.ConversationGroup && strings.TrimSpace(request.SenderName) != "" { query = request.SenderName + ": " + request.Query } + images := make([]weknora.Image, 0, len(request.Images)) + for _, image := range request.Images { + images = append(images, weknora.Image{Data: image.DataURI}) + } chat, err := s.weknora.AgentChat(ctx, sessionID, weknora.AgentChatRequest{ Query: query, AgentID: agentID, AgentEnabled: true, Channel: "im", EnableMemory: true, + Images: images, }) if err != nil { turn.Status, turn.ErrorMessage = "failed", err.Error() diff --git a/internal/qa/service_test.go b/internal/qa/service_test.go index 63ef6a7..b4c162d 100644 --- a/internal/qa/service_test.go +++ b/internal/qa/service_test.go @@ -63,6 +63,7 @@ type fakeWeKnora struct { mu sync.Mutex sessions int queries []string + requests []weknora.AgentChatRequest } func (f *fakeWeKnora) CreateSession(context.Context, string, string) (string, error) { @@ -75,6 +76,7 @@ func (f *fakeWeKnora) AgentChat(_ context.Context, sessionID string, request wek f.mu.Lock() defer f.mu.Unlock() f.queries = append(f.queries, request.Query) + f.requests = append(f.requests, request) return weknora.ChatResult{RequestID: "request-1", Answer: "answer", References: []core.Reference{{KnowledgeID: "kb-1"}}}, nil } @@ -102,4 +104,137 @@ func TestAskUsesSharedGroupSessionAndReplaysDuplicate(t *testing.T) { if client.sessions != 1 || len(client.queries) != 1 || client.queries[0] != "Alice: question" { t.Fatalf("sessions=%d queries=%v", client.sessions, client.queries) } + // sender prefix must not regress + if len(client.requests) != 1 { + t.Fatalf("expected 1 weknora request, got %d", len(client.requests)) + } + if client.requests[0].Query != "Alice: question" { + t.Fatalf("sender prefix regression: got %q", client.requests[0].Query) + } + // Turn.Media must be empty for no-image request + repo.mu.Lock() + stored := repo.turns["tenant:dingtalk:message-1"] + repo.mu.Unlock() + if len(stored.Media) != 0 { + t.Fatalf("expected empty Media, got %+v", stored.Media) + } +} + +func TestAskPropagatesImagesToWeKnora(t *testing.T) { + repo := newMemoryRepository() + client := &fakeWeKnora{} + service, err := New(repo, fixedResolver{}, client, session.GroupShared) + if err != nil { + t.Fatal(err) + } + req := core.QARequest{ + TenantKey: "tenant", Channel: "dingtalk", ConversationType: core.ConversationGroup, + ExternalChatID: "group-1", ExternalUserID: "user-1", + ExternalMessageID: "message-img", SenderName: "Alice", Query: "describe this", + Images: []core.ImageInput{ + {DataURI: "data:image/png;base64,AAAA", MIMEType: "image/png", SizeBytes: 1024, Width: 640, Height: 480}, + }, + } + result, err := service.Ask(context.Background(), req) + if err != nil { + t.Fatalf("Ask: %v", err) + } + if result.Answer != "answer" { + t.Fatalf("unexpected answer: %s", result.Answer) + } + client.mu.Lock() + defer client.mu.Unlock() + if len(client.requests) != 1 { + t.Fatalf("expected 1 weknora request, got %d", len(client.requests)) + } + chatReq := client.requests[0] + if len(chatReq.Images) != 1 { + t.Fatalf("expected 1 image in weknora request, got %d", len(chatReq.Images)) + } + if chatReq.Images[0].Data != "data:image/png;base64,AAAA" { + t.Fatalf("image Data not propagated: %q", chatReq.Images[0].Data) + } + // sender prefix must still work + if chatReq.Query != "Alice: describe this" { + t.Fatalf("sender prefix regression: %q", chatReq.Query) + } +} + +func TestAskDerivesTurnMediaFromImages(t *testing.T) { + repo := newMemoryRepository() + client := &fakeWeKnora{} + service, err := New(repo, fixedResolver{}, client, session.GroupShared) + if err != nil { + t.Fatal(err) + } + req := core.QARequest{ + TenantKey: "tenant", Channel: "dingtalk", ConversationType: core.ConversationGroup, + ExternalChatID: "group-1", ExternalUserID: "user-1", + ExternalMessageID: "message-media", SenderName: "Alice", Query: "q", + Images: []core.ImageInput{ + {DataURI: "data:image/png;base64,AAAA", MIMEType: "image/png", SizeBytes: 2048, Width: 640, Height: 480}, + }, + } + _, err = service.Ask(context.Background(), req) + if err != nil { + t.Fatalf("Ask: %v", err) + } + repo.mu.Lock() + stored := repo.turns["tenant:dingtalk:message-media"] + repo.mu.Unlock() + if len(stored.Media) != 1 { + t.Fatalf("expected 1 MediaAudit, got %d", len(stored.Media)) + } + m := stored.Media[0] + if m.Type != "image" || m.MIMEType != "image/png" || m.SizeBytes != 2048 || m.Width != 640 || m.Height != 480 { + t.Fatalf("MediaAudit metadata mismatch: %+v", m) + } + // DataURI must NOT be in Turn.Media + if m.Type == "data:image/png;base64,AAAA" { + t.Fatal("MediaAudit must not contain DataURI") + } +} + +func TestAskDuplicateReplayDoesNotCallWeKnoraAgain(t *testing.T) { + repo := newMemoryRepository() + client := &fakeWeKnora{} + service, err := New(repo, fixedResolver{}, client, session.GroupShared) + if err != nil { + t.Fatal(err) + } + req := core.QARequest{ + TenantKey: "tenant", Channel: "dingtalk", ConversationType: core.ConversationGroup, + ExternalChatID: "group-1", ExternalUserID: "user-1", + ExternalMessageID: "message-replay", SenderName: "Alice", Query: "q", + Images: []core.ImageInput{ + {DataURI: "data:image/png;base64,AAAA", MIMEType: "image/png", SizeBytes: 1024, Width: 640, Height: 480}, + }, + } + first, err := service.Ask(context.Background(), req) + if err != nil { + t.Fatalf("first Ask: %v", err) + } + second, err := service.Ask(context.Background(), req) + if err != nil { + t.Fatalf("replay Ask: %v", err) + } + if !second.Replayed { + t.Fatal("second should be replayed") + } + if first.Answer != second.Answer { + t.Fatal("replay must return same answer") + } + client.mu.Lock() + defer client.mu.Unlock() + // WeKnora must be called exactly once + if len(client.requests) != 1 { + t.Fatalf("duplicate must not call WeKnora again; got %d requests", len(client.requests)) + } + // Turn.Media must be preserved on replay + repo.mu.Lock() + stored := repo.turns["tenant:dingtalk:message-replay"] + repo.mu.Unlock() + if len(stored.Media) != 1 { + t.Fatalf("replay must preserve Turn.Media: %+v", stored.Media) + } } diff --git a/internal/store/postgres/store.go b/internal/store/postgres/store.go index c368179..dbc13b5 100644 --- a/internal/store/postgres/store.go +++ b/internal/store/postgres/store.go @@ -35,10 +35,18 @@ func (s *Store) Close() { } func (s *Store) BeginTurn(ctx context.Context, turn core.Turn) (core.Turn, bool, error) { + media := turn.Media + if media == nil { + media = []core.MediaAudit{} + } + mediaJSON, err := json.Marshal(media) + if err != nil { + return core.Turn{}, false, err + } command, err := s.pool.Exec(ctx, `INSERT INTO turns - (id, tenant_key, channel, external_message_id, session_key, question, status) - VALUES ($1,$2,$3,$4,$5,$6,'processing') ON CONFLICT DO NOTHING`, - turn.ID, turn.TenantKey, turn.Channel, turn.ExternalMessageID, turn.SessionKey, turn.Question) + (id, tenant_key, channel, external_message_id, session_key, question, status, media_json) + VALUES ($1,$2,$3,$4,$5,$6,'processing',$7) ON CONFLICT DO NOTHING`, + turn.ID, turn.TenantKey, turn.Channel, turn.ExternalMessageID, turn.SessionKey, turn.Question, mediaJSON) if err != nil { return core.Turn{}, false, err } @@ -84,14 +92,23 @@ func (s *Store) FinishTurn(ctx context.Context, turn core.Turn) error { if err != nil { return err } - _, err = s.pool.Exec(ctx, `UPDATE turns SET weknora_session_id=$2,weknora_request_id=$3,answer=$4,references_json=$5,status=$6,error_message=$7,updated_at=NOW() WHERE id=$1`, turn.ID, turn.WeKnoraSessionID, turn.WeKnoraRequestID, turn.Answer, references, turn.Status, turn.ErrorMessage) + media := turn.Media + if media == nil { + media = []core.MediaAudit{} + } + mediaJSON, err := json.Marshal(media) + if err != nil { + return err + } + _, err = s.pool.Exec(ctx, `UPDATE turns SET weknora_session_id=$2,weknora_request_id=$3,answer=$4,references_json=$5,media_json=$6,status=$7,error_message=$8,updated_at=NOW() WHERE id=$1`, turn.ID, turn.WeKnoraSessionID, turn.WeKnoraRequestID, turn.Answer, references, mediaJSON, turn.Status, turn.ErrorMessage) return err } func (s *Store) loadTurn(ctx context.Context, tenantKey, channel, messageID string) (core.Turn, error) { var turn core.Turn var references []byte - err := s.pool.QueryRow(ctx, `SELECT id,tenant_key,channel,external_message_id,session_key,weknora_session_id,weknora_request_id,question,answer,references_json,status,error_message FROM turns WHERE tenant_key=$1 AND channel=$2 AND external_message_id=$3`, tenantKey, channel, messageID).Scan(&turn.ID, &turn.TenantKey, &turn.Channel, &turn.ExternalMessageID, &turn.SessionKey, &turn.WeKnoraSessionID, &turn.WeKnoraRequestID, &turn.Question, &turn.Answer, &references, &turn.Status, &turn.ErrorMessage) + var mediaJSON []byte + err := s.pool.QueryRow(ctx, `SELECT id,tenant_key,channel,external_message_id,session_key,weknora_session_id,weknora_request_id,question,answer,references_json,media_json,status,error_message FROM turns WHERE tenant_key=$1 AND channel=$2 AND external_message_id=$3`, tenantKey, channel, messageID).Scan(&turn.ID, &turn.TenantKey, &turn.Channel, &turn.ExternalMessageID, &turn.SessionKey, &turn.WeKnoraSessionID, &turn.WeKnoraRequestID, &turn.Question, &turn.Answer, &references, &mediaJSON, &turn.Status, &turn.ErrorMessage) if err != nil { return core.Turn{}, err } @@ -100,5 +117,10 @@ func (s *Store) loadTurn(ctx context.Context, tenantKey, channel, messageID stri return core.Turn{}, err } } + if len(mediaJSON) > 0 { + if err := json.Unmarshal(mediaJSON, &turn.Media); err != nil { + return core.Turn{}, err + } + } return turn, nil } diff --git a/internal/store/postgres/store_contract_test.go b/internal/store/postgres/store_contract_test.go index 6a5ae33..8ca3244 100644 --- a/internal/store/postgres/store_contract_test.go +++ b/internal/store/postgres/store_contract_test.go @@ -60,7 +60,7 @@ func TestStoreSQLUsesReferencesJSONColumn(t *testing.T) { if !strings.Contains(update, "references_json=$5") { t.Error("FinishTurn UPDATE must write to references_json column") } - if !strings.Contains(selectTurns, "answer,references_json,status") { + if !strings.Contains(selectTurns, "answer,references_json,media_json,status") && !strings.Contains(selectTurns, "answer,references_json,status") { t.Error("loadTurn SELECT must read references_json column") } } diff --git a/internal/store/postgres/store_media_test.go b/internal/store/postgres/store_media_test.go new file mode 100644 index 0000000..73cd0a6 --- /dev/null +++ b/internal/store/postgres/store_media_test.go @@ -0,0 +1,99 @@ +package postgres + +import ( + "encoding/json" + "strings" + "testing" + + "gitea.zheromiku.top/Weknora/dobotai-platform/internal/core" +) + +func TestStoreSQLUsesMediaJSONColumn(t *testing.T) { + src := readStoreSource(t) + sqls := extractSQLLiterals(src) + + var insert, update, selectTurns string + for _, sql := range sqls { + upper := strings.ToUpper(sql) + if strings.Contains(upper, "INSERT") && strings.Contains(upper, "INTO TURNS") { + insert = sql + } + if strings.HasPrefix(strings.TrimSpace(upper), "UPDATE TURNS") { + update = sql + } + if strings.Contains(upper, "SELECT") && strings.Contains(upper, "FROM TURNS") { + selectTurns = sql + } + } + + if insert == "" { + t.Fatal("BeginTurn INSERT INTO turns statement not found") + } + if update == "" { + t.Fatal("FinishTurn UPDATE turns statement not found") + } + if selectTurns == "" { + t.Fatal("loadTurn SELECT FROM turns statement not found") + } + + if !strings.Contains(insert, "media_json") { + t.Error("BeginTurn INSERT must include media_json column") + } + if !strings.Contains(update, "media_json") { + t.Error("FinishTurn UPDATE must write to media_json column") + } + if !strings.Contains(selectTurns, "media_json") { + t.Error("loadTurn SELECT must read media_json column") + } +} + +func TestMediaAuditJSONExcludesSecrets(t *testing.T) { + ma := core.MediaAudit{ + Type: "image", + MIMEType: "image/png", + SizeBytes: 2048, + Width: 640, + Height: 480, + } + raw, err := json.Marshal(ma) + if err != nil { + t.Fatalf("marshal: %v", err) + } + s := string(raw) + for _, secret := range []string{"DataURI", "data_uri", "DownloadCode", "download_code", "URL", "url"} { + if strings.Contains(s, secret) { + t.Errorf("MediaAudit JSON must not contain secret field %q: %s", secret, s) + } + } + var fields map[string]any + if err := json.Unmarshal(raw, &fields); err != nil { + t.Fatalf("unmarshal fields: %v", err) + } + for _, key := range []string{"type", "mime_type", "size_bytes", "width", "height"} { + if _, ok := fields[key]; !ok { + t.Errorf("MediaAudit JSON missing %q field: %s", key, s) + } + } +} + +func TestMediaAuditNilMarshalsToNull(t *testing.T) { + // Nil slice marshals to "null" — the store must guard against this + // because the media_json column is NOT NULL with DEFAULT '[]'. + var media []core.MediaAudit + raw, err := json.Marshal(media) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if string(raw) != "null" { + t.Errorf("nil media slice marshals to %q (expected null); store must handle nil→[]", raw) + } +} + +func TestStoreHasNilMediaGuard(t *testing.T) { + src := readStoreSource(t) + // The store must ensure nil media is persisted as [] not null. + // Check for a nil guard in BeginTurn (e.g. "if turn.Media == nil" or "media == nil"). + if !strings.Contains(src, "edia == nil") && !strings.Contains(src, "edia != nil") && !strings.Contains(src, "len(turn.Media)") { + t.Error("store must guard against nil media slice (NOT NULL column)") + } +} diff --git a/internal/weknora/client_test.go b/internal/weknora/client_test.go index 1d41e73..07c32ee 100644 --- a/internal/weknora/client_test.go +++ b/internal/weknora/client_test.go @@ -2,6 +2,7 @@ package weknora import ( "context" + "encoding/json" "errors" "net/http" "net/http/httptest" @@ -63,3 +64,37 @@ func TestParseSSEErrorDominatesPartialAnswer(t *testing.T) { t.Fatalf("expected terminal vision error, got %v", err) } } + +func TestAgentChatRequestJSONSerializesImages(t *testing.T) { + req := AgentChatRequest{ + Query: "请描述并分析这张图片。", + AgentEnabled: true, + AgentID: "agent-1", + Channel: "im", + Images: []Image{ + {Data: "data:image/png;base64,iVBORw0KGgo="}, + }, + } + raw, err := json.Marshal(req) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if !strings.Contains(string(raw), `"images":[{"data":"data:image/png;base64,iVBORw0KGgo="}]`) { + t.Fatalf("images not serialized correctly: %s", raw) + } +} + +func TestAgentChatRequestJSONOmitsEmptyImages(t *testing.T) { + req := AgentChatRequest{ + Query: "hello", + AgentEnabled: true, + AgentID: "agent-1", + } + raw, err := json.Marshal(req) + if err != nil { + t.Fatalf("marshal: %v", err) + } + if strings.Contains(string(raw), `"images"`) { + t.Fatalf("empty images should be omitted: %s", raw) + } +}