Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,660 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultBatchImageZipMaxItems = 200
|
||||
defaultBatchImageZipMaxBytes = 512 * 1024 * 1024
|
||||
defaultBatchImageDownloadDuration = 10 * time.Minute
|
||||
defaultBatchImageDownloadConcurrency = 1
|
||||
batchImageDownloadScannerMaxLineBytes = 16 * 1024 * 1024
|
||||
)
|
||||
|
||||
var errBatchImageDownloadSizeExceeded = errors.New("batch image download size limit exceeded")
|
||||
|
||||
type BatchImageDownloadLimiter interface {
|
||||
Acquire(ctx context.Context, userID string, kind string) (BatchImageDownloadPermit, error)
|
||||
}
|
||||
|
||||
type BatchImageDownloadPermit interface {
|
||||
Release(ctx context.Context) error
|
||||
}
|
||||
|
||||
type BatchImageContentStream struct {
|
||||
Reader io.ReadCloser
|
||||
ContentType string
|
||||
Filename string
|
||||
ContentLength *int64
|
||||
}
|
||||
|
||||
type BatchImageZipOptions struct {
|
||||
Status string
|
||||
MaxItems int
|
||||
IncludeManifest bool
|
||||
}
|
||||
|
||||
type BatchImageZipResult struct {
|
||||
FileCount int
|
||||
ErrorCount int
|
||||
}
|
||||
|
||||
type BatchImageLineImages struct {
|
||||
CustomID string
|
||||
Images []BatchImageInlineImage
|
||||
ErrorCode string
|
||||
ErrorMessage string
|
||||
}
|
||||
|
||||
type BatchImageInlineImage struct {
|
||||
MimeType string
|
||||
Extension string
|
||||
Base64Data string
|
||||
}
|
||||
|
||||
type BatchImageDownloadService struct {
|
||||
Repo BatchImageRepository
|
||||
ProviderRegistry *BatchImageProviderRegistry
|
||||
AccountResolver BatchImageAccountResolver
|
||||
Limiter BatchImageDownloadLimiter
|
||||
Config *config.Config
|
||||
}
|
||||
|
||||
type batchImageDownloadLimitWriter struct {
|
||||
w io.Writer
|
||||
limit int64
|
||||
written int64
|
||||
}
|
||||
|
||||
func (w *batchImageDownloadLimitWriter) Write(p []byte) (int, error) {
|
||||
if w == nil || w.w == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
if w.limit > 0 && w.written+int64(len(p)) > w.limit {
|
||||
return 0, errBatchImageDownloadSizeExceeded
|
||||
}
|
||||
n, err := w.w.Write(p)
|
||||
w.written += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func NewBatchImageDownloadService(repo BatchImageRepository, accountRepo AccountRepository, limiter BatchImageDownloadLimiter, cfg *config.Config) *BatchImageDownloadService {
|
||||
return &BatchImageDownloadService{
|
||||
Repo: repo,
|
||||
ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
|
||||
AccountResolver: &BatchImageAccountRepositoryResolver{Repo: accountRepo},
|
||||
Limiter: limiter,
|
||||
Config: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) OpenItemContent(ctx context.Context, owner BatchImageOwner, batchID string, customID string, imageIndex int) (*BatchImageContentStream, error) {
|
||||
if imageIndex < 0 {
|
||||
return nil, ErrBatchImageItemImageIndexOutOfRange
|
||||
}
|
||||
job, err := s.getCompletedJob(ctx, owner, batchID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
item, err := s.Repo.GetBatchImageItemForDownload(ctx, job.BatchID, customID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if item.Status != BatchImageItemStatusSuccess {
|
||||
return nil, ErrBatchImageItemFailed
|
||||
}
|
||||
if imageIndex >= item.ImageCount {
|
||||
return nil, ErrBatchImageItemImageIndexOutOfRange
|
||||
}
|
||||
|
||||
permit, err := s.acquirePermit(ctx, owner.UserID, "item")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
releasePermit := true
|
||||
defer func() {
|
||||
if releasePermit && permit != nil {
|
||||
_ = permit.Release(ctx)
|
||||
}
|
||||
}()
|
||||
|
||||
provider, account, err := s.providerAndAccount(ctx, job)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r, _, err := provider.OpenResult(ctx, job, account)
|
||||
if err != nil {
|
||||
return nil, ErrBatchImageResultMissing.WithCause(err)
|
||||
}
|
||||
defer func() { _ = r.Close() }()
|
||||
|
||||
line, err := findBatchImageLineImages(r, item.CustomID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if imageIndex >= len(line.Images) {
|
||||
return nil, ErrBatchImageItemImageIndexOutOfRange
|
||||
}
|
||||
image := line.Images[imageIndex]
|
||||
if strings.TrimSpace(image.Base64Data) == "" {
|
||||
return nil, ErrBatchImageResultMissing
|
||||
}
|
||||
contentType := strings.TrimSpace(image.MimeType)
|
||||
if contentType == "" {
|
||||
contentType = "application/octet-stream"
|
||||
}
|
||||
extension := strings.TrimSpace(image.Extension)
|
||||
if extension == "" {
|
||||
extension = batchImageFileExtension(contentType)
|
||||
}
|
||||
if extension == "" {
|
||||
extension = "bin"
|
||||
}
|
||||
|
||||
reader := base64.NewDecoder(base64.StdEncoding, strings.NewReader(image.Base64Data))
|
||||
releasePermit = false
|
||||
return &BatchImageContentStream{
|
||||
Reader: &batchImagePermitReadCloser{Reader: reader, permit: permit},
|
||||
ContentType: contentType,
|
||||
Filename: BatchImageSafeDownloadFilename(item.CustomID, extension),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) StreamZip(ctx context.Context, owner BatchImageOwner, batchID string, opts BatchImageZipOptions, w io.Writer) (*BatchImageZipResult, error) {
|
||||
job, err := s.getCompletedJob(ctx, owner, batchID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
maxItems := opts.MaxItems
|
||||
if cap := s.maxZipItems(); maxItems <= 0 || maxItems > cap {
|
||||
// 客户端传入的 max_items 不得放大管理员配置的 ZIP 上限。
|
||||
maxItems = cap
|
||||
}
|
||||
if job.SuccessCount > maxItems {
|
||||
return nil, ErrBatchImageZipTooManyItems
|
||||
}
|
||||
successItems, err := s.Repo.ListBatchImageItemsForDownload(ctx, job.BatchID, BatchImageItemStatusSuccess, maxItems+1)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(successItems) > maxItems {
|
||||
return nil, ErrBatchImageZipTooManyItems
|
||||
}
|
||||
failedItems, err := s.Repo.ListBatchImageItemsForDownload(ctx, job.BatchID, BatchImageItemStatusFailed, maxItems)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
permit, err := s.acquirePermit(ctx, owner.UserID, "zip")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if permit != nil {
|
||||
defer func() { _ = permit.Release(ctx) }()
|
||||
}
|
||||
|
||||
provider, account, err := s.providerAndAccount(ctx, job)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r, _, err := provider.OpenResult(ctx, job, account)
|
||||
if err != nil {
|
||||
return nil, ErrBatchImageResultMissing.WithCause(err)
|
||||
}
|
||||
defer func() { _ = r.Close() }()
|
||||
|
||||
streamCtx := ctx
|
||||
cancel := func() {}
|
||||
if d := s.maxDownloadDuration(); d > 0 {
|
||||
streamCtx, cancel = context.WithTimeout(ctx, d)
|
||||
}
|
||||
defer cancel()
|
||||
|
||||
limitedWriter := &batchImageDownloadLimitWriter{w: w, limit: s.maxDownloadBytes()}
|
||||
zipWriter := zip.NewWriter(limitedWriter)
|
||||
result, manifestFiles, zipErrors, err := s.writeZipImages(streamCtx, zipWriter, r, successItems)
|
||||
if err != nil {
|
||||
_ = zipWriter.Close()
|
||||
if errors.Is(err, errBatchImageDownloadSizeExceeded) {
|
||||
return result, ErrBatchImageDownloadTooLarge.WithCause(err)
|
||||
}
|
||||
return result, ErrBatchImageDownloadFailed.WithCause(err)
|
||||
}
|
||||
zipErrors = append(zipErrors, batchImageZipErrorsFromItems(failedItems)...)
|
||||
if err := writeBatchImageZipJSON(zipWriter, "manifest.json", batchImageZipManifest{
|
||||
BatchID: job.BatchID,
|
||||
Model: job.Model,
|
||||
ItemCount: job.ItemCount,
|
||||
SuccessCount: job.SuccessCount,
|
||||
FailCount: job.FailCount,
|
||||
Files: manifestFiles,
|
||||
}); err != nil {
|
||||
_ = zipWriter.Close()
|
||||
if errors.Is(err, errBatchImageDownloadSizeExceeded) {
|
||||
return result, ErrBatchImageDownloadTooLarge.WithCause(err)
|
||||
}
|
||||
return result, ErrBatchImageDownloadFailed.WithCause(err)
|
||||
}
|
||||
if err := writeBatchImageZipJSON(zipWriter, "errors.json", zipErrors); err != nil {
|
||||
_ = zipWriter.Close()
|
||||
if errors.Is(err, errBatchImageDownloadSizeExceeded) {
|
||||
return result, ErrBatchImageDownloadTooLarge.WithCause(err)
|
||||
}
|
||||
return result, ErrBatchImageDownloadFailed.WithCause(err)
|
||||
}
|
||||
result.ErrorCount = len(zipErrors)
|
||||
if err := zipWriter.Close(); err != nil {
|
||||
if errors.Is(err, errBatchImageDownloadSizeExceeded) {
|
||||
return result, ErrBatchImageDownloadTooLarge.WithCause(err)
|
||||
}
|
||||
return result, ErrBatchImageDownloadFailed.WithCause(err)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) writeZipImages(ctx context.Context, zipWriter *zip.Writer, resultReader io.Reader, successItems []*BatchImageItem) (*BatchImageZipResult, []batchImageZipManifestFile, []batchImageZipError, error) {
|
||||
successByID := make(map[string]*BatchImageItem, len(successItems))
|
||||
missing := make(map[string]struct{}, len(successItems))
|
||||
for _, item := range successItems {
|
||||
if item == nil {
|
||||
continue
|
||||
}
|
||||
successByID[item.CustomID] = item
|
||||
missing[item.CustomID] = struct{}{}
|
||||
}
|
||||
scanner := bufio.NewScanner(resultReader)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), batchImageDownloadScannerMaxLineBytes)
|
||||
|
||||
result := &BatchImageZipResult{}
|
||||
var manifestFiles []batchImageZipManifestFile
|
||||
var zipErrors []batchImageZipError
|
||||
for scanner.Scan() {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return result, manifestFiles, zipErrors, err
|
||||
}
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
images, err := ExtractBatchImagePartsFromResultLine([]byte(line))
|
||||
if err != nil {
|
||||
return result, manifestFiles, zipErrors, err
|
||||
}
|
||||
item := successByID[images.CustomID]
|
||||
if item == nil {
|
||||
continue
|
||||
}
|
||||
delete(missing, images.CustomID)
|
||||
if len(images.Images) == 0 {
|
||||
zipErrors = append(zipErrors, batchImageZipError{CustomID: images.CustomID, Code: "EMPTY_IMAGE_OUTPUT", Message: "provider response contained no image output"})
|
||||
continue
|
||||
}
|
||||
for idx, image := range images.Images {
|
||||
extension := image.Extension
|
||||
if extension == "" {
|
||||
extension = "bin"
|
||||
}
|
||||
filename := batchImageZipImageFilename(item.CustomID, idx, extension)
|
||||
entry, err := zipWriter.CreateHeader(&zip.FileHeader{Name: filename, Method: zip.Deflate})
|
||||
if err != nil {
|
||||
return result, manifestFiles, zipErrors, err
|
||||
}
|
||||
decoder := base64.NewDecoder(base64.StdEncoding, strings.NewReader(image.Base64Data))
|
||||
if _, err := io.Copy(entry, decoder); err != nil {
|
||||
zipErrors = append(zipErrors, batchImageZipError{CustomID: item.CustomID, Code: "IMAGE_DECODE_FAILED", Message: "image data could not be decoded"})
|
||||
continue
|
||||
}
|
||||
result.FileCount++
|
||||
manifestFiles = append(manifestFiles, batchImageZipManifestFile{
|
||||
CustomID: item.CustomID,
|
||||
Filename: filename,
|
||||
MimeType: image.MimeType,
|
||||
ImageIndex: idx,
|
||||
})
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return result, manifestFiles, zipErrors, err
|
||||
}
|
||||
missingIDs := make([]string, 0, len(missing))
|
||||
for customID := range missing {
|
||||
missingIDs = append(missingIDs, customID)
|
||||
}
|
||||
sort.Strings(missingIDs)
|
||||
for _, customID := range missingIDs {
|
||||
zipErrors = append(zipErrors, batchImageZipError{CustomID: customID, Code: "RESULT_MISSING", Message: "provider result was not found for item"})
|
||||
}
|
||||
return result, manifestFiles, zipErrors, nil
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) getCompletedJob(ctx context.Context, owner BatchImageOwner, batchID string) (*BatchImageJob, error) {
|
||||
if s == nil || s.Repo == nil {
|
||||
return nil, ErrBatchImageDownloadFailed
|
||||
}
|
||||
job, err := s.Repo.GetBatchImageJobForDownload(ctx, owner.UserID, owner.APIKeyID, batchID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
switch job.Status {
|
||||
case BatchImageJobStatusCompleted:
|
||||
return job, nil
|
||||
case BatchImageJobStatusOutputDeleted:
|
||||
return nil, ErrBatchImageOutputDeleted
|
||||
default:
|
||||
return nil, ErrBatchImageNotReady
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) providerAndAccount(ctx context.Context, job *BatchImageJob) (BatchImageProvider, *Account, error) {
|
||||
if s == nil || s.ProviderRegistry == nil || s.AccountResolver == nil || job == nil {
|
||||
return nil, nil, ErrBatchImageDownloadFailed
|
||||
}
|
||||
provider, ok := s.ProviderRegistry.Get(job.Provider)
|
||||
if !ok || provider == nil {
|
||||
return nil, nil, ErrBatchImageUnsupportedProvider
|
||||
}
|
||||
if job.AccountID == nil || *job.AccountID <= 0 {
|
||||
return nil, nil, ErrBatchImageMissingAccountID
|
||||
}
|
||||
account, err := s.AccountResolver.ResolveBatchImageAccount(ctx, *job.AccountID)
|
||||
if err != nil {
|
||||
return nil, nil, ErrBatchImageDownloadFailed
|
||||
}
|
||||
if !provider.SupportsAccount(account) {
|
||||
return nil, nil, ErrBatchImageProviderUnsupportedAccount
|
||||
}
|
||||
return provider, account, nil
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) acquirePermit(ctx context.Context, userID int64, kind string) (BatchImageDownloadPermit, error) {
|
||||
if s == nil || s.Limiter == nil {
|
||||
return nil, nil
|
||||
}
|
||||
permit, err := s.Limiter.Acquire(ctx, fmt.Sprintf("%d", userID), kind)
|
||||
if err != nil {
|
||||
if infraerrors.Code(err) == http.StatusTooManyRequests {
|
||||
return nil, ErrBatchImageDownloadLimited
|
||||
}
|
||||
return nil, ErrBatchImageDownloadLimited.WithCause(err)
|
||||
}
|
||||
return permit, nil
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) maxZipItems() int {
|
||||
if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadItemsZip > 0 {
|
||||
return s.Config.BatchImage.MaxDownloadItemsZip
|
||||
}
|
||||
return defaultBatchImageZipMaxItems
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) maxDownloadBytes() int64 {
|
||||
if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadBytesPerRequest > 0 {
|
||||
return s.Config.BatchImage.MaxDownloadBytesPerRequest
|
||||
}
|
||||
return defaultBatchImageZipMaxBytes
|
||||
}
|
||||
|
||||
func (s *BatchImageDownloadService) maxDownloadDuration() time.Duration {
|
||||
if s != nil && s.Config != nil && s.Config.BatchImage.MaxDownloadDurationSeconds > 0 {
|
||||
return time.Duration(s.Config.BatchImage.MaxDownloadDurationSeconds) * time.Second
|
||||
}
|
||||
return defaultBatchImageDownloadDuration
|
||||
}
|
||||
|
||||
func ExtractBatchImagePartsFromResultLine(line []byte) (*BatchImageLineImages, error) {
|
||||
var obj map[string]any
|
||||
if err := json.Unmarshal(line, &obj); err != nil {
|
||||
return nil, ErrBatchImageIndexParseFailed.WithCause(err)
|
||||
}
|
||||
customID := batchImageFirstNonEmptyString(
|
||||
batchImageMapString(obj, "key"),
|
||||
batchImageMapString(obj, "custom_id"),
|
||||
batchImageMapString(obj, "customId"),
|
||||
batchImageNestedString(obj, "request", "key"),
|
||||
)
|
||||
if customID == "" {
|
||||
return nil, ErrBatchImageIndexParseFailed.WithCause(fmt.Errorf("missing custom id"))
|
||||
}
|
||||
out := &BatchImageLineImages{CustomID: customID}
|
||||
out.Images = append(out.Images, extractBatchImageInlineImages(batchImageNestedAny(obj, "response", "candidates"))...)
|
||||
out.Images = append(out.Images, extractBatchImageInlineImages(obj["candidates"])...)
|
||||
if len(out.Images) > 0 {
|
||||
return out, nil
|
||||
}
|
||||
if code, message, ok := batchImageFailureFromProviderFields(obj); ok {
|
||||
out.ErrorCode = code
|
||||
out.ErrorMessage = truncateBatchImageMessage(message, batchImageMaxErrorMessageLength)
|
||||
return out, nil
|
||||
}
|
||||
if _, hasResponse := obj["response"]; hasResponse || batchImageHasCandidates(obj) {
|
||||
out.ErrorCode = "EMPTY_IMAGE_OUTPUT"
|
||||
out.ErrorMessage = "provider response contained no image output"
|
||||
return out, nil
|
||||
}
|
||||
out.ErrorCode = "PROVIDER_ITEM_FAILED"
|
||||
out.ErrorMessage = "provider result line contained no image output"
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func extractBatchImageInlineImages(raw any) []BatchImageInlineImage {
|
||||
candidates, ok := raw.([]any)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
var images []BatchImageInlineImage
|
||||
for _, candidateRaw := range candidates {
|
||||
candidate, ok := candidateRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
parts, ok := batchImageNestedAny(candidate, "content", "parts").([]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
for _, partRaw := range parts {
|
||||
part, ok := partRaw.(map[string]any)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
inline, ok := firstMap(part["inlineData"], part["inline_data"])
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(batchImageMapString(inline, "data"))
|
||||
mime := strings.TrimSpace(batchImageFirstNonEmptyString(batchImageMapString(inline, "mimeType"), batchImageMapString(inline, "mime_type")))
|
||||
if data == "" || !strings.HasPrefix(strings.ToLower(mime), "image/") {
|
||||
continue
|
||||
}
|
||||
images = append(images, BatchImageInlineImage{
|
||||
MimeType: mime,
|
||||
Extension: batchImageFileExtension(mime),
|
||||
Base64Data: data,
|
||||
})
|
||||
}
|
||||
}
|
||||
return images
|
||||
}
|
||||
|
||||
func findBatchImageLineImages(r io.Reader, customID string) (*BatchImageLineImages, error) {
|
||||
scanner := bufio.NewScanner(r)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), batchImageDownloadScannerMaxLineBytes)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
parsed, err := ExtractBatchImagePartsFromResultLine([]byte(line))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if parsed.CustomID != customID {
|
||||
continue
|
||||
}
|
||||
if len(parsed.Images) == 0 {
|
||||
if parsed.ErrorCode != "" {
|
||||
return nil, ErrBatchImageItemFailed
|
||||
}
|
||||
return nil, ErrBatchImageResultMissing
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return nil, ErrBatchImageDownloadFailed.WithCause(err)
|
||||
}
|
||||
return nil, ErrBatchImageResultMissing
|
||||
}
|
||||
|
||||
func BatchImageSafeDownloadFilename(customID, extension string) string {
|
||||
base := sanitizeBatchImageFilenameBase(customID)
|
||||
extension = sanitizeBatchImageFilenameExtension(extension)
|
||||
if extension == "" {
|
||||
extension = "bin"
|
||||
}
|
||||
return base + "." + extension
|
||||
}
|
||||
|
||||
func BatchImageContentDispositionAttachment(filename string) string {
|
||||
filename = strings.ReplaceAll(filename, "\\", "_")
|
||||
filename = strings.ReplaceAll(filename, `"`, "_")
|
||||
filename = sanitizeBatchImageFilenameBase(strings.TrimSuffix(filename, filepath.Ext(filename))) + filepath.Ext(filename)
|
||||
return `attachment; filename="` + filename + `"`
|
||||
}
|
||||
|
||||
func sanitizeBatchImageFilenameBase(value string) string {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return "image"
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, r := range value {
|
||||
switch {
|
||||
case r == '/' || r == '\\' || r == ':' || r == 0:
|
||||
_ = b.WriteByte('_')
|
||||
case unicode.IsControl(r):
|
||||
_ = b.WriteByte('_')
|
||||
case unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_' || r == '-' || r == '.':
|
||||
_, _ = b.WriteRune(r)
|
||||
default:
|
||||
_ = b.WriteByte('_')
|
||||
}
|
||||
}
|
||||
out := strings.Trim(b.String(), ". ")
|
||||
for strings.Contains(out, "..") {
|
||||
out = strings.ReplaceAll(out, "..", "_")
|
||||
}
|
||||
out = strings.Trim(out, ". ")
|
||||
if out == "" {
|
||||
out = "image"
|
||||
}
|
||||
if len(out) > 120 {
|
||||
out = strings.TrimRight(out[:120], ". ")
|
||||
}
|
||||
if out == "" {
|
||||
out = "image"
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sanitizeBatchImageFilenameExtension(extension string) string {
|
||||
extension = strings.TrimPrefix(strings.TrimSpace(strings.ToLower(extension)), ".")
|
||||
var b strings.Builder
|
||||
for _, r := range extension {
|
||||
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
||||
_, _ = b.WriteRune(r)
|
||||
}
|
||||
}
|
||||
out := b.String()
|
||||
if len(out) > 12 {
|
||||
out = out[:12]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func batchImageZipImageFilename(customID string, imageIndex int, extension string) string {
|
||||
base := sanitizeBatchImageFilenameBase(customID)
|
||||
if imageIndex > 0 {
|
||||
base = fmt.Sprintf("%s_%d", base, imageIndex+1)
|
||||
}
|
||||
return "images/" + BatchImageSafeDownloadFilename(base, extension)
|
||||
}
|
||||
|
||||
func writeBatchImageZipJSON(zipWriter *zip.Writer, name string, value any) error {
|
||||
entry, err := zipWriter.CreateHeader(&zip.FileHeader{Name: name, Method: zip.Deflate})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
encoder := json.NewEncoder(entry)
|
||||
encoder.SetIndent("", " ")
|
||||
return encoder.Encode(value)
|
||||
}
|
||||
|
||||
type batchImageZipManifest struct {
|
||||
BatchID string `json:"batch_id"`
|
||||
Model string `json:"model"`
|
||||
ItemCount int `json:"item_count"`
|
||||
SuccessCount int `json:"success_count"`
|
||||
FailCount int `json:"fail_count"`
|
||||
Files []batchImageZipManifestFile `json:"files"`
|
||||
}
|
||||
|
||||
type batchImageZipManifestFile struct {
|
||||
CustomID string `json:"custom_id"`
|
||||
Filename string `json:"filename"`
|
||||
MimeType string `json:"mime_type"`
|
||||
ImageIndex int `json:"image_index"`
|
||||
}
|
||||
|
||||
type batchImageZipError struct {
|
||||
CustomID string `json:"custom_id"`
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
func batchImageZipErrorsFromItems(items []*BatchImageItem) []batchImageZipError {
|
||||
out := make([]batchImageZipError, 0, len(items))
|
||||
for _, item := range items {
|
||||
if item == nil {
|
||||
continue
|
||||
}
|
||||
out = append(out, batchImageZipError{
|
||||
CustomID: item.CustomID,
|
||||
Code: batchImageDerefString(item.ErrorCode),
|
||||
Message: sanitizeBatchImagePublicMessage(batchImageDerefString(item.ErrorMessage)),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type batchImagePermitReadCloser struct {
|
||||
io.Reader
|
||||
permit BatchImageDownloadPermit
|
||||
once sync.Once
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *batchImagePermitReadCloser) Close() error {
|
||||
r.once.Do(func() {
|
||||
if r.permit != nil {
|
||||
r.err = r.permit.Release(context.Background())
|
||||
}
|
||||
})
|
||||
return r.err
|
||||
}
|
||||
Reference in New Issue
Block a user