feat: publish outbound answer images to private OSS
This commit is contained in:
@@ -35,6 +35,18 @@ DINGTALK_FAQ_TABLE_SHEET_ID=
|
||||
DINGTALK_FAQ_TABLE_VIEW_ID=
|
||||
DINGTALK_FAQ_TABLE_OPERATOR_ID=
|
||||
|
||||
# Optional PDF-answer image egress for DingTalk cards. When any value is set,
|
||||
# all five values are required. DobotAI downloads only authenticated resource://
|
||||
# images from WEKNORA_BASE_URL /files, writes content-hashed copies under
|
||||
# dobotai/im-images/, and emits private OSS GET URLs valid for 2 hours.
|
||||
# Keep the bucket private and configure an OSS lifecycle rule to delete
|
||||
# dobotai/im-images/ objects after 1 day. Never commit real credentials.
|
||||
DOBOTAI_OUTBOUND_IMAGE_OSS_ENDPOINT=
|
||||
DOBOTAI_OUTBOUND_IMAGE_OSS_REGION=
|
||||
DOBOTAI_OUTBOUND_IMAGE_OSS_ACCESS_KEY_ID=
|
||||
DOBOTAI_OUTBOUND_IMAGE_OSS_ACCESS_KEY_SECRET=
|
||||
DOBOTAI_OUTBOUND_IMAGE_OSS_BUCKET=
|
||||
|
||||
# Optional support group for uncovered-answer notifications. This is the
|
||||
# group's openConversationId, not a webhook. Leave empty to record coverage
|
||||
# decisions and pending support cases without sending group notifications.
|
||||
|
||||
@@ -2,6 +2,7 @@ package main
|
||||
|
||||
import (
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
@@ -15,10 +16,12 @@ func TestMainWiresDirectoryResolverIntoDingTalkHandler(t *testing.T) {
|
||||
for _, required := range []string{
|
||||
"directory.NewDingTalkLookup",
|
||||
"directory.NewResolver",
|
||||
"Directory: directoryResolver",
|
||||
} {
|
||||
if !strings.Contains(src, required) {
|
||||
t.Errorf("main.go missing %q", required)
|
||||
}
|
||||
}
|
||||
if !regexp.MustCompile(`Directory:\s*directoryResolver`).MatchString(src) {
|
||||
t.Error("main.go does not wire directoryResolver into HandlerOptions.Directory")
|
||||
}
|
||||
}
|
||||
|
||||
+18
-1
@@ -181,6 +181,18 @@ func main() {
|
||||
if err != nil {
|
||||
log.Fatalf("DingTalk directory resolver init failed: %v", err)
|
||||
}
|
||||
outboundRewriter, err := buildOutboundImageRewriter(outboundImageConfig{
|
||||
WeKnoraBaseURL: os.Getenv("WEKNORA_BASE_URL"),
|
||||
WeKnoraAPIKey: os.Getenv("WEKNORA_API_KEY"),
|
||||
Endpoint: os.Getenv("DOBOTAI_OUTBOUND_IMAGE_OSS_ENDPOINT"),
|
||||
Region: os.Getenv("DOBOTAI_OUTBOUND_IMAGE_OSS_REGION"),
|
||||
AccessKeyID: os.Getenv("DOBOTAI_OUTBOUND_IMAGE_OSS_ACCESS_KEY_ID"),
|
||||
AccessKeySecret: os.Getenv("DOBOTAI_OUTBOUND_IMAGE_OSS_ACCESS_KEY_SECRET"),
|
||||
Bucket: os.Getenv("DOBOTAI_OUTBOUND_IMAGE_OSS_BUCKET"),
|
||||
})
|
||||
if err != nil {
|
||||
log.Fatalf("DingTalk outbound image init failed: %v", err)
|
||||
}
|
||||
feedbackService, err := feedback.NewService(dingCfg.TenantKey, "dingtalk", store)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
@@ -219,7 +231,12 @@ func main() {
|
||||
dingtalk.TenantConfig{TenantKey: dingCfg.TenantKey, Channel: "dingtalk"},
|
||||
qaService,
|
||||
sender,
|
||||
dingtalk.HandlerOptions{ImageDownloader: downloader, CardSender: cardSender, Directory: directoryResolver},
|
||||
dingtalk.HandlerOptions{
|
||||
ImageDownloader: downloader,
|
||||
CardSender: cardSender,
|
||||
Directory: directoryResolver,
|
||||
OutboundRewriter: outboundRewriter,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
|
||||
@@ -115,3 +115,32 @@ func TestMainWiresOptionalDingTalkCardDebugUser(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOutboundImageRewriterConfig(t *testing.T) {
|
||||
t.Run("disabled", func(t *testing.T) {
|
||||
rewriter, err := buildOutboundImageRewriter(outboundImageConfig{})
|
||||
if err != nil || rewriter != nil {
|
||||
t.Fatalf("rewriter=%v err=%v", rewriter, err)
|
||||
}
|
||||
})
|
||||
t.Run("partial configuration rejected", func(t *testing.T) {
|
||||
_, err := buildOutboundImageRewriter(outboundImageConfig{Bucket: "private-bucket"})
|
||||
if err == nil {
|
||||
t.Fatal("expected configuration error")
|
||||
}
|
||||
})
|
||||
t.Run("complete configuration", func(t *testing.T) {
|
||||
rewriter, err := buildOutboundImageRewriter(outboundImageConfig{
|
||||
WeKnoraBaseURL: "http://127.0.0.1:8080",
|
||||
WeKnoraAPIKey: "api-key",
|
||||
Endpoint: "https://oss-cn-shenzhen.aliyuncs.com",
|
||||
Region: "cn-shenzhen",
|
||||
AccessKeyID: "access-key-id",
|
||||
AccessKeySecret: "access-key-secret",
|
||||
Bucket: "private-bucket",
|
||||
})
|
||||
if err != nil || rewriter == nil {
|
||||
t.Fatalf("rewriter=%v err=%v", rewriter, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/channels/dingtalk"
|
||||
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/outboundimage"
|
||||
)
|
||||
|
||||
const (
|
||||
outboundImageMaxCount = 5
|
||||
outboundImageMaxBytes = 10 << 20
|
||||
outboundImageURLTTL = 2 * time.Hour
|
||||
outboundImagePrefix = "dobotai/im-images"
|
||||
)
|
||||
|
||||
type outboundImageConfig struct {
|
||||
WeKnoraBaseURL string
|
||||
WeKnoraAPIKey string
|
||||
Endpoint string
|
||||
Region string
|
||||
AccessKeyID string
|
||||
AccessKeySecret string
|
||||
Bucket string
|
||||
}
|
||||
|
||||
func buildOutboundImageRewriter(cfg outboundImageConfig) (dingtalk.OutboundContentRewriter, error) {
|
||||
ossValues := []string{cfg.Endpoint, cfg.Region, cfg.AccessKeyID, cfg.AccessKeySecret, cfg.Bucket}
|
||||
enabled := false
|
||||
for _, value := range ossValues {
|
||||
if strings.TrimSpace(value) != "" {
|
||||
enabled = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !enabled {
|
||||
return nil, nil
|
||||
}
|
||||
for _, value := range append([]string{cfg.WeKnoraBaseURL, cfg.WeKnoraAPIKey}, ossValues...) {
|
||||
if strings.TrimSpace(value) == "" {
|
||||
return nil, errors.New("outbound image: complete WeKnora and OSS configuration is required")
|
||||
}
|
||||
}
|
||||
transport := http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.Proxy = nil
|
||||
fetcher, err := outboundimage.NewWeKnoraFetcher(
|
||||
cfg.WeKnoraBaseURL,
|
||||
cfg.WeKnoraAPIKey,
|
||||
&http.Client{Timeout: 30 * time.Second, Transport: transport},
|
||||
outboundImageMaxBytes,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
store, err := outboundimage.NewAliyunOSSStore(
|
||||
cfg.Endpoint,
|
||||
cfg.Region,
|
||||
cfg.AccessKeyID,
|
||||
cfg.AccessKeySecret,
|
||||
cfg.Bucket,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
publisher := outboundimage.NewHashPublisher(store, outboundImagePrefix, outboundImageURLTTL)
|
||||
return outboundimage.NewRewriter(fetcher, publisher, outboundImageMaxCount), nil
|
||||
}
|
||||
@@ -3,6 +3,7 @@ module gitea.zheromiku.top/Weknora/dobotai-platform
|
||||
go 1.25.0
|
||||
|
||||
require (
|
||||
github.com/aliyun/alibabacloud-oss-go-sdk-v2 v1.5.1
|
||||
github.com/jackc/pgx/v5 v5.9.2
|
||||
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
|
||||
)
|
||||
@@ -15,4 +16,5 @@ require (
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
golang.org/x/sync v0.20.0 // indirect
|
||||
golang.org/x/text v0.37.0 // indirect
|
||||
golang.org/x/time v0.4.0 // indirect
|
||||
)
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
github.com/aliyun/alibabacloud-oss-go-sdk-v2 v1.5.1 h1:vtiFd0hhPAbyYJjztl0wYUq/PqEGkIlDmVuTIy6zw8Y=
|
||||
github.com/aliyun/alibabacloud-oss-go-sdk-v2 v1.5.1/go.mod h1:FTzydeQVmR24FI0D6XWUOMKckjXehM/jgMn1xC+DA9M=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
@@ -26,6 +28,8 @@ golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
|
||||
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
|
||||
golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
|
||||
golang.org/x/time v0.4.0 h1:Z81tqI5ddIoXDPvVQ7/7CC9TnLM7ubaFG2qXYd5BbYY=
|
||||
golang.org/x/time v0.4.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
|
||||
@@ -79,6 +79,10 @@ type UserDirectory interface {
|
||||
ResolveUser(ctx context.Context, tenantKey, staffID, fallbackName string) (directory.Profile, error)
|
||||
}
|
||||
|
||||
type OutboundContentRewriter interface {
|
||||
Rewrite(ctx context.Context, content string) string
|
||||
}
|
||||
|
||||
// FinalCardSender sends an AI card as the final reply. Implementations must
|
||||
// not log err, outTrackID, or content at any level.
|
||||
type FinalCardSender interface {
|
||||
@@ -93,9 +97,10 @@ type StreamingCardSender interface {
|
||||
}
|
||||
|
||||
type HandlerOptions struct {
|
||||
ImageDownloader ImageDownloader
|
||||
CardSender FinalCardSender
|
||||
Directory UserDirectory
|
||||
ImageDownloader ImageDownloader
|
||||
CardSender FinalCardSender
|
||||
Directory UserDirectory
|
||||
OutboundRewriter OutboundContentRewriter
|
||||
}
|
||||
|
||||
type Handler struct {
|
||||
@@ -105,6 +110,7 @@ type Handler struct {
|
||||
downloader ImageDownloader
|
||||
card FinalCardSender
|
||||
directory UserDirectory
|
||||
outboundRewriter OutboundContentRewriter
|
||||
now func() time.Time
|
||||
cardUpdateInterval time.Duration
|
||||
finalDeliveryTimeout time.Duration
|
||||
@@ -146,6 +152,7 @@ func NewHandlerWithOptions(cfg TenantConfig, qa QAAsker, reply MarkdownSender, o
|
||||
downloader: opts.ImageDownloader,
|
||||
card: opts.CardSender,
|
||||
directory: opts.Directory,
|
||||
outboundRewriter: opts.OutboundRewriter,
|
||||
now: time.Now,
|
||||
cardUpdateInterval: 500 * time.Millisecond,
|
||||
finalDeliveryTimeout: defaultHTTPTimeout,
|
||||
@@ -290,6 +297,7 @@ func (h *Handler) streamingCapable() (StreamingCardSender, StreamingQAAsker) {
|
||||
func (h *Handler) handleNonStreaming(ctx context.Context, webhookURL string, req core.QARequest) error {
|
||||
result, err := h.qa.Ask(ctx, req)
|
||||
if err != nil {
|
||||
logQAFailure("non_streaming", err)
|
||||
if replyErr := h.sendReply(ctx, webhookURL, qaErrorReply); replyErr != nil {
|
||||
return replyErr
|
||||
}
|
||||
@@ -298,13 +306,14 @@ func (h *Handler) handleNonStreaming(ctx context.Context, webhookURL string, req
|
||||
if result.Replayed && result.DeliveryStatus == core.DeliverySent {
|
||||
return nil
|
||||
}
|
||||
content := h.rewriteOutboundContent(ctx, result.Answer)
|
||||
|
||||
if h.card != nil {
|
||||
cardReq := CardRequest{
|
||||
ConversationType: req.ConversationType,
|
||||
ExternalUserID: req.ExternalUserID,
|
||||
ExternalChatID: req.ExternalChatID,
|
||||
Content: result.Answer,
|
||||
Content: content,
|
||||
}
|
||||
if outTrackID, cardErr := h.card.SendFinalCard(ctx, cardReq); cardErr == nil {
|
||||
h.recordDelivery(ctx, result, core.DeliverySent, core.DeliveryMethodCard, outTrackID, "")
|
||||
@@ -312,7 +321,7 @@ func (h *Handler) handleNonStreaming(ctx context.Context, webhookURL string, req
|
||||
}
|
||||
log.Printf("DingTalk card delivery failed, falling back to markdown")
|
||||
}
|
||||
if err := h.sendReply(ctx, webhookURL, result.Answer); err != nil {
|
||||
if err := h.sendReply(ctx, webhookURL, content); err != nil {
|
||||
h.recordDelivery(ctx, result, core.DeliveryFailed, "", "", core.DeliveryFailureChannelDelivery)
|
||||
return err
|
||||
}
|
||||
@@ -333,6 +342,7 @@ func (h *Handler) handleStreaming(ctx context.Context, webhookURL string, req co
|
||||
// Fallback: Ask → Markdown only (no SendFinalCard, no AskStream)
|
||||
result, askErr := h.qa.Ask(ctx, req)
|
||||
if askErr != nil {
|
||||
logQAFailure("streaming_fallback", askErr)
|
||||
if replyErr := h.sendReply(ctx, webhookURL, qaErrorReply); replyErr != nil {
|
||||
return replyErr
|
||||
}
|
||||
@@ -341,7 +351,8 @@ func (h *Handler) handleStreaming(ctx context.Context, webhookURL string, req co
|
||||
if result.Replayed && result.DeliveryStatus == core.DeliverySent {
|
||||
return nil
|
||||
}
|
||||
if replyErr := h.sendReply(ctx, webhookURL, result.Answer); replyErr != nil {
|
||||
content := h.rewriteOutboundContent(ctx, result.Answer)
|
||||
if replyErr := h.sendReply(ctx, webhookURL, content); replyErr != nil {
|
||||
h.recordDelivery(ctx, result, core.DeliveryFailed, "", "", core.DeliveryFailureChannelDelivery)
|
||||
return replyErr
|
||||
}
|
||||
@@ -369,6 +380,7 @@ func (h *Handler) handleStreaming(ctx context.Context, webhookURL string, req co
|
||||
|
||||
result, err := streamingQA.AskStream(ctx, req, observer)
|
||||
if err != nil || strings.TrimSpace(result.Answer) == "" {
|
||||
logQAFailure("streaming", errOrEmptyAnswer(err, result.Answer))
|
||||
return h.finalizeCardOrMarkdown(ctx, webhookURL, streamingCard, outTrackID, core.QAResult{}, qaErrorReply, true)
|
||||
}
|
||||
return h.finalizeCardOrMarkdown(ctx, webhookURL, streamingCard, outTrackID, result, result.Answer, false)
|
||||
@@ -386,12 +398,14 @@ func (h *Handler) handleEventStreaming(ctx context.Context, webhookURL string, r
|
||||
log.Printf("DingTalk event card start failed, falling back to markdown")
|
||||
result, askErr := h.qa.Ask(ctx, req)
|
||||
if askErr != nil {
|
||||
logQAFailure("event_fallback", askErr)
|
||||
return h.sendReply(ctx, webhookURL, qaErrorReply)
|
||||
}
|
||||
if result.Replayed && result.DeliveryStatus == core.DeliverySent {
|
||||
return nil
|
||||
}
|
||||
if replyErr := h.sendReply(ctx, webhookURL, result.Answer); replyErr != nil {
|
||||
content := h.rewriteOutboundContent(ctx, result.Answer)
|
||||
if replyErr := h.sendReply(ctx, webhookURL, content); replyErr != nil {
|
||||
h.recordDelivery(ctx, result, core.DeliveryFailed, "", "", core.DeliveryFailureChannelDelivery)
|
||||
return replyErr
|
||||
}
|
||||
@@ -421,6 +435,7 @@ func (h *Handler) handleEventStreaming(ctx context.Context, webhookURL string, r
|
||||
pump.Stop()
|
||||
|
||||
if err != nil || strings.TrimSpace(result.Answer) == "" {
|
||||
logQAFailure("event_streaming", errOrEmptyAnswer(err, result.Answer))
|
||||
return h.finalizeCardOrMarkdown(ctx, webhookURL, streamingCard, outTrackID, core.QAResult{}, qaErrorReply, true)
|
||||
}
|
||||
return h.finalizeCardOrMarkdown(ctx, webhookURL, streamingCard, outTrackID, result, result.Answer, false)
|
||||
@@ -432,6 +447,10 @@ func (h *Handler) finalizeCardOrMarkdown(ctx context.Context, webhookURL string,
|
||||
timeout = defaultHTTPTimeout
|
||||
}
|
||||
|
||||
if !isError {
|
||||
content = h.rewriteOutboundContent(ctx, content)
|
||||
}
|
||||
|
||||
// Final card convergence must still be attempted when the QA context was
|
||||
// canceled or timed out; otherwise the card can remain "processing".
|
||||
finalCtx, cancelFinal := context.WithTimeout(context.WithoutCancel(ctx), timeout)
|
||||
@@ -460,6 +479,19 @@ func (h *Handler) finalizeCardOrMarkdown(ctx context.Context, webhookURL string,
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *Handler) rewriteOutboundContent(ctx context.Context, content string) string {
|
||||
if h.outboundRewriter == nil {
|
||||
return content
|
||||
}
|
||||
timeout := h.finalDeliveryTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = defaultHTTPTimeout
|
||||
}
|
||||
rewriteCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), timeout)
|
||||
defer cancel()
|
||||
return h.outboundRewriter.Rewrite(rewriteCtx, content)
|
||||
}
|
||||
|
||||
func (h *Handler) recordDelivery(ctx context.Context, result core.QAResult, status, method, providerMessageID, failureCode string) {
|
||||
if strings.TrimSpace(result.TurnID) == "" || strings.TrimSpace(result.AnswerMessageID) == "" {
|
||||
return
|
||||
@@ -511,6 +543,50 @@ func renderProgressEvent(event weknora.StreamEvent) string {
|
||||
return ""
|
||||
}
|
||||
|
||||
func errOrEmptyAnswer(err error, answer string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strings.TrimSpace(answer) == "" {
|
||||
return weknora.ErrEmptyAnswer
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func logQAFailure(mode string, err error) {
|
||||
log.Printf("DingTalk QA failed: mode=%s stage=%s", mode, qaFailureStage(err))
|
||||
}
|
||||
|
||||
func qaFailureStage(err error) string {
|
||||
switch {
|
||||
case err == nil:
|
||||
return "unknown"
|
||||
case errors.Is(err, context.Canceled):
|
||||
return "context_canceled"
|
||||
case errors.Is(err, context.DeadlineExceeded):
|
||||
return "context_deadline"
|
||||
case errors.Is(err, core.ErrTurnInProgress):
|
||||
return "duplicate_or_in_progress"
|
||||
case errors.Is(err, weknora.ErrEmptyAnswer):
|
||||
return "empty_answer"
|
||||
}
|
||||
message := strings.ToLower(err.Error())
|
||||
switch {
|
||||
case strings.Contains(message, "begin turn"):
|
||||
return "begin_turn"
|
||||
case strings.Contains(message, "agent") && strings.Contains(message, "binding"):
|
||||
return "agent_binding"
|
||||
case strings.Contains(message, "session"):
|
||||
return "session"
|
||||
case strings.Contains(message, "weknora agent chat"):
|
||||
return "agent_chat"
|
||||
case strings.Contains(message, "required") || strings.Contains(message, "unsupported"):
|
||||
return "validation"
|
||||
default:
|
||||
return "unknown"
|
||||
}
|
||||
}
|
||||
|
||||
func imageFailureStage(err error) string {
|
||||
message := strings.ToLower(err.Error())
|
||||
switch {
|
||||
|
||||
@@ -726,6 +726,46 @@ func TestHandleCardSuccessDirectFieldsCorrectNoMarkdown(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleNonStreamingRewritesFinalCardAndMarkdownFallback(t *testing.T) {
|
||||
resource := "resource://AbCdEfGhIjKlMnOpQrStUv"
|
||||
signed := "https://img.example/signed.jpg"
|
||||
answer := "说明\n"
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
cardErr error
|
||||
wantReply bool
|
||||
}{
|
||||
{name: "card"},
|
||||
{name: "markdown fallback", cardErr: errors.New("card failed"), wantReply: true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
qa := &fakeQA{result: core.QAResult{Answer: answer}}
|
||||
reply := &fakeReply{}
|
||||
card := &fakeCard{returnID: "track", err: tc.cardErr}
|
||||
h := newTestHandlerWithCard(t, qa, reply, card)
|
||||
rewriter := &fakeOutboundRewriter{from: resource, to: signed}
|
||||
h.outboundRewriter = rewriter
|
||||
|
||||
if err := h.Handle(context.Background(), directTextCallback()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !card.called || reply.called != tc.wantReply {
|
||||
t.Fatalf("card=%t reply=%t", card.called, reply.called)
|
||||
}
|
||||
content := card.lastReq.Content
|
||||
if tc.wantReply {
|
||||
content = reply.lastText
|
||||
}
|
||||
if strings.Contains(content, resource) || !strings.Contains(content, signed) {
|
||||
t.Fatalf("outbound content=%q", content)
|
||||
}
|
||||
if rewriter.calls != 1 {
|
||||
t.Fatalf("rewrite calls=%d, want 1", rewriter.calls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleCardSuccessGroupChatIDCorrect(t *testing.T) {
|
||||
qa := &fakeQA{result: core.QAResult{Answer: "group card answer"}}
|
||||
reply := &fakeReply{}
|
||||
@@ -940,6 +980,17 @@ type cardUpdate struct {
|
||||
isError bool
|
||||
}
|
||||
|
||||
type fakeOutboundRewriter struct {
|
||||
calls int
|
||||
from string
|
||||
to string
|
||||
}
|
||||
|
||||
func (f *fakeOutboundRewriter) Rewrite(_ context.Context, content string) string {
|
||||
f.calls++
|
||||
return strings.ReplaceAll(content, f.from, f.to)
|
||||
}
|
||||
|
||||
func (f *fakeStreamingCard) SendFinalCard(_ context.Context, req CardRequest) (string, error) {
|
||||
f.finalCalled = true
|
||||
f.finalReq = req
|
||||
@@ -1123,11 +1174,14 @@ func TestHandleStreamingThrottleClockAdvances(t *testing.T) {
|
||||
|
||||
func TestHandleStreamingStartFailsFallbackAskMarkdown(t *testing.T) {
|
||||
card := &fakeStreamingCard{startErr: errors.New("card API unavailable")}
|
||||
resource := "resource://AbCdEfGhIjKlMnOpQrStUv"
|
||||
qa := &fakeStreamingQA{
|
||||
result: core.QAResult{Answer: "**fallback markdown**"},
|
||||
result: core.QAResult{Answer: "**fallback markdown**\n"},
|
||||
}
|
||||
reply := &fakeReply{}
|
||||
h := newTestHandlerWithStreaming(t, qa, reply, card)
|
||||
rewriter := &fakeOutboundRewriter{from: resource, to: "https://img.example/signed.jpg"}
|
||||
h.outboundRewriter = rewriter
|
||||
|
||||
err := h.Handle(context.Background(), directTextCallback())
|
||||
if err != nil {
|
||||
@@ -1156,8 +1210,11 @@ func TestHandleStreamingStartFailsFallbackAskMarkdown(t *testing.T) {
|
||||
if !reply.called {
|
||||
t.Fatal("SendMarkdown must be called on Start failure")
|
||||
}
|
||||
if reply.lastText != "**fallback markdown**" {
|
||||
t.Fatalf("markdown text = %q, want fallback markdown", reply.lastText)
|
||||
if strings.Contains(reply.lastText, resource) || !strings.Contains(reply.lastText, "https://img.example/signed.jpg") {
|
||||
t.Fatalf("markdown text = %q, want rewritten fallback markdown", reply.lastText)
|
||||
}
|
||||
if rewriter.calls != 1 {
|
||||
t.Fatalf("rewrite calls=%d, want 1", rewriter.calls)
|
||||
}
|
||||
|
||||
// No card updates
|
||||
@@ -2005,5 +2062,32 @@ func TestTriggerPersonSnapshotRegression(t *testing.T) {
|
||||
t.Fatalf("SenderName = %q, want 张三", qa.lastReq.SenderName)
|
||||
}
|
||||
})
|
||||
|
||||
}
|
||||
|
||||
func TestHandleStreamingRewritesOnlyFinalOutboundImageContent(t *testing.T) {
|
||||
card := &fakeStreamingCard{startID: "track-image-final"}
|
||||
resource := "resource://AbCdEfGhIjKlMnOpQrStUv"
|
||||
answer := "操作说明\n"
|
||||
qa := &fakeStreamingQA{chunks: []string{answer}, result: core.QAResult{Answer: answer}}
|
||||
reply := &fakeReply{}
|
||||
rewriter := &fakeOutboundRewriter{from: resource, to: "https://img.example/signed.jpg"}
|
||||
h := newTestHandlerWithStreaming(t, qa, reply, card)
|
||||
h.outboundRewriter = rewriter
|
||||
|
||||
if err := h.Handle(context.Background(), directTextCallback()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(card.updates) < 2 {
|
||||
t.Fatalf("updates=%d", len(card.updates))
|
||||
}
|
||||
if !strings.Contains(card.updates[0].content, resource) {
|
||||
t.Fatalf("streaming intermediate content should remain untouched: %q", card.updates[0].content)
|
||||
}
|
||||
final := card.updates[len(card.updates)-1]
|
||||
if !final.isFinalize || strings.Contains(final.content, resource) || !strings.Contains(final.content, "https://img.example/signed.jpg") {
|
||||
t.Fatalf("final content=%q finalize=%t", final.content, final.isFinalize)
|
||||
}
|
||||
if rewriter.calls != 1 {
|
||||
t.Fatalf("rewrite calls=%d, want 1", rewriter.calls)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss"
|
||||
"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss/credentials"
|
||||
)
|
||||
|
||||
type AliyunOSSStore struct {
|
||||
client *oss.Client
|
||||
bucket string
|
||||
}
|
||||
|
||||
func NewAliyunOSSStore(endpoint, region, accessKeyID, accessKeySecret, bucket string) (*AliyunOSSStore, error) {
|
||||
endpoint = strings.TrimSpace(endpoint)
|
||||
region = strings.TrimSpace(region)
|
||||
accessKeyID = strings.TrimSpace(accessKeyID)
|
||||
accessKeySecret = strings.TrimSpace(accessKeySecret)
|
||||
bucket = strings.TrimSpace(bucket)
|
||||
if endpoint == "" || region == "" || accessKeyID == "" || accessKeySecret == "" || bucket == "" {
|
||||
return nil, errors.New("outbound image: OSS endpoint, region, credentials, and bucket are required")
|
||||
}
|
||||
provider := credentials.NewStaticCredentialsProvider(accessKeyID, accessKeySecret, "")
|
||||
config := oss.LoadDefaultConfig().
|
||||
WithCredentialsProvider(provider).
|
||||
WithRegion(region).
|
||||
WithEndpoint(endpoint)
|
||||
return &AliyunOSSStore{client: oss.NewClient(config), bucket: bucket}, nil
|
||||
}
|
||||
|
||||
func (s *AliyunOSSStore) Put(ctx context.Context, key string, data []byte, contentType string) error {
|
||||
if s == nil || s.client == nil || strings.TrimSpace(key) == "" || len(data) == 0 {
|
||||
return errors.New("outbound image: invalid OSS put request")
|
||||
}
|
||||
_, err := s.client.PutObject(ctx, &oss.PutObjectRequest{
|
||||
Bucket: oss.Ptr(s.bucket),
|
||||
Key: oss.Ptr(key),
|
||||
Body: bytes.NewReader(data),
|
||||
ContentType: oss.Ptr(contentType),
|
||||
})
|
||||
if err != nil {
|
||||
return errors.New("outbound image: OSS put failed")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *AliyunOSSStore) PresignGet(ctx context.Context, key string, ttl time.Duration) (string, error) {
|
||||
if s == nil || s.client == nil || strings.TrimSpace(key) == "" || ttl <= 0 {
|
||||
return "", errors.New("outbound image: invalid OSS presign request")
|
||||
}
|
||||
result, err := s.client.Presign(ctx, &oss.GetObjectRequest{
|
||||
Bucket: oss.Ptr(s.bucket),
|
||||
Key: oss.Ptr(key),
|
||||
}, oss.PresignExpires(ttl))
|
||||
if err != nil {
|
||||
return "", errors.New("outbound image: OSS presign failed")
|
||||
}
|
||||
return result.URL, nil
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"path"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type ObjectStore interface {
|
||||
Put(ctx context.Context, key string, data []byte, contentType string) error
|
||||
PresignGet(ctx context.Context, key string, ttl time.Duration) (string, error)
|
||||
}
|
||||
|
||||
type HashPublisher struct {
|
||||
store ObjectStore
|
||||
prefix string
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func NewHashPublisher(store ObjectStore, prefix string, ttl time.Duration) *HashPublisher {
|
||||
return &HashPublisher{store: store, prefix: strings.Trim(strings.TrimSpace(prefix), "/"), ttl: ttl}
|
||||
}
|
||||
|
||||
func (p *HashPublisher) Publish(ctx context.Context, data []byte, contentType string) (string, error) {
|
||||
if p == nil || p.store == nil || len(data) == 0 || p.prefix == "" || p.ttl <= 0 {
|
||||
return "", errors.New("outbound image: invalid publisher configuration")
|
||||
}
|
||||
extension := imageExtension(contentType)
|
||||
if extension == "" {
|
||||
return "", errors.New("outbound image: unsupported image content type")
|
||||
}
|
||||
sum := sha256.Sum256(data)
|
||||
digest := hex.EncodeToString(sum[:])
|
||||
key := path.Join(p.prefix, digest[:2], digest+extension)
|
||||
if err := p.store.Put(ctx, key, data, contentType); err != nil {
|
||||
return "", errors.New("outbound image: publish object failed")
|
||||
}
|
||||
signedURL, err := p.store.PresignGet(ctx, key, p.ttl)
|
||||
if err != nil {
|
||||
return "", errors.New("outbound image: sign object URL failed")
|
||||
}
|
||||
if !strings.HasPrefix(strings.TrimSpace(signedURL), "https://") {
|
||||
return "", errors.New("outbound image: signed object URL must use HTTPS")
|
||||
}
|
||||
return signedURL, nil
|
||||
}
|
||||
|
||||
func imageExtension(contentType string) string {
|
||||
switch strings.ToLower(strings.TrimSpace(strings.Split(contentType, ";")[0])) {
|
||||
case "image/jpeg", "image/jpg":
|
||||
return ".jpg"
|
||||
case "image/png":
|
||||
return ".png"
|
||||
case "image/webp":
|
||||
return ".webp"
|
||||
case "image/gif":
|
||||
return ".gif"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type fakeObjectStore struct {
|
||||
key string
|
||||
data []byte
|
||||
contentType string
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
func (s *fakeObjectStore) Put(_ context.Context, key string, data []byte, contentType string) error {
|
||||
s.key = key
|
||||
s.data = append([]byte(nil), data...)
|
||||
s.contentType = contentType
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *fakeObjectStore) PresignGet(_ context.Context, key string, ttl time.Duration) (string, error) {
|
||||
s.ttl = ttl
|
||||
return "https://private-oss.example/" + key + "?signature=short", nil
|
||||
}
|
||||
|
||||
func TestHashPublisherUsesDeterministicTemporaryKeyAndTwoHourHTTPSURL(t *testing.T) {
|
||||
store := &fakeObjectStore{}
|
||||
publisher := NewHashPublisher(store, "dobotai/im-images", 2*time.Hour)
|
||||
|
||||
first, err := publisher.Publish(context.Background(), []byte("same-image"), "image/jpeg")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
firstKey := store.key
|
||||
second, err := publisher.Publish(context.Background(), []byte("same-image"), "image/jpeg")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if first != second || firstKey != store.key {
|
||||
t.Fatalf("publish is not deterministic: url1=%q url2=%q key1=%q key2=%q", first, second, firstKey, store.key)
|
||||
}
|
||||
if !strings.HasPrefix(store.key, "dobotai/im-images/") || !strings.HasSuffix(store.key, ".jpg") {
|
||||
t.Fatalf("key=%q", store.key)
|
||||
}
|
||||
if strings.Contains(store.key, "same-image") {
|
||||
t.Fatalf("key leaks content: %q", store.key)
|
||||
}
|
||||
if store.ttl != 2*time.Hour || !strings.HasPrefix(first, "https://") {
|
||||
t.Fatalf("ttl=%s url=%q", store.ttl, first)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var resourceMarkdownImageRE = regexp.MustCompile(`!\[[^\]]*\]\((resource://[0-9A-Za-z_-]+)\)`)
|
||||
|
||||
type Fetcher interface {
|
||||
Fetch(ctx context.Context, reference string) ([]byte, string, error)
|
||||
}
|
||||
|
||||
type Publisher interface {
|
||||
Publish(ctx context.Context, data []byte, contentType string) (string, error)
|
||||
}
|
||||
|
||||
type Rewriter struct {
|
||||
fetcher Fetcher
|
||||
publisher Publisher
|
||||
maxImages int
|
||||
}
|
||||
|
||||
func NewRewriter(fetcher Fetcher, publisher Publisher, maxImages int) *Rewriter {
|
||||
return &Rewriter{fetcher: fetcher, publisher: publisher, maxImages: maxImages}
|
||||
}
|
||||
|
||||
func (r *Rewriter) Rewrite(ctx context.Context, markdown string) string {
|
||||
if r == nil || r.fetcher == nil || r.publisher == nil || r.maxImages <= 0 {
|
||||
return resourceMarkdownImageRE.ReplaceAllString(markdown, "")
|
||||
}
|
||||
urls := map[string]string{}
|
||||
attempted := 0
|
||||
return resourceMarkdownImageRE.ReplaceAllStringFunc(markdown, func(image string) string {
|
||||
match := resourceMarkdownImageRE.FindStringSubmatch(image)
|
||||
if len(match) != 2 {
|
||||
return ""
|
||||
}
|
||||
reference := match[1]
|
||||
if signedURL, ok := urls[reference]; ok {
|
||||
if signedURL == "" {
|
||||
return ""
|
||||
}
|
||||
return strings.Replace(image, reference, signedURL, 1)
|
||||
}
|
||||
if attempted >= r.maxImages {
|
||||
urls[reference] = ""
|
||||
return ""
|
||||
}
|
||||
attempted++
|
||||
data, contentType, err := r.fetcher.Fetch(ctx, reference)
|
||||
if err != nil {
|
||||
urls[reference] = ""
|
||||
return ""
|
||||
}
|
||||
signedURL, err := r.publisher.Publish(ctx, data, contentType)
|
||||
if err != nil || !strings.HasPrefix(signedURL, "https://") {
|
||||
urls[reference] = ""
|
||||
return ""
|
||||
}
|
||||
urls[reference] = signedURL
|
||||
return strings.Replace(image, reference, signedURL, 1)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type fakeFetcher struct {
|
||||
calls map[string]int
|
||||
fail map[string]bool
|
||||
}
|
||||
|
||||
func (f *fakeFetcher) Fetch(_ context.Context, ref string) ([]byte, string, error) {
|
||||
if f.calls == nil {
|
||||
f.calls = map[string]int{}
|
||||
}
|
||||
f.calls[ref]++
|
||||
if f.fail[ref] {
|
||||
return nil, "", errors.New("fetch failed")
|
||||
}
|
||||
return []byte("image:" + ref), "image/png", nil
|
||||
}
|
||||
|
||||
type fakePublisher struct {
|
||||
calls int
|
||||
}
|
||||
|
||||
func (p *fakePublisher) Publish(_ context.Context, data []byte, contentType string) (string, error) {
|
||||
p.calls++
|
||||
if contentType != "image/png" || len(data) == 0 {
|
||||
return "", errors.New("bad image")
|
||||
}
|
||||
return "https://img.example/" + string(rune('0'+p.calls)) + ".png?sig=short", nil
|
||||
}
|
||||
|
||||
func TestRewriteMarkdownResourceImagesDeduplicatesAndPreservesOtherURLs(t *testing.T) {
|
||||
fetcher := &fakeFetcher{}
|
||||
publisher := &fakePublisher{}
|
||||
r := NewRewriter(fetcher, publisher, 5)
|
||||
input := "步骤一\n\n重复\n\n外部\n"
|
||||
|
||||
got := r.Rewrite(context.Background(), input)
|
||||
|
||||
if strings.Contains(got, "resource://") {
|
||||
t.Fatalf("resource reference remained: %s", got)
|
||||
}
|
||||
if strings.Count(got, "https://img.example/1.png?sig=short") != 2 {
|
||||
t.Fatalf("deduplicated signed URL count mismatch: %s", got)
|
||||
}
|
||||
if !strings.Contains(got, "") {
|
||||
t.Fatalf("ordinary HTTPS image changed: %s", got)
|
||||
}
|
||||
if fetcher.calls["resource://AbCdEfGhIjKlMnOpQrStUv"] != 1 || publisher.calls != 1 {
|
||||
t.Fatalf("calls fetch=%v publish=%d", fetcher.calls, publisher.calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteMarkdownResourceImagesDropsFailuresAndImagesBeyondLimit(t *testing.T) {
|
||||
failed := "resource://FailedRefAbCdEfGhIjKlMn"
|
||||
fetcher := &fakeFetcher{fail: map[string]bool{failed: true}}
|
||||
publisher := &fakePublisher{}
|
||||
r := NewRewriter(fetcher, publisher, 3)
|
||||
input := "保留正文\n\n\n\n"
|
||||
|
||||
got := r.Rewrite(context.Background(), input)
|
||||
|
||||
if !strings.Contains(got, "保留正文") {
|
||||
t.Fatalf("text was lost: %s", got)
|
||||
}
|
||||
if strings.Contains(got, "resource://") || strings.Contains(got, "坏图") || strings.Contains(got, "![三]") {
|
||||
t.Fatalf("broken or over-limit image remained: %s", got)
|
||||
}
|
||||
if publisher.calls != 2 {
|
||||
t.Fatalf("publisher calls=%d, want 2", publisher.calls)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type WeKnoraFetcher struct {
|
||||
baseURL *url.URL
|
||||
apiKey string
|
||||
client *http.Client
|
||||
maxBytes int64
|
||||
}
|
||||
|
||||
func NewWeKnoraFetcher(baseURL, apiKey string, client *http.Client, maxBytes int64) (*WeKnoraFetcher, error) {
|
||||
parsed, err := url.Parse(strings.TrimRight(strings.TrimSpace(baseURL), "/"))
|
||||
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
||||
return nil, errors.New("outbound image: invalid WeKnora base URL")
|
||||
}
|
||||
if strings.TrimSpace(apiKey) == "" {
|
||||
return nil, errors.New("outbound image: WeKnora API key is required")
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
return nil, errors.New("outbound image: positive max bytes is required")
|
||||
}
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
return &WeKnoraFetcher{baseURL: parsed, apiKey: apiKey, client: client, maxBytes: maxBytes}, nil
|
||||
}
|
||||
|
||||
func (f *WeKnoraFetcher) Fetch(ctx context.Context, reference string) ([]byte, string, error) {
|
||||
if !strings.HasPrefix(reference, "resource://") || !resourceMarkdownImageRE.MatchString("") {
|
||||
return nil, "", errors.New("outbound image: unsupported resource reference")
|
||||
}
|
||||
u := *f.baseURL
|
||||
u.Path = "/files"
|
||||
u.RawPath = ""
|
||||
query := u.Query()
|
||||
query.Set("file_path", reference)
|
||||
u.RawQuery = query.Encode()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
return nil, "", errors.New("outbound image: create WeKnora request failed")
|
||||
}
|
||||
req.Header.Set("X-API-Key", f.apiKey)
|
||||
resp, err := f.client.Do(req)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("outbound image: WeKnora request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4096))
|
||||
return nil, "", fmt.Errorf("outbound image: WeKnora HTTP %d", resp.StatusCode)
|
||||
}
|
||||
contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0])
|
||||
if !strings.HasPrefix(contentType, "image/") {
|
||||
return nil, "", errors.New("outbound image: WeKnora response is not an image")
|
||||
}
|
||||
if length := strings.TrimSpace(resp.Header.Get("Content-Length")); length != "" {
|
||||
if n, convErr := strconv.ParseInt(length, 10, 64); convErr == nil && n > f.maxBytes {
|
||||
return nil, "", errors.New("outbound image: WeKnora image exceeds size limit")
|
||||
}
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(resp.Body, f.maxBytes+1))
|
||||
if err != nil {
|
||||
return nil, "", errors.New("outbound image: read WeKnora response failed")
|
||||
}
|
||||
if int64(len(data)) == 0 || int64(len(data)) > f.maxBytes {
|
||||
return nil, "", errors.New("outbound image: WeKnora image exceeds size limit")
|
||||
}
|
||||
return data, contentType, nil
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package outboundimage
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestWeKnoraFetcherDownloadsAuthenticatedResourceImage(t *testing.T) {
|
||||
var called bool
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||
called = true
|
||||
if req.URL.Path != "/files" {
|
||||
t.Fatalf("path=%q", req.URL.Path)
|
||||
}
|
||||
if got := req.URL.Query().Get("file_path"); got != "resource://AbCdEfGhIjKlMnOpQrStUv" {
|
||||
t.Fatalf("file_path=%q", got)
|
||||
}
|
||||
if got := req.Header.Get("X-API-Key"); got != "secret-test-key" {
|
||||
t.Fatalf("api key header=%q", got)
|
||||
}
|
||||
w.Header().Set("Content-Type", "image/jpeg")
|
||||
w.Header().Set("Content-Length", "4")
|
||||
_, _ = io.WriteString(w, "JPEG")
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
fetcher, err := NewWeKnoraFetcher(server.URL+"/api/v1", "secret-test-key", server.Client(), 10<<20)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, contentType, err := fetcher.Fetch(context.Background(), "resource://AbCdEfGhIjKlMnOpQrStUv")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !called || string(data) != "JPEG" || contentType != "image/jpeg" {
|
||||
t.Fatalf("called=%t data=%q type=%q", called, data, contentType)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWeKnoraFetcherRejectsUnsafeReferenceNonImageAndOversize(t *testing.T) {
|
||||
t.Run("unsafe reference", func(t *testing.T) {
|
||||
fetcher, err := NewWeKnoraFetcher("https://weknora.example", "key", http.DefaultClient, 4)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := fetcher.Fetch(context.Background(), "https://evil.example/a.png"); err == nil {
|
||||
t.Fatal("expected unsafe reference rejection")
|
||||
}
|
||||
})
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
contentType string
|
||||
body string
|
||||
}{
|
||||
{name: "non image", contentType: "text/plain", body: "TEXT"},
|
||||
{name: "oversize", contentType: "image/png", body: "12345"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", tc.contentType)
|
||||
_, _ = io.WriteString(w, tc.body)
|
||||
}))
|
||||
defer server.Close()
|
||||
fetcher, err := NewWeKnoraFetcher(server.URL, "key", server.Client(), 4)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, _, err := fetcher.Fetch(context.Background(), "resource://AbCdEfGhIjKlMnOpQrStUv"); err == nil {
|
||||
t.Fatal("expected response rejection")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user