From d6467f6eb029c0547f7b0d7b51768179de631103 Mon Sep 17 00:00:00 2001 From: wucm667 Date: Fri, 31 Jul 2026 15:04:56 +0800 Subject: [PATCH] fix(images): decode data URLs during task offload --- backend/internal/service/image_storage.go | 67 +++++++++++++- .../internal/service/image_storage_test.go | 89 +++++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) diff --git a/backend/internal/service/image_storage.go b/backend/internal/service/image_storage.go index a2f4882fa4..b160846b47 100644 --- a/backend/internal/service/image_storage.go +++ b/backend/internal/service/image_storage.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "io" + "mime" "net/http" "strconv" "strings" @@ -126,6 +127,9 @@ func (u *ImageResultUploader) fetchImageBytes(ctx context.Context, item map[stri var rawURL string if err := json.Unmarshal(raw, &rawURL); err == nil { if rawURL = strings.TrimSpace(rawURL); rawURL != "" { + if len(rawURL) >= len("data:") && strings.EqualFold(rawURL[:len("data:")], "data:") { + return u.decodeImageDataURL(rawURL) + } return u.download(ctx, rawURL) } } @@ -133,6 +137,60 @@ func (u *ImageResultUploader) fetchImageBytes(ctx context.Context, item map[stri return nil, "", errors.New("image item has neither b64_json nor url") } +func (u *ImageResultUploader) decodeImageDataURL(rawURL string) ([]byte, string, error) { + header, payload, ok := strings.Cut(rawURL[len("data:"):], ",") + if !ok { + return nil, "", errors.New("decode image data URL: missing comma separator") + } + + parts := strings.Split(header, ";") + if strings.TrimSpace(parts[0]) == "" { + return nil, "", errors.New("decode image data URL: missing media type") + } + base64Index := len(parts) - 1 + if base64Index < 1 || !strings.EqualFold(strings.TrimSpace(parts[base64Index]), "base64") { + for i := 1; i < base64Index; i++ { + if strings.EqualFold(strings.TrimSpace(parts[i]), "base64") { + return nil, "", errors.New("decode image data URL: base64 marker must be the final header token") + } + } + return nil, "", errors.New("decode image data URL: payload is not base64 encoded") + } + for i := 1; i < base64Index; i++ { + if strings.EqualFold(strings.TrimSpace(parts[i]), "base64") { + return nil, "", errors.New("decode image data URL: duplicate base64 marker") + } + } + mediaTypeHeader := strings.Join(parts[:base64Index], ";") + declaredType, _, err := mime.ParseMediaType(mediaTypeHeader) + if err != nil { + return nil, "", fmt.Errorf("decode image data URL: invalid media type: %w", err) + } + declaredType = strings.ToLower(declaredType) + if !strings.HasPrefix(declaredType, "image/") { + return nil, "", fmt.Errorf("decode image data URL: media type %q is not an image", declaredType) + } + + limit := u.maxDownloadBytes + if limit <= 0 { + limit = defaultImageMaxDownloadBytes + } + decoder := base64.NewDecoder(base64.StdEncoding, strings.NewReader(payload)) + data, err := io.ReadAll(io.LimitReader(decoder, limit+1)) + if err != nil { + return nil, "", fmt.Errorf("decode image data URL base64 payload: %w", err) + } + if int64(len(data)) > limit { + return nil, "", fmt.Errorf("decoded image data URL exceeds %d bytes", limit) + } + + contentType := detectedImageContentType(data) + if contentType == "" { + contentType = declaredType + } + return data, contentType, nil +} + func (u *ImageResultUploader) download(ctx context.Context, rawURL string) ([]byte, string, error) { req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) if err != nil { @@ -169,11 +227,18 @@ func (u *ImageResultUploader) buildKey(taskID string, index int, contentType str } func detectImageContentType(data []byte) string { + if ct := detectedImageContentType(data); ct != "" { + return ct + } + return "image/png" +} + +func detectedImageContentType(data []byte) string { ct := strings.TrimSpace(strings.Split(http.DetectContentType(data), ";")[0]) if strings.HasPrefix(ct, "image/") { return ct } - return "image/png" + return "" } func extensionForContentType(ct string) string { diff --git a/backend/internal/service/image_storage_test.go b/backend/internal/service/image_storage_test.go index fc9e5043e5..3a320671ee 100644 --- a/backend/internal/service/image_storage_test.go +++ b/backend/internal/service/image_storage_test.go @@ -29,6 +29,12 @@ type fakeImageStorage struct { err error } +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + func (f *fakeImageStorage) Save(_ context.Context, key, contentType string, data []byte) (string, error) { if f.err != nil { return "", f.err @@ -91,6 +97,89 @@ func TestImageResultUploaderRewritesURL(t *testing.T) { require.JSONEq(t, `"https://cdn.test/images/imgtask_xyz-0.png"`, string(parsed.Data[0]["url"])) } +func TestImageResultUploaderRewritesImageDataURLWithoutHTTP(t *testing.T) { + httpCalls := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + httpCalls++ + return nil, errors.New("HTTP must not be called for data URLs") + })} + storage := &fakeImageStorage{} + uploader := NewImageResultUploader(storage, "images/", 0, client) + b64 := base64.StdEncoding.EncodeToString(pngBytes) + result := json.RawMessage(`{"data":[{"url":"DATA:image/jpeg;name=photo.jpg;BaSe64,` + b64 + `","revised_prompt":"kept"}]}`) + + out, err := uploader.Rewrite(context.Background(), "imgtask_data", result) + require.NoError(t, err) + require.Zero(t, httpCalls) + require.Len(t, storage.saved, 1) + require.Equal(t, pngBytes, storage.saved[0].data) + require.Equal(t, "image/png", storage.saved[0].contentType, "detected bytes take precedence over a conflicting declaration") + require.Equal(t, "images/imgtask_data-0.png", storage.saved[0].key) + + var parsed struct { + Data []map[string]json.RawMessage `json:"data"` + } + require.NoError(t, json.Unmarshal(out, &parsed)) + require.JSONEq(t, `"https://cdn.test/images/imgtask_data-0.png"`, string(parsed.Data[0]["url"])) + require.JSONEq(t, `"kept"`, string(parsed.Data[0]["revised_prompt"])) +} + +func TestImageResultUploaderDataURLValidation(t *testing.T) { + tests := []struct { + name string + url string + wantErr string + }{ + {name: "missing comma", url: "data:image/png;base64", wantErr: "missing comma"}, + {name: "non image", url: "data:text/plain;base64,aGVsbG8=", wantErr: "is not an image"}, + {name: "non base64", url: "data:image/png,raw", wantErr: "not base64"}, + {name: "invalid base64", url: "data:image/png;base64,%%%", wantErr: "base64 payload"}, + {name: "invalid media type", url: "data:image/png;bad parameter;base64,aGVsbG8=", wantErr: "invalid media type"}, + {name: "parameter after base64", url: "data:image/png;base64;name=photo.png,aGVsbG8=", wantErr: "base64 marker must be the final header token"}, + {name: "duplicate base64 marker", url: "data:image/png;base64;base64,aGVsbG8=", wantErr: "duplicate base64 marker"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + httpCalls := 0 + client := &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) { + httpCalls++ + return nil, errors.New("HTTP must not be called for data URLs") + })} + uploader := NewImageResultUploader(&fakeImageStorage{}, "images/", 0, client) + result, err := json.Marshal(map[string]any{"data": []map[string]string{{"url": tt.url}}}) + require.NoError(t, err) + + _, err = uploader.Rewrite(context.Background(), "imgtask_bad", result) + require.ErrorContains(t, err, tt.wantErr) + require.Zero(t, httpCalls) + }) + } +} + +func TestImageResultUploaderRejectsOversizedImageDataURL(t *testing.T) { + storage := &fakeImageStorage{} + uploader := NewImageResultUploader(storage, "images/", 3, nil) + payload := base64.StdEncoding.EncodeToString([]byte("four")) + result := json.RawMessage(`{"data":[{"url":"data:image/png;base64,` + payload + `"}]}`) + + _, err := uploader.Rewrite(context.Background(), "imgtask_large", result) + require.ErrorContains(t, err, "decoded image data URL exceeds 3 bytes") + require.Empty(t, storage.saved) +} + +func TestImageResultUploaderB64JSONTakesPrecedenceOverDataURL(t *testing.T) { + storage := &fakeImageStorage{} + uploader := NewImageResultUploader(storage, "images/", 0, nil) + b64 := base64.StdEncoding.EncodeToString(pngBytes) + result := json.RawMessage(`{"data":[{"b64_json":"` + b64 + `","url":"data:text/plain,not-an-image"}]}`) + + _, err := uploader.Rewrite(context.Background(), "imgtask_precedence", result) + require.NoError(t, err) + require.Len(t, storage.saved, 1) + require.Equal(t, pngBytes, storage.saved[0].data) +} + func TestImageResultUploaderPropagatesStorageError(t *testing.T) { storage := &fakeImageStorage{err: errors.New("bucket unreachable")} uploader := NewImageResultUploader(storage, "images/", 0, nil)