image_url_compat.go

package openai

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/service"
	"github.com/QuantumNous/new-api/setting/system_setting"
)

const maxImageURLCompatBytes int64 = 40 * 1024 * 1024

const maxImageURLCompatDownloadAttempts = 3

type imageURLCompatDownloadFunc func(string, ...string) (*http.Response, error)

type imageURLCompatStats struct {
	Converted       int
	DownloadedBytes int64
	DownloadTime    time.Duration
	EncodeTime      time.Duration
	TotalTime       time.Duration
}

func downloadOpenAIImageURLCompat(rawURL string, reason ...string) (*http.Response, error) {
	parsedURL, err := url.Parse(rawURL)
	if err != nil {
		return nil, fmt.Errorf("parse image compatibility URL: %w", err)
	}

	proxyURL, err := http.ProxyFromEnvironment(&http.Request{URL: parsedURL})
	if err != nil {
		return nil, fmt.Errorf("resolve image compatibility proxy: %w", err)
	}
	if proxyURL == nil {
		return service.DoDownloadRequest(rawURL, reason...)
	}
	if err := validateImageURLCompatProxyTarget(parsedURL); err != nil {
		return nil, fmt.Errorf("request reject: %w", err)
	}

	baseClient := service.GetHttpClient()
	client := &http.Client{
		Transport: baseClient.Transport,
		Timeout:   baseClient.Timeout,
		Jar:       baseClient.Jar,
		CheckRedirect: func(request *http.Request, via []*http.Request) error {
			if len(via) >= 10 {
				return fmt.Errorf("stopped after 10 redirects")
			}
			redirectProxy, err := http.ProxyFromEnvironment(request)
			if err != nil {
				return fmt.Errorf("resolve redirect proxy: %w", err)
			}
			if redirectProxy == nil {
				return fmt.Errorf("image compatibility redirect left the configured proxy")
			}
			return validateImageURLCompatProxyTarget(request.URL)
		},
	}

	common.SysLog(fmt.Sprintf(
		"downloading image compatibility URL through configured proxy: %s, reason: %s",
		common.MaskSensitiveInfo(rawURL),
		strings.Join(reason, ", "),
	))
	return retryImageURLCompatDownload(
		func() (*http.Response, error) { return client.Get(rawURL) },
		time.Sleep,
	)
}

func retryImageURLCompatDownload(request func() (*http.Response, error), wait func(time.Duration)) (*http.Response, error) {
	var lastErr error
	for attempt := 1; attempt <= maxImageURLCompatDownloadAttempts; attempt++ {
		response, err := request()
		if err == nil && response != nil && response.StatusCode < http.StatusInternalServerError {
			return response, nil
		}
		if response != nil && response.Body != nil {
			response.Body.Close()
		}
		if err != nil {
			lastErr = err
		} else if response == nil {
			lastErr = fmt.Errorf("image compatibility download returned an empty response")
		} else {
			lastErr = fmt.Errorf("image compatibility download returned HTTP %d", response.StatusCode)
		}
		if attempt < maxImageURLCompatDownloadAttempts {
			wait(time.Duration(attempt) * 300 * time.Millisecond)
		}
	}
	return nil, lastErr
}

// Clash Fake-IP addresses are local DNS placeholders. When an HTTP proxy is
// selected, validate the hostname and port without resolving that placeholder.
func validateImageURLCompatProxyTarget(parsedURL *url.URL) error {
	if parsedURL == nil || parsedURL.Hostname() == "" {
		return fmt.Errorf("invalid image compatibility URL")
	}
	if parsedURL.Scheme != "http" && parsedURL.Scheme != "https" {
		return fmt.Errorf("unsupported image compatibility URL scheme: %s", parsedURL.Scheme)
	}
	if parsedURL.User != nil {
		return fmt.Errorf("image compatibility URL must not contain credentials")
	}

	port := 80
	if parsedURL.Scheme == "https" {
		port = 443
	}
	if parsedURL.Port() != "" {
		parsedPort, err := strconv.Atoi(parsedURL.Port())
		if err != nil {
			return fmt.Errorf("invalid image compatibility port: %w", err)
		}
		port = parsedPort
	}

	fetchSetting := system_setting.GetFetchSetting()
	if fetchSetting == nil || !fetchSetting.EnableSSRFProtection {
		return nil
	}
	protection, err := common.NewSSRFProtectionFromFetchSetting(
		fetchSetting.AllowPrivateIp,
		fetchSetting.DomainFilterMode,
		fetchSetting.IpFilterMode,
		fetchSetting.DomainList,
		fetchSetting.IpList,
		fetchSetting.AllowedPorts,
		fetchSetting.ApplyIPFilterForDomain,
	)
	if err != nil {
		return err
	}
	return protection.ValidateNetworkTarget(parsedURL.Hostname(), port)
}

func (s imageURLCompatStats) ServerTiming() string {
	return fmt.Sprintf(
		"image-url-download;dur=%.3f, image-b64-encode;dur=%.3f, image-url-compat;dur=%.3f",
		float64(s.DownloadTime.Microseconds())/1000,
		float64(s.EncodeTime.Microseconds())/1000,
		float64(s.TotalTime.Microseconds())/1000,
	)
}

func (s imageURLCompatStats) LogMessage() string {
	return fmt.Sprintf(
		"image URL compatibility converted=%d bytes=%d download_ms=%.3f encode_ms=%.3f total_ms=%.3f",
		s.Converted,
		s.DownloadedBytes,
		float64(s.DownloadTime.Microseconds())/1000,
		float64(s.EncodeTime.Microseconds())/1000,
		float64(s.TotalTime.Microseconds())/1000,
	)
}

func convertOpenAIImageURLsToBase64(responseBody []byte, download imageURLCompatDownloadFunc) ([]byte, imageURLCompatStats, error) {
	var stats imageURLCompatStats
	if !common.GetEnvOrDefaultBool("CODEX_IMAGE_URL_TO_B64", false) {
		return responseBody, stats, nil
	}

	startedAt := time.Now()
	var payload map[string]json.RawMessage
	if err := common.Unmarshal(responseBody, &payload); err != nil {
		return nil, stats, fmt.Errorf("decode image compatibility response: %w", err)
	}
	dataJSON, ok := payload["data"]
	if !ok {
		return responseBody, stats, nil
	}

	var images []map[string]json.RawMessage
	if err := common.Unmarshal(dataJSON, &images); err != nil {
		return nil, stats, fmt.Errorf("decode image compatibility data: %w", err)
	}

	for _, image := range images {
		if jsonString(image["b64_json"]) != "" {
			continue
		}

		var imageURL string
		var sourceField string
		for _, field := range []string{"url", "result_url", "image_url"} {
			if value := jsonString(image[field]); value != "" {
				imageURL = value
				sourceField = field
				break
			}
		}
		if imageURL == "" {
			continue
		}

		downloadStartedAt := time.Now()
		imageResponse, err := download(imageURL, "OpenAI image URL compatibility")
		if err != nil {
			return nil, stats, fmt.Errorf("download image compatibility URL: %w", err)
		}
		imageBytes, readErr := readImageURLCompatResponse(imageResponse)
		stats.DownloadTime += time.Since(downloadStartedAt)
		if readErr != nil {
			return nil, stats, readErr
		}

		encodeStartedAt := time.Now()
		encodedJSON, err := common.Marshal(base64.StdEncoding.EncodeToString(imageBytes))
		if err != nil {
			return nil, stats, fmt.Errorf("encode image compatibility response: %w", err)
		}
		image["b64_json"] = encodedJSON
		delete(image, sourceField)
		stats.EncodeTime += time.Since(encodeStartedAt)
		stats.Converted++
		stats.DownloadedBytes += int64(len(imageBytes))
	}

	if stats.Converted == 0 {
		return responseBody, stats, nil
	}

	encodeStartedAt := time.Now()
	dataJSON, err := common.Marshal(images)
	if err != nil {
		return nil, stats, fmt.Errorf("encode image compatibility data: %w", err)
	}
	payload["data"] = dataJSON
	convertedBody, err := common.Marshal(payload)
	stats.EncodeTime += time.Since(encodeStartedAt)
	stats.TotalTime = time.Since(startedAt)
	if err != nil {
		return nil, stats, fmt.Errorf("encode image compatibility payload: %w", err)
	}
	return convertedBody, stats, nil
}

func jsonString(raw json.RawMessage) string {
	if len(raw) == 0 {
		return ""
	}
	var value string
	if err := common.Unmarshal(raw, &value); err != nil {
		return ""
	}
	return strings.TrimSpace(value)
}

func readImageURLCompatResponse(response *http.Response) ([]byte, error) {
	if response == nil || response.Body == nil {
		return nil, fmt.Errorf("image compatibility download returned an empty response")
	}
	defer response.Body.Close()
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return nil, fmt.Errorf("image compatibility download returned HTTP %d", response.StatusCode)
	}
	if response.ContentLength > maxImageURLCompatBytes {
		return nil, fmt.Errorf("image compatibility download exceeds %d bytes", maxImageURLCompatBytes)
	}

	imageBytes, err := io.ReadAll(io.LimitReader(response.Body, maxImageURLCompatBytes+1))
	if err != nil {
		return nil, fmt.Errorf("read image compatibility download: %w", err)
	}
	if int64(len(imageBytes)) > maxImageURLCompatBytes {
		return nil, fmt.Errorf("image compatibility download exceeds %d bytes", maxImageURLCompatBytes)
	}
	if len(imageBytes) == 0 {
		return nil, fmt.Errorf("image compatibility download returned an empty body")
	}
	return imageBytes, nil
}