fix: improve error handling for image URL fetching and validation

This commit is contained in:
Laisky.Cai
2025-04-01 10:14:46 +00:00
parent ecadb19791
commit 747e84851e

View File

@@ -3,7 +3,6 @@ package image
import ( import (
"bytes" "bytes"
"encoding/base64" "encoding/base64"
"github.com/songquanpeng/one-api/common/client"
"image" "image"
_ "image/gif" _ "image/gif"
_ "image/jpeg" _ "image/jpeg"
@@ -13,6 +12,9 @@ import (
"strings" "strings"
"sync" "sync"
"github.com/pkg/errors"
"github.com/songquanpeng/one-api/common/client"
_ "golang.org/x/image/webp" _ "golang.org/x/image/webp"
) )
@@ -22,27 +24,36 @@ var dataURLPattern = regexp.MustCompile(`data:image/([^;]+);base64,(.*)`)
func IsImageUrl(url string) (bool, error) { func IsImageUrl(url string) (bool, error) {
resp, err := client.UserContentRequestHTTPClient.Head(url) resp, err := client.UserContentRequestHTTPClient.Head(url)
if err != nil { if err != nil {
return false, err return false, errors.Wrapf(err, "failed to fetch image URL: %s", url)
} }
if !strings.HasPrefix(resp.Header.Get("Content-Type"), "image/") {
return false, nil if resp.StatusCode != http.StatusOK {
return false, errors.Errorf("failed to fetch image URL: %s, status code: %d", url, resp.StatusCode)
} }
if resp.ContentLength > 10*1024*1024 {
return false, errors.Errorf("image size should not exceed 10MB: %s, size: %d", url, resp.ContentLength)
}
return true, nil return true, nil
} }
func GetImageSizeFromUrl(url string) (width int, height int, err error) { func GetImageSizeFromUrl(url string) (width int, height int, err error) {
isImage, err := IsImageUrl(url) isImage, err := IsImageUrl(url)
if err != nil {
return 0, 0, errors.Wrap(err, "failed to fetch image URL")
}
if !isImage { if !isImage {
return return 0, 0, errors.New("not an image URL")
} }
resp, err := client.UserContentRequestHTTPClient.Get(url) resp, err := client.UserContentRequestHTTPClient.Get(url)
if err != nil { if err != nil {
return return 0, 0, errors.Wrap(err, "failed to get image from URL")
} }
defer resp.Body.Close() defer resp.Body.Close()
img, _, err := image.DecodeConfig(resp.Body) img, _, err := image.DecodeConfig(resp.Body)
if err != nil { if err != nil {
return return 0, 0, errors.Wrap(err, "failed to decode image")
} }
return img.Width, img.Height, nil return img.Width, img.Height, nil
} }
@@ -58,22 +69,26 @@ func GetImageFromUrl(url string) (mimeType string, data string, err error) {
} }
isImage, err := IsImageUrl(url) isImage, err := IsImageUrl(url)
if !isImage { if err != nil {
return return mimeType, data, errors.Wrap(err, "failed to fetch image URL")
} }
if !isImage {
return mimeType, data, errors.New("not an image URL")
}
resp, err := http.Get(url) resp, err := http.Get(url)
if err != nil { if err != nil {
return return mimeType, data, errors.Wrap(err, "failed to get image from URL")
} }
defer resp.Body.Close() defer resp.Body.Close()
buffer := bytes.NewBuffer(nil) buffer := bytes.NewBuffer(nil)
_, err = buffer.ReadFrom(resp.Body) _, err = buffer.ReadFrom(resp.Body)
if err != nil { if err != nil {
return return mimeType, data, errors.Wrap(err, "failed to read image data from response")
} }
mimeType = resp.Header.Get("Content-Type") mimeType = resp.Header.Get("Content-Type")
data = base64.StdEncoding.EncodeToString(buffer.Bytes()) data = base64.StdEncoding.EncodeToString(buffer.Bytes())
return return mimeType, data, nil
} }
var ( var (