fix(images): decode data URLs during task offload

2026-10-14 12:02:27 +08:00 · 2026-07-31 15:04:56 +08:00
parent f9d2791693
commit d6467f6eb0
2 changed files with 155 additions and 1 deletions
@@ -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 {
@@ -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)