refactor: 彻底移除 Sora 视频生成模块（全栈清理）
## 后端变更 - 删除 21 个 sora_*.go 服务文件（service/handler/repository/routes） - 删除 Sora 相关 migration 文件（046/047/063/090） - 清理 config 中的 sora_* 配置项和平台常量 - 清理 wire 依赖注入中的 Sora 组件 - 修复 wire_gen.go 语法错误（缺少逗号和闭合括号） - 移除 go.mod 中的 go-sora2api 依赖 - 更新 ent schema usage_log.go 注释 ## 前端变更 - 删除 SoraView、SoraAdminView 及 8 个 Sora 子组件 - 删除 sora API 层和路由配置 - 清理 UserEditModal 中的 Sora 存储配额 UI - 清理 types/index.ts 中 Sora 相关类型定义 - 清理 stores/app.ts 默认配置 - 清理 i18n 翻译文件 en.ts/zh.ts (~110 行) - 更新相关测试文件 ## 文档更新 - README.md / README_CN.md / README_JA.md: 移除 Sora 状态说明和配置段落 - PROJECT_DIFF.md: 移除 Sora 相关差异描述 ## 验证结果 - ✅ Go 编译通过 (go build ./...) - ✅ TypeScript 类型检查通过 (vue-tsc --noEmit) - ✅ 后端测试全通过 (0 failures) - ✅ 前端测试全通过 (59 files, 329 tests, 0 failures) - ✅ 前端生产构建成功 (23.81s)
This commit is contained in:
@@ -24,7 +24,6 @@ const (
|
||||
PlatformOpenAI = domain.PlatformOpenAI
|
||||
PlatformGemini = domain.PlatformGemini
|
||||
PlatformAntigravity = domain.PlatformAntigravity
|
||||
PlatformSora = domain.PlatformSora // 从本地版本合并
|
||||
)
|
||||
|
||||
// Account type constants
|
||||
@@ -253,19 +252,6 @@ const (
|
||||
SettingKeyEnableMetadataPassthrough = "enable_metadata_passthrough"
|
||||
// SettingKeyEnableCCHSigning 是否对 billing header 中的 cch 进行 xxHash64 签名（默认 false）
|
||||
SettingKeyEnableCCHSigning = "enable_cch_signing"
|
||||
|
||||
// Sora S3 存储配置 (从本地版本合并)
|
||||
SettingKeySoraS3Enabled = "sora_s3_enabled" // 是否启用 Sora S3 存储
|
||||
SettingKeySoraS3Endpoint = "sora_s3_endpoint" // S3 端点地址
|
||||
SettingKeySoraS3Region = "sora_s3_region" // S3 区域
|
||||
SettingKeySoraS3Bucket = "sora_s3_bucket" // S3 存储桶名称
|
||||
SettingKeySoraS3AccessKeyID = "sora_s3_access_key_id" // S3 Access Key ID
|
||||
SettingKeySoraS3SecretAccessKey = "sora_s3_secret_access_key" // S3 Secret Access Key（加密存储）
|
||||
SettingKeySoraS3Prefix = "sora_s3_prefix" // S3 对象键前缀
|
||||
SettingKeySoraS3ForcePathStyle = "sora_s3_force_path_style" // 是否强制 Path Style（兼容 MinIO 等）
|
||||
SettingKeySoraS3CDNURL = "sora_s3_cdn_url" // CDN 加速 URL（可选）
|
||||
SettingKeySoraS3Profiles = "sora_s3_profiles" // Sora S3 多配置（JSON）
|
||||
SettingKeySoraDefaultStorageQuotaBytes = "sora_default_storage_quota_bytes" // Sora 默认存储配额（字节）
|
||||
)
|
||||
|
||||
// AdminAPIKeyPrefix is the prefix for admin API keys (distinct from user "sk-" keys).
|
||||
|
||||
@@ -503,10 +503,6 @@ type ForwardResult struct {
|
||||
// 图片生成计费字段（图片生成模型使用）
|
||||
ImageCount int // 生成的图片数量
|
||||
ImageSize string // 图片尺寸 "1K", "2K", "4K"
|
||||
|
||||
// Sora 媒体字段 (从本地版本合并)
|
||||
MediaType string // image / video / prompt
|
||||
MediaURL string // 生成后的媒体地址（可选）
|
||||
}
|
||||
|
||||
// UpstreamFailoverError indicates an upstream error that should trigger account failover.
|
||||
|
||||
@@ -24,8 +24,6 @@ import (
|
||||
var (
|
||||
ErrRegistrationDisabled = infraerrors.Forbidden("REGISTRATION_DISABLED", "registration is currently disabled")
|
||||
ErrSettingNotFound = infraerrors.NotFound("SETTING_NOT_FOUND", "setting not found")
|
||||
ErrSoraS3ProfileNotFound = infraerrors.NotFound("SORA_S3_PROFILE_NOT_FOUND", "sora s3 profile not found") // 从本地版本合并
|
||||
ErrSoraS3ProfileExists = infraerrors.Conflict("SORA_S3_PROFILE_EXISTS", "sora s3 profile already exists") // 从本地版本合并
|
||||
ErrDefaultSubGroupInvalid = infraerrors.BadRequest(
|
||||
"DEFAULT_SUBSCRIPTION_GROUP_INVALID",
|
||||
"default subscription group must exist and be subscription type",
|
||||
@@ -2128,315 +2126,3 @@ func (s *SettingService) SetStreamTimeoutSettings(ctx context.Context, settings
|
||||
|
||||
return s.settingRepo.Set(ctx, SettingKeyStreamTimeoutSettings, string(data))
|
||||
}
|
||||
|
||||
// Sora S3 存储配置 (从本地版本合并)
|
||||
type soraS3ProfilesStore struct {
|
||||
ActiveProfileID string `json:"active_profile_id"`
|
||||
Items []soraS3ProfileStoreItem `json:"items"`
|
||||
}
|
||||
|
||||
type soraS3ProfileStoreItem struct {
|
||||
ProfileID string `json:"profile_id"`
|
||||
Name string `json:"name"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
Region string `json:"region"`
|
||||
Bucket string `json:"bucket"`
|
||||
AccessKeyID string `json:"access_key_id"`
|
||||
SecretAccessKey string `json:"secret_access_key"`
|
||||
Prefix string `json:"prefix"`
|
||||
ForcePathStyle bool `json:"force_path_style"`
|
||||
CDNURL string `json:"cdn_url"`
|
||||
DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
// GetSoraS3Settings 获取 Sora S3 存储配置（兼容旧单配置语义：返回当前激活配置）
|
||||
func (s *SettingService) GetSoraS3Settings(ctx context.Context) (*SoraS3Settings, error) {
|
||||
profiles, err := s.ListSoraS3Profiles(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
activeProfile := pickActiveSoraS3Profile(profiles.Items, profiles.ActiveProfileID)
|
||||
if activeProfile == nil {
|
||||
return &SoraS3Settings{}, nil
|
||||
}
|
||||
|
||||
return &SoraS3Settings{
|
||||
Enabled: activeProfile.Enabled,
|
||||
Endpoint: activeProfile.Endpoint,
|
||||
Region: activeProfile.Region,
|
||||
Bucket: activeProfile.Bucket,
|
||||
AccessKeyID: activeProfile.AccessKeyID,
|
||||
SecretAccessKey: activeProfile.SecretAccessKey,
|
||||
SecretAccessKeyConfigured: activeProfile.SecretAccessKeyConfigured,
|
||||
Prefix: activeProfile.Prefix,
|
||||
ForcePathStyle: activeProfile.ForcePathStyle,
|
||||
CDNURL: activeProfile.CDNURL,
|
||||
DefaultStorageQuotaBytes: activeProfile.DefaultStorageQuotaBytes,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ListSoraS3Profiles 获取 Sora S3 多配置列表
|
||||
func (s *SettingService) ListSoraS3Profiles(ctx context.Context) (*SoraS3ProfileList, error) {
|
||||
store, err := s.loadSoraS3ProfilesStore(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return convertSoraS3ProfilesStore(store), nil
|
||||
}
|
||||
|
||||
func (s *SettingService) loadSoraS3ProfilesStore(ctx context.Context) (*soraS3ProfilesStore, error) {
|
||||
raw, err := s.settingRepo.GetValue(ctx, SettingKeySoraS3Profiles)
|
||||
if err == nil {
|
||||
trimmed := strings.TrimSpace(raw)
|
||||
if trimmed == "" {
|
||||
return &soraS3ProfilesStore{}, nil
|
||||
}
|
||||
var store soraS3ProfilesStore
|
||||
if unmarshalErr := json.Unmarshal([]byte(trimmed), &store); unmarshalErr != nil {
|
||||
legacy, legacyErr := s.getLegacySoraS3Settings(ctx)
|
||||
if legacyErr != nil {
|
||||
return nil, fmt.Errorf("unmarshal sora s3 profiles: %w", unmarshalErr)
|
||||
}
|
||||
if isEmptyLegacySoraS3Settings(legacy) {
|
||||
return &soraS3ProfilesStore{}, nil
|
||||
}
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
return &soraS3ProfilesStore{
|
||||
ActiveProfileID: "default",
|
||||
Items: []soraS3ProfileStoreItem{
|
||||
{
|
||||
ProfileID: "default",
|
||||
Name: "Default",
|
||||
Enabled: legacy.Enabled,
|
||||
Endpoint: strings.TrimSpace(legacy.Endpoint),
|
||||
Region: strings.TrimSpace(legacy.Region),
|
||||
Bucket: strings.TrimSpace(legacy.Bucket),
|
||||
AccessKeyID: strings.TrimSpace(legacy.AccessKeyID),
|
||||
SecretAccessKey: legacy.SecretAccessKey,
|
||||
Prefix: strings.TrimSpace(legacy.Prefix),
|
||||
ForcePathStyle: legacy.ForcePathStyle,
|
||||
CDNURL: strings.TrimSpace(legacy.CDNURL),
|
||||
DefaultStorageQuotaBytes: maxInt64(legacy.DefaultStorageQuotaBytes, 0),
|
||||
UpdatedAt: now,
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
normalized := normalizeSoraS3ProfilesStore(store)
|
||||
return &normalized, nil
|
||||
}
|
||||
|
||||
if !errors.Is(err, ErrSettingNotFound) {
|
||||
return nil, fmt.Errorf("get sora s3 profiles: %w", err)
|
||||
}
|
||||
|
||||
legacy, legacyErr := s.getLegacySoraS3Settings(ctx)
|
||||
if legacyErr != nil {
|
||||
return nil, legacyErr
|
||||
}
|
||||
if isEmptyLegacySoraS3Settings(legacy) {
|
||||
return &soraS3ProfilesStore{}, nil
|
||||
}
|
||||
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
return &soraS3ProfilesStore{
|
||||
ActiveProfileID: "default",
|
||||
Items: []soraS3ProfileStoreItem{
|
||||
{
|
||||
ProfileID: "default",
|
||||
Name: "Default",
|
||||
Enabled: legacy.Enabled,
|
||||
Endpoint: strings.TrimSpace(legacy.Endpoint),
|
||||
Region: strings.TrimSpace(legacy.Region),
|
||||
Bucket: strings.TrimSpace(legacy.Bucket),
|
||||
AccessKeyID: strings.TrimSpace(legacy.AccessKeyID),
|
||||
SecretAccessKey: legacy.SecretAccessKey,
|
||||
Prefix: strings.TrimSpace(legacy.Prefix),
|
||||
ForcePathStyle: legacy.ForcePathStyle,
|
||||
CDNURL: strings.TrimSpace(legacy.CDNURL),
|
||||
DefaultStorageQuotaBytes: maxInt64(legacy.DefaultStorageQuotaBytes, 0),
|
||||
UpdatedAt: now,
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *SettingService) getLegacySoraS3Settings(ctx context.Context) (*SoraS3Settings, error) {
|
||||
keys := []string{
|
||||
SettingKeySoraS3Enabled,
|
||||
SettingKeySoraS3Endpoint,
|
||||
SettingKeySoraS3Region,
|
||||
SettingKeySoraS3Bucket,
|
||||
SettingKeySoraS3AccessKeyID,
|
||||
SettingKeySoraS3SecretAccessKey,
|
||||
SettingKeySoraS3Prefix,
|
||||
SettingKeySoraS3ForcePathStyle,
|
||||
SettingKeySoraS3CDNURL,
|
||||
SettingKeySoraDefaultStorageQuotaBytes,
|
||||
}
|
||||
|
||||
values, err := s.settingRepo.GetMultiple(ctx, keys)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get legacy sora s3 settings: %w", err)
|
||||
}
|
||||
|
||||
result := &SoraS3Settings{
|
||||
Enabled: values[SettingKeySoraS3Enabled] == "true",
|
||||
Endpoint: values[SettingKeySoraS3Endpoint],
|
||||
Region: values[SettingKeySoraS3Region],
|
||||
Bucket: values[SettingKeySoraS3Bucket],
|
||||
AccessKeyID: values[SettingKeySoraS3AccessKeyID],
|
||||
SecretAccessKey: values[SettingKeySoraS3SecretAccessKey],
|
||||
SecretAccessKeyConfigured: values[SettingKeySoraS3SecretAccessKey] != "",
|
||||
Prefix: values[SettingKeySoraS3Prefix],
|
||||
ForcePathStyle: values[SettingKeySoraS3ForcePathStyle] == "true",
|
||||
CDNURL: values[SettingKeySoraS3CDNURL],
|
||||
}
|
||||
if v, parseErr := strconv.ParseInt(values[SettingKeySoraDefaultStorageQuotaBytes], 10, 64); parseErr == nil {
|
||||
result.DefaultStorageQuotaBytes = v
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func normalizeSoraS3ProfilesStore(store soraS3ProfilesStore) soraS3ProfilesStore {
|
||||
seen := make(map[string]struct{}, len(store.Items))
|
||||
normalized := soraS3ProfilesStore{
|
||||
ActiveProfileID: strings.TrimSpace(store.ActiveProfileID),
|
||||
Items: make([]soraS3ProfileStoreItem, 0, len(store.Items)),
|
||||
}
|
||||
now := time.Now().UTC().Format(time.RFC3339)
|
||||
|
||||
for idx := range store.Items {
|
||||
item := store.Items[idx]
|
||||
item.ProfileID = strings.TrimSpace(item.ProfileID)
|
||||
if item.ProfileID == "" {
|
||||
item.ProfileID = fmt.Sprintf("profile-%d", idx+1)
|
||||
}
|
||||
if _, exists := seen[item.ProfileID]; exists {
|
||||
continue
|
||||
}
|
||||
seen[item.ProfileID] = struct{}{}
|
||||
|
||||
item.Name = strings.TrimSpace(item.Name)
|
||||
if item.Name == "" {
|
||||
item.Name = item.ProfileID
|
||||
}
|
||||
item.Endpoint = strings.TrimSpace(item.Endpoint)
|
||||
item.Region = strings.TrimSpace(item.Region)
|
||||
item.Bucket = strings.TrimSpace(item.Bucket)
|
||||
item.AccessKeyID = strings.TrimSpace(item.AccessKeyID)
|
||||
item.Prefix = strings.TrimSpace(item.Prefix)
|
||||
item.CDNURL = strings.TrimSpace(item.CDNURL)
|
||||
item.DefaultStorageQuotaBytes = maxInt64(item.DefaultStorageQuotaBytes, 0)
|
||||
item.UpdatedAt = strings.TrimSpace(item.UpdatedAt)
|
||||
if item.UpdatedAt == "" {
|
||||
item.UpdatedAt = now
|
||||
}
|
||||
normalized.Items = append(normalized.Items, item)
|
||||
}
|
||||
|
||||
if len(normalized.Items) == 0 {
|
||||
normalized.ActiveProfileID = ""
|
||||
return normalized
|
||||
}
|
||||
|
||||
if findSoraS3ProfileIndex(normalized.Items, normalized.ActiveProfileID) >= 0 {
|
||||
return normalized
|
||||
}
|
||||
|
||||
normalized.ActiveProfileID = normalized.Items[0].ProfileID
|
||||
return normalized
|
||||
}
|
||||
|
||||
func convertSoraS3ProfilesStore(store *soraS3ProfilesStore) *SoraS3ProfileList {
|
||||
if store == nil {
|
||||
return &SoraS3ProfileList{}
|
||||
}
|
||||
items := make([]SoraS3Profile, 0, len(store.Items))
|
||||
for idx := range store.Items {
|
||||
item := store.Items[idx]
|
||||
items = append(items, SoraS3Profile{
|
||||
ProfileID: item.ProfileID,
|
||||
Name: item.Name,
|
||||
IsActive: item.ProfileID == store.ActiveProfileID,
|
||||
Enabled: item.Enabled,
|
||||
Endpoint: item.Endpoint,
|
||||
Region: item.Region,
|
||||
Bucket: item.Bucket,
|
||||
AccessKeyID: item.AccessKeyID,
|
||||
SecretAccessKey: item.SecretAccessKey,
|
||||
SecretAccessKeyConfigured: item.SecretAccessKey != "",
|
||||
Prefix: item.Prefix,
|
||||
ForcePathStyle: item.ForcePathStyle,
|
||||
CDNURL: item.CDNURL,
|
||||
DefaultStorageQuotaBytes: item.DefaultStorageQuotaBytes,
|
||||
UpdatedAt: item.UpdatedAt,
|
||||
})
|
||||
}
|
||||
return &SoraS3ProfileList{
|
||||
ActiveProfileID: store.ActiveProfileID,
|
||||
Items: items,
|
||||
}
|
||||
}
|
||||
|
||||
func pickActiveSoraS3Profile(items []SoraS3Profile, activeProfileID string) *SoraS3Profile {
|
||||
for idx := range items {
|
||||
if items[idx].ProfileID == activeProfileID {
|
||||
return &items[idx]
|
||||
}
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return nil
|
||||
}
|
||||
return &items[0]
|
||||
}
|
||||
|
||||
func findSoraS3ProfileIndex(items []soraS3ProfileStoreItem, profileID string) int {
|
||||
for idx := range items {
|
||||
if items[idx].ProfileID == profileID {
|
||||
return idx
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func isEmptyLegacySoraS3Settings(settings *SoraS3Settings) bool {
|
||||
if settings == nil {
|
||||
return true
|
||||
}
|
||||
if settings.Enabled {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.Endpoint) != "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.Region) != "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.Bucket) != "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.AccessKeyID) != "" {
|
||||
return false
|
||||
}
|
||||
if settings.SecretAccessKey != "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.Prefix) != "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(settings.CDNURL) != "" {
|
||||
return false
|
||||
}
|
||||
return settings.DefaultStorageQuotaBytes == 0
|
||||
}
|
||||
|
||||
func maxInt64(value int64, min int64) int64 {
|
||||
if value < min {
|
||||
return min
|
||||
}
|
||||
return value
|
||||
}
|
||||
|
||||
@@ -113,46 +113,6 @@ type DefaultSubscriptionSetting struct {
|
||||
ValidityDays int `json:"validity_days"`
|
||||
}
|
||||
|
||||
// SoraS3Settings Sora S3 存储配置 (从本地版本合并)
|
||||
type SoraS3Settings struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
Region string `json:"region"`
|
||||
Bucket string `json:"bucket"`
|
||||
AccessKeyID string `json:"access_key_id"`
|
||||
SecretAccessKey string `json:"secret_access_key"` // 仅内部使用，不直接返回前端
|
||||
SecretAccessKeyConfigured bool `json:"secret_access_key_configured"` // 前端展示用
|
||||
Prefix string `json:"prefix"`
|
||||
ForcePathStyle bool `json:"force_path_style"`
|
||||
CDNURL string `json:"cdn_url"`
|
||||
DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"`
|
||||
}
|
||||
|
||||
// SoraS3Profile Sora S3 多配置项（服务内部模型）(从本地版本合并)
|
||||
type SoraS3Profile struct {
|
||||
ProfileID string `json:"profile_id"`
|
||||
Name string `json:"name"`
|
||||
IsActive bool `json:"is_active"`
|
||||
Enabled bool `json:"enabled"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
Region string `json:"region"`
|
||||
Bucket string `json:"bucket"`
|
||||
AccessKeyID string `json:"access_key_id"`
|
||||
SecretAccessKey string `json:"-"` // 仅内部使用，不直接返回前端
|
||||
SecretAccessKeyConfigured bool `json:"secret_access_key_configured"` // 前端展示用
|
||||
Prefix string `json:"prefix"`
|
||||
ForcePathStyle bool `json:"force_path_style"`
|
||||
CDNURL string `json:"cdn_url"`
|
||||
DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"`
|
||||
UpdatedAt string `json:"updated_at"`
|
||||
}
|
||||
|
||||
// SoraS3ProfileList Sora S3 多配置列表 (从本地版本合并)
|
||||
type SoraS3ProfileList struct {
|
||||
ActiveProfileID string `json:"active_profile_id"`
|
||||
Items []SoraS3Profile `json:"items"`
|
||||
}
|
||||
|
||||
type PublicSettings struct {
|
||||
RegistrationEnabled bool
|
||||
EmailVerifyEnabled bool
|
||||
|
||||
@@ -1,40 +0,0 @@
|
||||
package service
|
||||
|
||||
import "context"
|
||||
|
||||
// SoraAccountRepository Sora 账号扩展表仓储接口
|
||||
// 用于管理 sora_accounts 表，与 accounts 主表形成双表结构。
|
||||
//
|
||||
// 设计说明：
|
||||
// - sora_accounts 表存储 Sora 账号的 OAuth 凭证副本
|
||||
// - Sora gateway 优先读取此表的字段以获得更好的查询性能
|
||||
// - 主表 accounts 通过 credentials JSON 字段也存储相同信息
|
||||
// - Token 刷新时需要同时更新两个表以保持数据一致性
|
||||
type SoraAccountRepository interface {
|
||||
// Upsert 创建或更新 Sora 账号扩展信息
|
||||
// accountID: 关联的 accounts.id
|
||||
// updates: 要更新的字段，支持 access_token、refresh_token、session_token
|
||||
//
|
||||
// 如果记录不存在则创建，存在则更新。
|
||||
// 用于：
|
||||
// 1. 创建 Sora 账号时初始化扩展表
|
||||
// 2. Token 刷新时同步更新扩展表
|
||||
Upsert(ctx context.Context, accountID int64, updates map[string]any) error
|
||||
|
||||
// GetByAccountID 根据账号 ID 获取 Sora 扩展信息
|
||||
// 返回 nil, nil 表示记录不存在（非错误）
|
||||
GetByAccountID(ctx context.Context, accountID int64) (*SoraAccount, error)
|
||||
|
||||
// Delete 删除 Sora 账号扩展信息
|
||||
// 通常由外键 ON DELETE CASCADE 自动处理，此方法用于手动清理
|
||||
Delete(ctx context.Context, accountID int64) error
|
||||
}
|
||||
|
||||
// SoraAccount Sora 账号扩展信息
|
||||
// 对应 sora_accounts 表，存储 Sora 账号的 OAuth 凭证副本
|
||||
type SoraAccount struct {
|
||||
AccountID int64 // 关联的 accounts.id
|
||||
AccessToken string // OAuth access_token
|
||||
RefreshToken string // OAuth refresh_token
|
||||
SessionToken string // Session token（可选，用于 ST→AT 兜底）
|
||||
}
|
||||
@@ -1,117 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// SoraClient 定义直连 Sora 的任务操作接口。
|
||||
type SoraClient interface {
|
||||
Enabled() bool
|
||||
UploadImage(ctx context.Context, account *Account, data []byte, filename string) (string, error)
|
||||
CreateImageTask(ctx context.Context, account *Account, req SoraImageRequest) (string, error)
|
||||
CreateVideoTask(ctx context.Context, account *Account, req SoraVideoRequest) (string, error)
|
||||
CreateStoryboardTask(ctx context.Context, account *Account, req SoraStoryboardRequest) (string, error)
|
||||
UploadCharacterVideo(ctx context.Context, account *Account, data []byte) (string, error)
|
||||
GetCameoStatus(ctx context.Context, account *Account, cameoID string) (*SoraCameoStatus, error)
|
||||
DownloadCharacterImage(ctx context.Context, account *Account, imageURL string) ([]byte, error)
|
||||
UploadCharacterImage(ctx context.Context, account *Account, data []byte) (string, error)
|
||||
FinalizeCharacter(ctx context.Context, account *Account, req SoraCharacterFinalizeRequest) (string, error)
|
||||
SetCharacterPublic(ctx context.Context, account *Account, cameoID string) error
|
||||
DeleteCharacter(ctx context.Context, account *Account, characterID string) error
|
||||
PostVideoForWatermarkFree(ctx context.Context, account *Account, generationID string) (string, error)
|
||||
DeletePost(ctx context.Context, account *Account, postID string) error
|
||||
GetWatermarkFreeURLCustom(ctx context.Context, account *Account, parseURL, parseToken, postID string) (string, error)
|
||||
EnhancePrompt(ctx context.Context, account *Account, prompt, expansionLevel string, durationS int) (string, error)
|
||||
GetImageTask(ctx context.Context, account *Account, taskID string) (*SoraImageTaskStatus, error)
|
||||
GetVideoTask(ctx context.Context, account *Account, taskID string) (*SoraVideoTaskStatus, error)
|
||||
}
|
||||
|
||||
// SoraImageRequest 图片生成请求参数
|
||||
type SoraImageRequest struct {
|
||||
Prompt string
|
||||
Width int
|
||||
Height int
|
||||
MediaID string
|
||||
}
|
||||
|
||||
// SoraVideoRequest 视频生成请求参数
|
||||
type SoraVideoRequest struct {
|
||||
Prompt string
|
||||
Orientation string
|
||||
Frames int
|
||||
Model string
|
||||
Size string
|
||||
VideoCount int
|
||||
MediaID string
|
||||
RemixTargetID string
|
||||
CameoIDs []string
|
||||
}
|
||||
|
||||
// SoraStoryboardRequest 分镜视频生成请求参数
|
||||
type SoraStoryboardRequest struct {
|
||||
Prompt string
|
||||
Orientation string
|
||||
Frames int
|
||||
Model string
|
||||
Size string
|
||||
MediaID string
|
||||
}
|
||||
|
||||
// SoraImageTaskStatus 图片任务状态
|
||||
type SoraImageTaskStatus struct {
|
||||
ID string
|
||||
Status string
|
||||
ProgressPct float64
|
||||
URLs []string
|
||||
ErrorMsg string
|
||||
}
|
||||
|
||||
// SoraVideoTaskStatus 视频任务状态
|
||||
type SoraVideoTaskStatus struct {
|
||||
ID string
|
||||
Status string
|
||||
ProgressPct int
|
||||
URLs []string
|
||||
GenerationID string
|
||||
ErrorMsg string
|
||||
}
|
||||
|
||||
// SoraCameoStatus 角色处理中间态
|
||||
type SoraCameoStatus struct {
|
||||
Status string
|
||||
StatusMessage string
|
||||
DisplayNameHint string
|
||||
UsernameHint string
|
||||
ProfileAssetURL string
|
||||
InstructionSetHint any
|
||||
InstructionSet any
|
||||
}
|
||||
|
||||
// SoraCharacterFinalizeRequest 角色定稿请求参数
|
||||
type SoraCharacterFinalizeRequest struct {
|
||||
CameoID string
|
||||
Username string
|
||||
DisplayName string
|
||||
ProfileAssetPointer string
|
||||
InstructionSet any
|
||||
}
|
||||
|
||||
// SoraUpstreamError 上游错误
|
||||
type SoraUpstreamError struct {
|
||||
StatusCode int
|
||||
Message string
|
||||
Headers http.Header
|
||||
Body []byte
|
||||
}
|
||||
|
||||
func (e *SoraUpstreamError) Error() string {
|
||||
if e == nil {
|
||||
return "sora upstream error"
|
||||
}
|
||||
if e.Message != "" {
|
||||
return fmt.Sprintf("sora upstream error: %d %s", e.StatusCode, e.Message)
|
||||
}
|
||||
return fmt.Sprintf("sora upstream error: %d", e.StatusCode)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,564 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var _ SoraClient = (*stubSoraClientForPoll)(nil)
|
||||
|
||||
type stubSoraClientForPoll struct {
|
||||
imageStatus *SoraImageTaskStatus
|
||||
videoStatus *SoraVideoTaskStatus
|
||||
imageCalls int
|
||||
videoCalls int
|
||||
enhanced string
|
||||
enhanceErr error
|
||||
storyboard bool
|
||||
videoReq SoraVideoRequest
|
||||
parseErr error
|
||||
postCalls int
|
||||
deleteCalls int
|
||||
}
|
||||
|
||||
func (s *stubSoraClientForPoll) Enabled() bool { return true }
|
||||
func (s *stubSoraClientForPoll) UploadImage(ctx context.Context, account *Account, data []byte, filename string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) CreateImageTask(ctx context.Context, account *Account, req SoraImageRequest) (string, error) {
|
||||
return "task-image", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) CreateVideoTask(ctx context.Context, account *Account, req SoraVideoRequest) (string, error) {
|
||||
s.videoReq = req
|
||||
return "task-video", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) CreateStoryboardTask(ctx context.Context, account *Account, req SoraStoryboardRequest) (string, error) {
|
||||
s.storyboard = true
|
||||
return "task-video", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) UploadCharacterVideo(ctx context.Context, account *Account, data []byte) (string, error) {
|
||||
return "cameo-1", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) GetCameoStatus(ctx context.Context, account *Account, cameoID string) (*SoraCameoStatus, error) {
|
||||
return &SoraCameoStatus{
|
||||
Status: "finalized",
|
||||
StatusMessage: "Completed",
|
||||
DisplayNameHint: "Character",
|
||||
UsernameHint: "user.character",
|
||||
ProfileAssetURL: "https://example.com/avatar.webp",
|
||||
}, nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) DownloadCharacterImage(ctx context.Context, account *Account, imageURL string) ([]byte, error) {
|
||||
return []byte("avatar"), nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) UploadCharacterImage(ctx context.Context, account *Account, data []byte) (string, error) {
|
||||
return "asset-pointer", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) FinalizeCharacter(ctx context.Context, account *Account, req SoraCharacterFinalizeRequest) (string, error) {
|
||||
return "character-1", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) SetCharacterPublic(ctx context.Context, account *Account, cameoID string) error {
|
||||
return nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) DeleteCharacter(ctx context.Context, account *Account, characterID string) error {
|
||||
return nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) PostVideoForWatermarkFree(ctx context.Context, account *Account, generationID string) (string, error) {
|
||||
s.postCalls++
|
||||
return "s_post", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) DeletePost(ctx context.Context, account *Account, postID string) error {
|
||||
s.deleteCalls++
|
||||
return nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) GetWatermarkFreeURLCustom(ctx context.Context, account *Account, parseURL, parseToken, postID string) (string, error) {
|
||||
if s.parseErr != nil {
|
||||
return "", s.parseErr
|
||||
}
|
||||
return "https://example.com/no-watermark.mp4", nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) EnhancePrompt(ctx context.Context, account *Account, prompt, expansionLevel string, durationS int) (string, error) {
|
||||
if s.enhanced != "" {
|
||||
return s.enhanced, s.enhanceErr
|
||||
}
|
||||
return "enhanced prompt", s.enhanceErr
|
||||
}
|
||||
func (s *stubSoraClientForPoll) GetImageTask(ctx context.Context, account *Account, taskID string) (*SoraImageTaskStatus, error) {
|
||||
s.imageCalls++
|
||||
return s.imageStatus, nil
|
||||
}
|
||||
func (s *stubSoraClientForPoll) GetVideoTask(ctx context.Context, account *Account, taskID string) (*SoraVideoTaskStatus, error) {
|
||||
s.videoCalls++
|
||||
return s.videoStatus, nil
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_PollImageTaskCompleted(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
imageStatus: &SoraImageTaskStatus{
|
||||
Status: "completed",
|
||||
URLs: []string{"https://example.com/a.png"},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
service := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
|
||||
urls, err := service.pollImageTask(context.Background(), nil, &Account{ID: 1}, "task", false)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"https://example.com/a.png"}, urls)
|
||||
require.Equal(t, 1, client.imageCalls)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardPromptEnhance(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
enhanced: "cinematic prompt",
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{
|
||||
ID: 1,
|
||||
Platform: PlatformSora,
|
||||
Status: StatusActive,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"prompt-enhance-short-10s": "prompt-enhance-short-15s",
|
||||
},
|
||||
},
|
||||
}
|
||||
body := []byte(`{"model":"prompt-enhance-short-10s","messages":[{"role":"user","content":"cat running"}],"stream":false}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "prompt", result.MediaType)
|
||||
require.Equal(t, "prompt-enhance-short-10s", result.Model)
|
||||
require.Equal(t, "prompt-enhance-short-15s", result.UpstreamModel)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardStoryboardPrompt(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
videoStatus: &SoraVideoTaskStatus{
|
||||
Status: "completed",
|
||||
URLs: []string{"https://example.com/v.mp4"},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{ID: 1, Platform: PlatformSora, Status: StatusActive}
|
||||
body := []byte(`{"model":"sora2-landscape-10s","messages":[{"role":"user","content":"[5.0s]猫猫跳伞 [5.0s]猫猫落地"}],"stream":false}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, client.storyboard)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardVideoCount(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
videoStatus: &SoraVideoTaskStatus{
|
||||
Status: "completed",
|
||||
URLs: []string{"https://example.com/v.mp4"},
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{ID: 1, Platform: PlatformSora, Status: StatusActive}
|
||||
body := []byte(`{"model":"sora2-landscape-10s","messages":[{"role":"user","content":"cat running"}],"video_count":3,"stream":false}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, 3, client.videoReq.VideoCount)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardCharacterOnly(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{ID: 1, Platform: PlatformSora, Status: StatusActive}
|
||||
body := []byte(`{"model":"sora2-landscape-10s","video":"aGVsbG8=","stream":false}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "prompt", result.MediaType)
|
||||
require.Equal(t, 0, client.videoCalls)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardWatermarkFallback(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
videoStatus: &SoraVideoTaskStatus{
|
||||
Status: "completed",
|
||||
URLs: []string{"https://example.com/original.mp4"},
|
||||
GenerationID: "gen_1",
|
||||
},
|
||||
parseErr: errors.New("parse failed"),
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{ID: 1, Platform: PlatformSora, Status: StatusActive}
|
||||
body := []byte(`{"model":"sora2-landscape-10s","messages":[{"role":"user","content":"cat running"}],"stream":false,"watermark_free":true,"watermark_parse_method":"custom","watermark_parse_url":"https://parser.example.com","watermark_parse_token":"token","watermark_fallback_on_failure":true}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "https://example.com/original.mp4", result.MediaURL)
|
||||
require.Equal(t, 1, client.postCalls)
|
||||
require.Equal(t, 0, client.deleteCalls)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_ForwardWatermarkCustomSuccessAndDelete(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
videoStatus: &SoraVideoTaskStatus{
|
||||
Status: "completed",
|
||||
URLs: []string{"https://example.com/original.mp4"},
|
||||
GenerationID: "gen_1",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
account := &Account{ID: 1, Platform: PlatformSora, Status: StatusActive}
|
||||
body := []byte(`{"model":"sora2-landscape-10s","messages":[{"role":"user","content":"cat running"}],"stream":false,"watermark_free":true,"watermark_parse_method":"custom","watermark_parse_url":"https://parser.example.com","watermark_parse_token":"token","watermark_delete_post":true}`)
|
||||
|
||||
result, err := svc.Forward(context.Background(), nil, account, body, false)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, "https://example.com/no-watermark.mp4", result.MediaURL)
|
||||
require.Equal(t, 1, client.postCalls)
|
||||
require.Equal(t, 1, client.deleteCalls)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_PollVideoTaskFailed(t *testing.T) {
|
||||
client := &stubSoraClientForPoll{
|
||||
videoStatus: &SoraVideoTaskStatus{
|
||||
Status: "failed",
|
||||
ErrorMsg: "reject",
|
||||
},
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 1,
|
||||
MaxPollAttempts: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
service := NewSoraGatewayService(client, nil, nil, cfg)
|
||||
|
||||
status, err := service.pollVideoTaskDetailed(context.Background(), nil, &Account{ID: 1}, "task", false)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, status)
|
||||
require.Contains(t, err.Error(), "reject")
|
||||
require.Equal(t, 1, client.videoCalls)
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_BuildSoraMediaURLSigned(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Gateway: config.GatewayConfig{
|
||||
SoraMediaSigningKey: "test-key",
|
||||
SoraMediaSignedURLTTLSeconds: 600,
|
||||
},
|
||||
}
|
||||
service := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
|
||||
url := service.buildSoraMediaURL("/image/2025/01/01/a.png", "")
|
||||
require.Contains(t, url, "/sora/media-signed")
|
||||
require.Contains(t, url, "expires=")
|
||||
require.Contains(t, url, "sig=")
|
||||
}
|
||||
|
||||
func TestNormalizeSoraMediaURLs_Empty(t *testing.T) {
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
result := svc.normalizeSoraMediaURLs(nil)
|
||||
require.Empty(t, result)
|
||||
|
||||
result = svc.normalizeSoraMediaURLs([]string{})
|
||||
require.Empty(t, result)
|
||||
}
|
||||
|
||||
func TestNormalizeSoraMediaURLs_HTTPUrls(t *testing.T) {
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
urls := []string{"https://example.com/a.png", "http://example.com/b.mp4"}
|
||||
result := svc.normalizeSoraMediaURLs(urls)
|
||||
require.Equal(t, urls, result)
|
||||
}
|
||||
|
||||
func TestNormalizeSoraMediaURLs_LocalPaths(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
svc := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
urls := []string{"/image/2025/01/a.png", "video/2025/01/b.mp4"}
|
||||
result := svc.normalizeSoraMediaURLs(urls)
|
||||
require.Len(t, result, 2)
|
||||
require.Contains(t, result[0], "/sora/media")
|
||||
require.Contains(t, result[1], "/sora/media")
|
||||
}
|
||||
|
||||
func TestNormalizeSoraMediaURLs_SkipsBlank(t *testing.T) {
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
urls := []string{"https://example.com/a.png", "", " ", "https://example.com/b.png"}
|
||||
result := svc.normalizeSoraMediaURLs(urls)
|
||||
require.Len(t, result, 2)
|
||||
}
|
||||
|
||||
func TestBuildSoraContent_Image(t *testing.T) {
|
||||
content := buildSoraContent("image", []string{"https://a.com/1.png", "https://a.com/2.png"})
|
||||
require.Contains(t, content, "")
|
||||
require.Contains(t, content, "")
|
||||
}
|
||||
|
||||
func TestBuildSoraContent_Video(t *testing.T) {
|
||||
content := buildSoraContent("video", []string{"https://a.com/v.mp4"})
|
||||
require.Contains(t, content, "<video src='https://a.com/v.mp4'")
|
||||
}
|
||||
|
||||
func TestBuildSoraContent_VideoEmpty(t *testing.T) {
|
||||
content := buildSoraContent("video", nil)
|
||||
require.Empty(t, content)
|
||||
}
|
||||
|
||||
func TestBuildSoraContent_Prompt(t *testing.T) {
|
||||
content := buildSoraContent("prompt", nil)
|
||||
require.Empty(t, content)
|
||||
}
|
||||
|
||||
func TestSoraImageSizeFromModel(t *testing.T) {
|
||||
require.Equal(t, "360", soraImageSizeFromModel("gpt-image"))
|
||||
require.Equal(t, "540", soraImageSizeFromModel("gpt-image-landscape"))
|
||||
require.Equal(t, "540", soraImageSizeFromModel("gpt-image-portrait"))
|
||||
require.Equal(t, "540", soraImageSizeFromModel("something-landscape"))
|
||||
require.Equal(t, "360", soraImageSizeFromModel("unknown-model"))
|
||||
}
|
||||
|
||||
func TestFirstMediaURL(t *testing.T) {
|
||||
require.Equal(t, "", firstMediaURL(nil))
|
||||
require.Equal(t, "", firstMediaURL([]string{}))
|
||||
require.Equal(t, "a", firstMediaURL([]string{"a", "b"}))
|
||||
}
|
||||
|
||||
func TestSoraProErrorMessage(t *testing.T) {
|
||||
require.Contains(t, soraProErrorMessage("sora2pro-hd", ""), "Pro-HD")
|
||||
require.Contains(t, soraProErrorMessage("sora2pro", ""), "Pro")
|
||||
require.Empty(t, soraProErrorMessage("sora-basic", ""))
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_WriteSoraError_StreamEscapesJSON(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
svc.writeSoraError(c, http.StatusBadGateway, "upstream_error", "invalid \"prompt\"\nline2", true)
|
||||
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "event: error\n")
|
||||
require.Contains(t, body, "data: [DONE]\n\n")
|
||||
|
||||
lines := strings.Split(body, "\n")
|
||||
require.GreaterOrEqual(t, len(lines), 2)
|
||||
require.Equal(t, "event: error", lines[0])
|
||||
require.True(t, strings.HasPrefix(lines[1], "data: "))
|
||||
|
||||
data := strings.TrimPrefix(lines[1], "data: ")
|
||||
var parsed map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(data), &parsed))
|
||||
errObj, ok := parsed["error"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "upstream_error", errObj["type"])
|
||||
require.Equal(t, "invalid \"prompt\"\nline2", errObj["message"])
|
||||
}
|
||||
|
||||
func TestSoraGatewayService_HandleSoraRequestError_FailoverHeadersCloned(t *testing.T) {
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
sourceHeaders := http.Header{}
|
||||
sourceHeaders.Set("cf-ray", "9d01b0e9ecc35829-SEA")
|
||||
|
||||
err := svc.handleSoraRequestError(
|
||||
context.Background(),
|
||||
&Account{ID: 1, Platform: PlatformSora},
|
||||
&SoraUpstreamError{
|
||||
StatusCode: http.StatusForbidden,
|
||||
Message: "forbidden",
|
||||
Headers: sourceHeaders,
|
||||
Body: []byte(`<!DOCTYPE html><title>Just a moment...</title>`),
|
||||
},
|
||||
"sora2-landscape-10s",
|
||||
nil,
|
||||
false,
|
||||
)
|
||||
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.NotNil(t, failoverErr.ResponseHeaders)
|
||||
require.Equal(t, "9d01b0e9ecc35829-SEA", failoverErr.ResponseHeaders.Get("cf-ray"))
|
||||
|
||||
sourceHeaders.Set("cf-ray", "mutated-after-return")
|
||||
require.Equal(t, "9d01b0e9ecc35829-SEA", failoverErr.ResponseHeaders.Get("cf-ray"))
|
||||
}
|
||||
|
||||
func TestShouldFailoverUpstreamError(t *testing.T) {
|
||||
svc := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
require.True(t, svc.shouldFailoverUpstreamError(401))
|
||||
require.True(t, svc.shouldFailoverUpstreamError(404))
|
||||
require.True(t, svc.shouldFailoverUpstreamError(429))
|
||||
require.True(t, svc.shouldFailoverUpstreamError(500))
|
||||
require.True(t, svc.shouldFailoverUpstreamError(502))
|
||||
require.False(t, svc.shouldFailoverUpstreamError(200))
|
||||
require.False(t, svc.shouldFailoverUpstreamError(400))
|
||||
}
|
||||
|
||||
func TestWithSoraTimeout_NilService(t *testing.T) {
|
||||
var svc *SoraGatewayService
|
||||
ctx, cancel := svc.withSoraTimeout(context.Background(), false)
|
||||
require.NotNil(t, ctx)
|
||||
require.Nil(t, cancel)
|
||||
}
|
||||
|
||||
func TestWithSoraTimeout_ZeroTimeout(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
svc := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
ctx, cancel := svc.withSoraTimeout(context.Background(), false)
|
||||
require.NotNil(t, ctx)
|
||||
require.Nil(t, cancel)
|
||||
}
|
||||
|
||||
func TestWithSoraTimeout_PositiveTimeout(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Gateway: config.GatewayConfig{
|
||||
SoraRequestTimeoutSeconds: 30,
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
ctx, cancel := svc.withSoraTimeout(context.Background(), false)
|
||||
require.NotNil(t, ctx)
|
||||
require.NotNil(t, cancel)
|
||||
cancel()
|
||||
}
|
||||
|
||||
func TestPollInterval(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
PollIntervalSeconds: 5,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
require.Equal(t, 5*time.Second, svc.pollInterval())
|
||||
|
||||
// 默认值
|
||||
svc2 := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
require.True(t, svc2.pollInterval() > 0)
|
||||
}
|
||||
|
||||
func TestPollMaxAttempts(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Client: config.SoraClientConfig{
|
||||
MaxPollAttempts: 100,
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraGatewayService(nil, nil, nil, cfg)
|
||||
require.Equal(t, 100, svc.pollMaxAttempts())
|
||||
|
||||
// 默认值
|
||||
svc2 := NewSoraGatewayService(nil, nil, nil, &config.Config{})
|
||||
require.True(t, svc2.pollMaxAttempts() > 0)
|
||||
}
|
||||
|
||||
func TestDecodeSoraImageInput_BlockPrivateURL(t *testing.T) {
|
||||
_, _, err := decodeSoraImageInput(context.Background(), "http://127.0.0.1/internal.png")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestDecodeSoraImageInput_DataURL(t *testing.T) {
|
||||
encoded := "data:image/png;base64,aGVsbG8="
|
||||
data, filename, err := decodeSoraImageInput(context.Background(), encoded)
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, data)
|
||||
require.Contains(t, filename, ".png")
|
||||
}
|
||||
|
||||
func TestDecodeBase64WithLimit_ExceedLimit(t *testing.T) {
|
||||
data, err := decodeBase64WithLimit("aGVsbG8=", 3)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, data)
|
||||
}
|
||||
|
||||
func TestParseSoraWatermarkOptions_NumericBool(t *testing.T) {
|
||||
body := map[string]any{
|
||||
"watermark_free": float64(1),
|
||||
"watermark_fallback_on_failure": float64(0),
|
||||
}
|
||||
opts := parseSoraWatermarkOptions(body)
|
||||
require.True(t, opts.Enabled)
|
||||
require.False(t, opts.FallbackOnFailure)
|
||||
}
|
||||
|
||||
func TestParseSoraVideoCount(t *testing.T) {
|
||||
require.Equal(t, 1, parseSoraVideoCount(nil))
|
||||
require.Equal(t, 2, parseSoraVideoCount(map[string]any{"video_count": float64(2)}))
|
||||
require.Equal(t, 3, parseSoraVideoCount(map[string]any{"videos": "5"}))
|
||||
require.Equal(t, 1, parseSoraVideoCount(map[string]any{"n_variants": 0}))
|
||||
}
|
||||
@@ -1,532 +0,0 @@
|
||||
//nolint:unused
|
||||
package service
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var soraSSEDataRe = regexp.MustCompile(`^data:\s*`)
|
||||
var soraImageMarkdownRe = regexp.MustCompile(`!\[[^\]]*\]\(([^)]+)\)`)
|
||||
var soraVideoHTMLRe = regexp.MustCompile(`(?i)<video[^>]+src=['"]([^'"]+)['"]`)
|
||||
|
||||
const soraRewriteBufferLimit = 2048
|
||||
|
||||
type soraStreamingResult struct {
|
||||
mediaType string
|
||||
mediaURLs []string
|
||||
imageCount int
|
||||
imageSize string
|
||||
firstTokenMs *int
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) setUpstreamRequestError(c *gin.Context, account *Account, err error) {
|
||||
safeErr := sanitizeUpstreamErrorMessage(err.Error())
|
||||
setOpsUpstreamError(c, 0, safeErr, "")
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: 0,
|
||||
Kind: "request_error",
|
||||
Message: safeErr,
|
||||
})
|
||||
if c != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{
|
||||
"error": gin.H{
|
||||
"type": "upstream_error",
|
||||
"message": "Upstream request failed",
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account) {
|
||||
if s.rateLimitService == nil || account == nil || resp == nil {
|
||||
return
|
||||
}
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, reqModel string) (*ForwardResult, error) {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
_ = resp.Body.Close()
|
||||
resp.Body = io.NopCloser(bytes.NewReader(respBody))
|
||||
|
||||
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody))
|
||||
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
||||
if msg := soraProErrorMessage(reqModel, upstreamMsg); msg != "" {
|
||||
upstreamMsg = msg
|
||||
}
|
||||
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
}
|
||||
upstreamDetail = truncateString(string(respBody), maxBytes)
|
||||
}
|
||||
setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: resp.StatusCode,
|
||||
UpstreamRequestID: resp.Header.Get("x-request-id"),
|
||||
Kind: "http_error",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
|
||||
if c != nil {
|
||||
responsePayload := s.buildErrorPayload(respBody, upstreamMsg)
|
||||
c.JSON(resp.StatusCode, responsePayload)
|
||||
}
|
||||
if upstreamMsg == "" {
|
||||
return nil, fmt.Errorf("upstream error: %d", resp.StatusCode)
|
||||
}
|
||||
return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg)
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) buildErrorPayload(respBody []byte, overrideMessage string) map[string]any {
|
||||
if len(respBody) > 0 {
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(respBody, &payload); err == nil {
|
||||
if errObj, ok := payload["error"].(map[string]any); ok {
|
||||
if overrideMessage != "" {
|
||||
errObj["message"] = overrideMessage
|
||||
}
|
||||
payload["error"] = errObj
|
||||
return payload
|
||||
}
|
||||
}
|
||||
}
|
||||
return map[string]any{
|
||||
"error": map[string]any{
|
||||
"type": "upstream_error",
|
||||
"message": overrideMessage,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel string, clientStream bool) (*soraStreamingResult, error) {
|
||||
if resp == nil {
|
||||
return nil, errors.New("empty response")
|
||||
}
|
||||
|
||||
if clientStream {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
if v := resp.Header.Get("x-request-id"); v != "" {
|
||||
c.Header("x-request-id", v)
|
||||
}
|
||||
}
|
||||
|
||||
w := c.Writer
|
||||
flusher, _ := w.(http.Flusher)
|
||||
|
||||
contentBuilder := strings.Builder{}
|
||||
var firstTokenMs *int
|
||||
var upstreamError error
|
||||
rewriteBuffer := ""
|
||||
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
|
||||
maxLineSize = s.cfg.Gateway.MaxLineSize
|
||||
}
|
||||
scanner.Buffer(make([]byte, 64*1024), maxLineSize)
|
||||
|
||||
sendLine := func(line string) error {
|
||||
if !clientStream {
|
||||
return nil
|
||||
}
|
||||
if _, err := fmt.Fprintf(w, "%s\n", line); err != nil {
|
||||
return err
|
||||
}
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
if soraSSEDataRe.MatchString(line) {
|
||||
data := soraSSEDataRe.ReplaceAllString(line, "")
|
||||
if data == "[DONE]" {
|
||||
if rewriteBuffer != "" {
|
||||
flushLine, flushContent, err := s.flushSoraRewriteBuffer(rewriteBuffer, originalModel)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if flushLine != "" {
|
||||
if flushContent != "" {
|
||||
if _, err := contentBuilder.WriteString(flushContent); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := sendLine(flushLine); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
rewriteBuffer = ""
|
||||
}
|
||||
if err := sendLine("data: [DONE]"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
break
|
||||
}
|
||||
updatedLine, contentDelta, errEvent := s.processSoraSSEData(data, originalModel, &rewriteBuffer)
|
||||
if errEvent != nil && upstreamError == nil {
|
||||
upstreamError = errEvent
|
||||
}
|
||||
if contentDelta != "" {
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
firstTokenMs = &ms
|
||||
}
|
||||
if _, err := contentBuilder.WriteString(contentDelta); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if err := sendLine(updatedLine); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := sendLine(line); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if errors.Is(err, bufio.ErrTooLong) {
|
||||
if clientStream {
|
||||
_, _ = fmt.Fprintf(w, "event: error\ndata: {\"error\":\"response_too_large\"}\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if ctx.Err() == context.DeadlineExceeded && s.rateLimitService != nil && account != nil {
|
||||
s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel)
|
||||
}
|
||||
if clientStream {
|
||||
_, _ = fmt.Fprintf(w, "event: error\ndata: {\"error\":\"stream_read_error\"}\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
content := contentBuilder.String()
|
||||
mediaType, mediaURLs := s.extractSoraMedia(content)
|
||||
if mediaType == "" && isSoraPromptEnhanceModel(originalModel) {
|
||||
mediaType = "prompt"
|
||||
}
|
||||
imageSize := ""
|
||||
imageCount := 0
|
||||
if mediaType == "image" {
|
||||
imageSize = soraImageSizeFromModel(originalModel)
|
||||
imageCount = len(mediaURLs)
|
||||
}
|
||||
|
||||
if upstreamError != nil && !clientStream {
|
||||
if c != nil {
|
||||
c.JSON(http.StatusBadGateway, map[string]any{
|
||||
"error": map[string]any{
|
||||
"type": "upstream_error",
|
||||
"message": upstreamError.Error(),
|
||||
},
|
||||
})
|
||||
}
|
||||
return nil, upstreamError
|
||||
}
|
||||
|
||||
if !clientStream {
|
||||
response := buildSoraNonStreamResponse(content, originalModel)
|
||||
if len(mediaURLs) > 0 {
|
||||
response["media_url"] = mediaURLs[0]
|
||||
if len(mediaURLs) > 1 {
|
||||
response["media_urls"] = mediaURLs
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
return &soraStreamingResult{
|
||||
mediaType: mediaType,
|
||||
mediaURLs: mediaURLs,
|
||||
imageCount: imageCount,
|
||||
imageSize: imageSize,
|
||||
firstTokenMs: firstTokenMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) processSoraSSEData(data string, originalModel string, rewriteBuffer *string) (string, string, error) {
|
||||
if strings.TrimSpace(data) == "" {
|
||||
return "data: ", "", nil
|
||||
}
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(data), &payload); err != nil {
|
||||
return "data: " + data, "", nil
|
||||
}
|
||||
|
||||
if errObj, ok := payload["error"].(map[string]any); ok {
|
||||
if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" {
|
||||
return "data: " + data, "", errors.New(msg)
|
||||
}
|
||||
}
|
||||
|
||||
if model, ok := payload["model"].(string); ok && model != "" && originalModel != "" {
|
||||
payload["model"] = originalModel
|
||||
}
|
||||
|
||||
contentDelta, updated := extractSoraContent(payload)
|
||||
if updated {
|
||||
var rewritten string
|
||||
if rewriteBuffer != nil {
|
||||
rewritten = s.rewriteSoraContentWithBuffer(contentDelta, rewriteBuffer)
|
||||
} else {
|
||||
rewritten = s.rewriteSoraContent(contentDelta)
|
||||
}
|
||||
if rewritten != contentDelta {
|
||||
applySoraContent(payload, rewritten)
|
||||
contentDelta = rewritten
|
||||
}
|
||||
}
|
||||
|
||||
updatedData, err := jsonMarshalRaw(payload)
|
||||
if err != nil {
|
||||
return "data: " + data, contentDelta, nil
|
||||
}
|
||||
return "data: " + string(updatedData), contentDelta, nil
|
||||
}
|
||||
|
||||
func extractSoraContent(payload map[string]any) (string, bool) {
|
||||
choices, ok := payload["choices"].([]any)
|
||||
if !ok || len(choices) == 0 {
|
||||
return "", false
|
||||
}
|
||||
choice, ok := choices[0].(map[string]any)
|
||||
if !ok {
|
||||
return "", false
|
||||
}
|
||||
if delta, ok := choice["delta"].(map[string]any); ok {
|
||||
if content, ok := delta["content"].(string); ok {
|
||||
return content, true
|
||||
}
|
||||
}
|
||||
if message, ok := choice["message"].(map[string]any); ok {
|
||||
if content, ok := message["content"].(string); ok {
|
||||
return content, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func applySoraContent(payload map[string]any, content string) {
|
||||
choices, ok := payload["choices"].([]any)
|
||||
if !ok || len(choices) == 0 {
|
||||
return
|
||||
}
|
||||
choice, ok := choices[0].(map[string]any)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if delta, ok := choice["delta"].(map[string]any); ok {
|
||||
delta["content"] = content
|
||||
choice["delta"] = delta
|
||||
return
|
||||
}
|
||||
if message, ok := choice["message"].(map[string]any); ok {
|
||||
message["content"] = content
|
||||
choice["message"] = message
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) rewriteSoraContentWithBuffer(contentDelta string, buffer *string) string {
|
||||
if buffer == nil {
|
||||
return s.rewriteSoraContent(contentDelta)
|
||||
}
|
||||
if contentDelta == "" && *buffer == "" {
|
||||
return ""
|
||||
}
|
||||
combined := *buffer + contentDelta
|
||||
rewritten := s.rewriteSoraContent(combined)
|
||||
bufferStart := s.findSoraRewriteBufferStart(rewritten)
|
||||
if bufferStart < 0 {
|
||||
*buffer = ""
|
||||
return rewritten
|
||||
}
|
||||
if len(rewritten)-bufferStart > soraRewriteBufferLimit {
|
||||
bufferStart = len(rewritten) - soraRewriteBufferLimit
|
||||
}
|
||||
output := rewritten[:bufferStart]
|
||||
*buffer = rewritten[bufferStart:]
|
||||
return output
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) findSoraRewriteBufferStart(content string) int {
|
||||
minIndex := -1
|
||||
start := 0
|
||||
for {
|
||||
idx := strings.Index(content[start:], "![")
|
||||
if idx < 0 {
|
||||
break
|
||||
}
|
||||
idx += start
|
||||
if !hasSoraImageMatchAt(content, idx) {
|
||||
if minIndex == -1 || idx < minIndex {
|
||||
minIndex = idx
|
||||
}
|
||||
}
|
||||
start = idx + 2
|
||||
}
|
||||
lower := strings.ToLower(content)
|
||||
start = 0
|
||||
for {
|
||||
idx := strings.Index(lower[start:], "<video")
|
||||
if idx < 0 {
|
||||
break
|
||||
}
|
||||
idx += start
|
||||
if !hasSoraVideoMatchAt(content, idx) {
|
||||
if minIndex == -1 || idx < minIndex {
|
||||
minIndex = idx
|
||||
}
|
||||
}
|
||||
start = idx + len("<video")
|
||||
}
|
||||
return minIndex
|
||||
}
|
||||
|
||||
func hasSoraImageMatchAt(content string, idx int) bool {
|
||||
if idx < 0 || idx >= len(content) {
|
||||
return false
|
||||
}
|
||||
loc := soraImageMarkdownRe.FindStringIndex(content[idx:])
|
||||
return loc != nil && loc[0] == 0
|
||||
}
|
||||
|
||||
func hasSoraVideoMatchAt(content string, idx int) bool {
|
||||
if idx < 0 || idx >= len(content) {
|
||||
return false
|
||||
}
|
||||
loc := soraVideoHTMLRe.FindStringIndex(content[idx:])
|
||||
return loc != nil && loc[0] == 0
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) rewriteSoraContent(content string) string {
|
||||
if content == "" {
|
||||
return content
|
||||
}
|
||||
content = soraImageMarkdownRe.ReplaceAllStringFunc(content, func(match string) string {
|
||||
sub := soraImageMarkdownRe.FindStringSubmatch(match)
|
||||
if len(sub) < 2 {
|
||||
return match
|
||||
}
|
||||
rewritten := s.rewriteSoraURL(sub[1])
|
||||
if rewritten == sub[1] {
|
||||
return match
|
||||
}
|
||||
return strings.Replace(match, sub[1], rewritten, 1)
|
||||
})
|
||||
content = soraVideoHTMLRe.ReplaceAllStringFunc(content, func(match string) string {
|
||||
sub := soraVideoHTMLRe.FindStringSubmatch(match)
|
||||
if len(sub) < 2 {
|
||||
return match
|
||||
}
|
||||
rewritten := s.rewriteSoraURL(sub[1])
|
||||
if rewritten == sub[1] {
|
||||
return match
|
||||
}
|
||||
return strings.Replace(match, sub[1], rewritten, 1)
|
||||
})
|
||||
return content
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) flushSoraRewriteBuffer(buffer string, originalModel string) (string, string, error) {
|
||||
if buffer == "" {
|
||||
return "", "", nil
|
||||
}
|
||||
rewritten := s.rewriteSoraContent(buffer)
|
||||
payload := map[string]any{
|
||||
"choices": []any{
|
||||
map[string]any{
|
||||
"delta": map[string]any{
|
||||
"content": rewritten,
|
||||
},
|
||||
"index": 0,
|
||||
},
|
||||
},
|
||||
}
|
||||
if originalModel != "" {
|
||||
payload["model"] = originalModel
|
||||
}
|
||||
updatedData, err := jsonMarshalRaw(payload)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return "data: " + string(updatedData), rewritten, nil
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) rewriteSoraURL(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return raw
|
||||
}
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
path := parsed.Path
|
||||
if !strings.HasPrefix(path, "/tmp/") && !strings.HasPrefix(path, "/static/") {
|
||||
return raw
|
||||
}
|
||||
return s.buildSoraMediaURL(path, parsed.RawQuery)
|
||||
}
|
||||
|
||||
func (s *SoraGatewayService) extractSoraMedia(content string) (string, []string) {
|
||||
if content == "" {
|
||||
return "", nil
|
||||
}
|
||||
if match := soraVideoHTMLRe.FindStringSubmatch(content); len(match) > 1 {
|
||||
return "video", []string{match[1]}
|
||||
}
|
||||
imageMatches := soraImageMarkdownRe.FindAllStringSubmatch(content, -1)
|
||||
if len(imageMatches) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
urls := make([]string, 0, len(imageMatches))
|
||||
for _, match := range imageMatches {
|
||||
if len(match) > 1 {
|
||||
urls = append(urls, match[1])
|
||||
}
|
||||
}
|
||||
return "image", urls
|
||||
}
|
||||
|
||||
func isSoraPromptEnhanceModel(model string) bool {
|
||||
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(model)), "prompt-enhance")
|
||||
}
|
||||
@@ -1,63 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SoraGeneration 代表一条 Sora 客户端生成记录。
|
||||
type SoraGeneration struct {
|
||||
ID int64 `json:"id"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID *int64 `json:"api_key_id,omitempty"`
|
||||
Model string `json:"model"`
|
||||
Prompt string `json:"prompt"`
|
||||
MediaType string `json:"media_type"` // video / image
|
||||
Status string `json:"status"` // pending / generating / completed / failed / cancelled
|
||||
MediaURL string `json:"media_url"` // 主媒体 URL（预签名或 CDN）
|
||||
MediaURLs []string `json:"media_urls"` // 多图时的 URL 数组
|
||||
FileSizeBytes int64 `json:"file_size_bytes"`
|
||||
StorageType string `json:"storage_type"` // s3 / local / upstream / none
|
||||
S3ObjectKeys []string `json:"s3_object_keys"` // S3 object key 数组
|
||||
UpstreamTaskID string `json:"upstream_task_id"`
|
||||
ErrorMessage string `json:"error_message"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt *time.Time `json:"completed_at,omitempty"`
|
||||
}
|
||||
|
||||
// Sora 生成记录状态常量
|
||||
const (
|
||||
SoraGenStatusPending = "pending"
|
||||
SoraGenStatusGenerating = "generating"
|
||||
SoraGenStatusCompleted = "completed"
|
||||
SoraGenStatusFailed = "failed"
|
||||
SoraGenStatusCancelled = "cancelled"
|
||||
)
|
||||
|
||||
// Sora 存储类型常量
|
||||
const (
|
||||
SoraStorageTypeS3 = "s3"
|
||||
SoraStorageTypeLocal = "local"
|
||||
SoraStorageTypeUpstream = "upstream"
|
||||
SoraStorageTypeNone = "none"
|
||||
)
|
||||
|
||||
// SoraGenerationListParams 查询生成记录的参数。
|
||||
type SoraGenerationListParams struct {
|
||||
UserID int64
|
||||
Status string // 可选筛选
|
||||
StorageType string // 可选筛选
|
||||
MediaType string // 可选筛选
|
||||
Page int
|
||||
PageSize int
|
||||
}
|
||||
|
||||
// SoraGenerationRepository 生成记录持久化接口。
|
||||
type SoraGenerationRepository interface {
|
||||
Create(ctx context.Context, gen *SoraGeneration) error
|
||||
GetByID(ctx context.Context, id int64) (*SoraGeneration, error)
|
||||
Update(ctx context.Context, gen *SoraGeneration) error
|
||||
Delete(ctx context.Context, id int64) error
|
||||
List(ctx context.Context, params SoraGenerationListParams) ([]*SoraGeneration, int64, error)
|
||||
CountByUserAndStatus(ctx context.Context, userID int64, statuses []string) (int64, error)
|
||||
}
|
||||
@@ -1,341 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrSoraGenerationConcurrencyLimit 表示用户进行中的任务数超限。
|
||||
ErrSoraGenerationConcurrencyLimit = errors.New("sora generation concurrent limit exceeded")
|
||||
// ErrSoraGenerationStateConflict 表示状态已发生变化（例如任务已取消）。
|
||||
ErrSoraGenerationStateConflict = errors.New("sora generation state conflict")
|
||||
// ErrSoraGenerationNotActive 表示任务不在可取消状态。
|
||||
ErrSoraGenerationNotActive = errors.New("sora generation is not active")
|
||||
)
|
||||
|
||||
const soraGenerationActiveLimit = 3
|
||||
|
||||
type soraGenerationRepoAtomicCreator interface {
|
||||
CreatePendingWithLimit(ctx context.Context, gen *SoraGeneration, activeStatuses []string, maxActive int64) error
|
||||
}
|
||||
|
||||
type soraGenerationRepoConditionalUpdater interface {
|
||||
UpdateGeneratingIfPending(ctx context.Context, id int64, upstreamTaskID string) (bool, error)
|
||||
UpdateCompletedIfActive(ctx context.Context, id int64, mediaURL string, mediaURLs []string, storageType string, s3Keys []string, fileSizeBytes int64, completedAt time.Time) (bool, error)
|
||||
UpdateFailedIfActive(ctx context.Context, id int64, errMsg string, completedAt time.Time) (bool, error)
|
||||
UpdateCancelledIfActive(ctx context.Context, id int64, completedAt time.Time) (bool, error)
|
||||
UpdateStorageIfCompleted(ctx context.Context, id int64, mediaURL string, mediaURLs []string, storageType string, s3Keys []string, fileSizeBytes int64) (bool, error)
|
||||
}
|
||||
|
||||
// SoraGenerationService 管理 Sora 客户端的生成记录 CRUD。
|
||||
type SoraGenerationService struct {
|
||||
genRepo SoraGenerationRepository
|
||||
s3Storage *SoraS3Storage
|
||||
quotaService *SoraQuotaService
|
||||
}
|
||||
|
||||
// NewSoraGenerationService 创建生成记录服务。
|
||||
func NewSoraGenerationService(
|
||||
genRepo SoraGenerationRepository,
|
||||
s3Storage *SoraS3Storage,
|
||||
quotaService *SoraQuotaService,
|
||||
) *SoraGenerationService {
|
||||
return &SoraGenerationService{
|
||||
genRepo: genRepo,
|
||||
s3Storage: s3Storage,
|
||||
quotaService: quotaService,
|
||||
}
|
||||
}
|
||||
|
||||
// CreatePending 创建一条 pending 状态的生成记录。
|
||||
func (s *SoraGenerationService) CreatePending(ctx context.Context, userID int64, apiKeyID *int64, model, prompt, mediaType string) (*SoraGeneration, error) {
|
||||
gen := &SoraGeneration{
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
Model: model,
|
||||
Prompt: prompt,
|
||||
MediaType: mediaType,
|
||||
Status: SoraGenStatusPending,
|
||||
StorageType: SoraStorageTypeNone,
|
||||
}
|
||||
if atomicCreator, ok := s.genRepo.(soraGenerationRepoAtomicCreator); ok {
|
||||
if err := atomicCreator.CreatePendingWithLimit(
|
||||
ctx,
|
||||
gen,
|
||||
[]string{SoraGenStatusPending, SoraGenStatusGenerating},
|
||||
soraGenerationActiveLimit,
|
||||
); err != nil {
|
||||
if errors.Is(err, ErrSoraGenerationConcurrencyLimit) {
|
||||
return nil, err
|
||||
}
|
||||
return nil, fmt.Errorf("create generation: %w", err)
|
||||
}
|
||||
logger.LegacyPrintf("service.sora_gen", "[SoraGen] 创建记录 id=%d user=%d model=%s", gen.ID, userID, model)
|
||||
return gen, nil
|
||||
}
|
||||
|
||||
if err := s.genRepo.Create(ctx, gen); err != nil {
|
||||
return nil, fmt.Errorf("create generation: %w", err)
|
||||
}
|
||||
logger.LegacyPrintf("service.sora_gen", "[SoraGen] 创建记录 id=%d user=%d model=%s", gen.ID, userID, model)
|
||||
return gen, nil
|
||||
}
|
||||
|
||||
// MarkGenerating 标记为生成中。
|
||||
func (s *SoraGenerationService) MarkGenerating(ctx context.Context, id int64, upstreamTaskID string) error {
|
||||
if updater, ok := s.genRepo.(soraGenerationRepoConditionalUpdater); ok {
|
||||
updated, err := updater.UpdateGeneratingIfPending(ctx, id, upstreamTaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !updated {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.Status != SoraGenStatusPending {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
gen.Status = SoraGenStatusGenerating
|
||||
gen.UpstreamTaskID = upstreamTaskID
|
||||
return s.genRepo.Update(ctx, gen)
|
||||
}
|
||||
|
||||
// MarkCompleted 标记为已完成。
|
||||
func (s *SoraGenerationService) MarkCompleted(ctx context.Context, id int64, mediaURL string, mediaURLs []string, storageType string, s3Keys []string, fileSizeBytes int64) error {
|
||||
now := time.Now()
|
||||
if updater, ok := s.genRepo.(soraGenerationRepoConditionalUpdater); ok {
|
||||
updated, err := updater.UpdateCompletedIfActive(ctx, id, mediaURL, mediaURLs, storageType, s3Keys, fileSizeBytes, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !updated {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.Status != SoraGenStatusPending && gen.Status != SoraGenStatusGenerating {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
gen.Status = SoraGenStatusCompleted
|
||||
gen.MediaURL = mediaURL
|
||||
gen.MediaURLs = mediaURLs
|
||||
gen.StorageType = storageType
|
||||
gen.S3ObjectKeys = s3Keys
|
||||
gen.FileSizeBytes = fileSizeBytes
|
||||
gen.CompletedAt = &now
|
||||
return s.genRepo.Update(ctx, gen)
|
||||
}
|
||||
|
||||
// MarkFailed 标记为失败。
|
||||
func (s *SoraGenerationService) MarkFailed(ctx context.Context, id int64, errMsg string) error {
|
||||
now := time.Now()
|
||||
if updater, ok := s.genRepo.(soraGenerationRepoConditionalUpdater); ok {
|
||||
updated, err := updater.UpdateFailedIfActive(ctx, id, errMsg, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !updated {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.Status != SoraGenStatusPending && gen.Status != SoraGenStatusGenerating {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
gen.Status = SoraGenStatusFailed
|
||||
gen.ErrorMessage = errMsg
|
||||
gen.CompletedAt = &now
|
||||
return s.genRepo.Update(ctx, gen)
|
||||
}
|
||||
|
||||
// MarkCancelled 标记为已取消。
|
||||
func (s *SoraGenerationService) MarkCancelled(ctx context.Context, id int64) error {
|
||||
now := time.Now()
|
||||
if updater, ok := s.genRepo.(soraGenerationRepoConditionalUpdater); ok {
|
||||
updated, err := updater.UpdateCancelledIfActive(ctx, id, now)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !updated {
|
||||
return ErrSoraGenerationNotActive
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.Status != SoraGenStatusPending && gen.Status != SoraGenStatusGenerating {
|
||||
return ErrSoraGenerationNotActive
|
||||
}
|
||||
gen.Status = SoraGenStatusCancelled
|
||||
gen.CompletedAt = &now
|
||||
return s.genRepo.Update(ctx, gen)
|
||||
}
|
||||
|
||||
// UpdateStorageForCompleted 更新已完成记录的存储信息（不重置 completed_at）。
|
||||
func (s *SoraGenerationService) UpdateStorageForCompleted(
|
||||
ctx context.Context,
|
||||
id int64,
|
||||
mediaURL string,
|
||||
mediaURLs []string,
|
||||
storageType string,
|
||||
s3Keys []string,
|
||||
fileSizeBytes int64,
|
||||
) error {
|
||||
if updater, ok := s.genRepo.(soraGenerationRepoConditionalUpdater); ok {
|
||||
updated, err := updater.UpdateStorageIfCompleted(ctx, id, mediaURL, mediaURLs, storageType, s3Keys, fileSizeBytes)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !updated {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.Status != SoraGenStatusCompleted {
|
||||
return ErrSoraGenerationStateConflict
|
||||
}
|
||||
gen.MediaURL = mediaURL
|
||||
gen.MediaURLs = mediaURLs
|
||||
gen.StorageType = storageType
|
||||
gen.S3ObjectKeys = s3Keys
|
||||
gen.FileSizeBytes = fileSizeBytes
|
||||
return s.genRepo.Update(ctx, gen)
|
||||
}
|
||||
|
||||
// GetByID 获取记录详情（含权限校验）。
|
||||
func (s *SoraGenerationService) GetByID(ctx context.Context, id, userID int64) (*SoraGeneration, error) {
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if gen.UserID != userID {
|
||||
return nil, fmt.Errorf("无权访问此生成记录")
|
||||
}
|
||||
return gen, nil
|
||||
}
|
||||
|
||||
// List 查询生成记录列表（分页 + 筛选）。
|
||||
func (s *SoraGenerationService) List(ctx context.Context, params SoraGenerationListParams) ([]*SoraGeneration, int64, error) {
|
||||
if params.Page <= 0 {
|
||||
params.Page = 1
|
||||
}
|
||||
if params.PageSize <= 0 {
|
||||
params.PageSize = 20
|
||||
}
|
||||
if params.PageSize > 100 {
|
||||
params.PageSize = 100
|
||||
}
|
||||
return s.genRepo.List(ctx, params)
|
||||
}
|
||||
|
||||
// Delete 删除记录（联动 S3/本地文件清理 + 配额释放）。
|
||||
func (s *SoraGenerationService) Delete(ctx context.Context, id, userID int64) error {
|
||||
gen, err := s.genRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if gen.UserID != userID {
|
||||
return fmt.Errorf("无权删除此生成记录")
|
||||
}
|
||||
|
||||
// 清理 S3 文件
|
||||
if gen.StorageType == SoraStorageTypeS3 && len(gen.S3ObjectKeys) > 0 && s.s3Storage != nil {
|
||||
if err := s.s3Storage.DeleteObjects(ctx, gen.S3ObjectKeys); err != nil {
|
||||
logger.LegacyPrintf("service.sora_gen", "[SoraGen] S3 清理失败 id=%d err=%v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
// 释放配额（S3/本地均释放）
|
||||
if gen.FileSizeBytes > 0 && (gen.StorageType == SoraStorageTypeS3 || gen.StorageType == SoraStorageTypeLocal) && s.quotaService != nil {
|
||||
if err := s.quotaService.ReleaseUsage(ctx, userID, gen.FileSizeBytes); err != nil {
|
||||
logger.LegacyPrintf("service.sora_gen", "[SoraGen] 配额释放失败 id=%d err=%v", id, err)
|
||||
}
|
||||
}
|
||||
|
||||
return s.genRepo.Delete(ctx, id)
|
||||
}
|
||||
|
||||
// CountActiveByUser 统计用户进行中的任务数（用于并发限制）。
|
||||
func (s *SoraGenerationService) CountActiveByUser(ctx context.Context, userID int64) (int64, error) {
|
||||
return s.genRepo.CountByUserAndStatus(ctx, userID, []string{SoraGenStatusPending, SoraGenStatusGenerating})
|
||||
}
|
||||
|
||||
// ResolveMediaURLs 为 S3 记录动态生成预签名 URL。
|
||||
func (s *SoraGenerationService) ResolveMediaURLs(ctx context.Context, gen *SoraGeneration) error {
|
||||
if gen == nil || gen.StorageType != SoraStorageTypeS3 || s.s3Storage == nil {
|
||||
return nil
|
||||
}
|
||||
if len(gen.S3ObjectKeys) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
urls := make([]string, len(gen.S3ObjectKeys))
|
||||
var wg sync.WaitGroup
|
||||
var firstErr error
|
||||
var errMu sync.Mutex
|
||||
|
||||
for idx, key := range gen.S3ObjectKeys {
|
||||
wg.Add(1)
|
||||
go func(i int, objectKey string) {
|
||||
defer wg.Done()
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
errMu.Lock()
|
||||
if firstErr == nil {
|
||||
firstErr = fmt.Errorf("goroutine panic fetching S3 URL: %v", r)
|
||||
}
|
||||
errMu.Unlock()
|
||||
}
|
||||
}()
|
||||
url, err := s.s3Storage.GetAccessURL(ctx, objectKey)
|
||||
if err != nil {
|
||||
errMu.Lock()
|
||||
if firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
errMu.Unlock()
|
||||
return
|
||||
}
|
||||
urls[i] = url
|
||||
}(idx, key)
|
||||
}
|
||||
wg.Wait()
|
||||
if firstErr != nil {
|
||||
return firstErr
|
||||
}
|
||||
|
||||
gen.MediaURL = urls[0]
|
||||
gen.MediaURLs = urls
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,876 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ==================== Stub: SoraGenerationRepository ====================
|
||||
|
||||
var _ SoraGenerationRepository = (*stubGenRepo)(nil)
|
||||
|
||||
type stubGenRepo struct {
|
||||
gens map[int64]*SoraGeneration
|
||||
nextID int64
|
||||
createErr error
|
||||
getErr error
|
||||
updateErr error
|
||||
deleteErr error
|
||||
listErr error
|
||||
countErr error
|
||||
countValue int64
|
||||
}
|
||||
|
||||
func newStubGenRepo() *stubGenRepo {
|
||||
return &stubGenRepo{gens: make(map[int64]*SoraGeneration), nextID: 1}
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) Create(_ context.Context, gen *SoraGeneration) error {
|
||||
if r.createErr != nil {
|
||||
return r.createErr
|
||||
}
|
||||
gen.ID = r.nextID
|
||||
gen.CreatedAt = time.Now()
|
||||
r.nextID++
|
||||
r.gens[gen.ID] = gen
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) GetByID(_ context.Context, id int64) (*SoraGeneration, error) {
|
||||
if r.getErr != nil {
|
||||
return nil, r.getErr
|
||||
}
|
||||
if gen, ok := r.gens[id]; ok {
|
||||
return gen, nil
|
||||
}
|
||||
return nil, fmt.Errorf("not found")
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) Update(_ context.Context, gen *SoraGeneration) error {
|
||||
if r.updateErr != nil {
|
||||
return r.updateErr
|
||||
}
|
||||
r.gens[gen.ID] = gen
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) Delete(_ context.Context, id int64) error {
|
||||
if r.deleteErr != nil {
|
||||
return r.deleteErr
|
||||
}
|
||||
delete(r.gens, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) List(_ context.Context, params SoraGenerationListParams) ([]*SoraGeneration, int64, error) {
|
||||
if r.listErr != nil {
|
||||
return nil, 0, r.listErr
|
||||
}
|
||||
var result []*SoraGeneration
|
||||
for _, gen := range r.gens {
|
||||
if gen.UserID != params.UserID {
|
||||
continue
|
||||
}
|
||||
if params.Status != "" && gen.Status != params.Status {
|
||||
continue
|
||||
}
|
||||
if params.StorageType != "" && gen.StorageType != params.StorageType {
|
||||
continue
|
||||
}
|
||||
if params.MediaType != "" && gen.MediaType != params.MediaType {
|
||||
continue
|
||||
}
|
||||
result = append(result, gen)
|
||||
}
|
||||
return result, int64(len(result)), nil
|
||||
}
|
||||
|
||||
func (r *stubGenRepo) CountByUserAndStatus(_ context.Context, userID int64, statuses []string) (int64, error) {
|
||||
if r.countErr != nil {
|
||||
return 0, r.countErr
|
||||
}
|
||||
if r.countValue > 0 {
|
||||
return r.countValue, nil
|
||||
}
|
||||
var count int64
|
||||
statusSet := make(map[string]struct{})
|
||||
for _, s := range statuses {
|
||||
statusSet[s] = struct{}{}
|
||||
}
|
||||
for _, gen := range r.gens {
|
||||
if gen.UserID == userID {
|
||||
if _, ok := statusSet[gen.Status]; ok {
|
||||
count++
|
||||
}
|
||||
}
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
// ==================== Stub: UserRepository (用于 SoraQuotaService) ====================
|
||||
|
||||
var _ UserRepository = (*stubUserRepoForQuota)(nil)
|
||||
|
||||
type stubUserRepoForQuota struct {
|
||||
users map[int64]*User
|
||||
updateErr error
|
||||
}
|
||||
|
||||
func newStubUserRepoForQuota() *stubUserRepoForQuota {
|
||||
return &stubUserRepoForQuota{users: make(map[int64]*User)}
|
||||
}
|
||||
|
||||
func (r *stubUserRepoForQuota) GetByID(_ context.Context, id int64) (*User, error) {
|
||||
if u, ok := r.users[id]; ok {
|
||||
return u, nil
|
||||
}
|
||||
return nil, fmt.Errorf("user not found")
|
||||
}
|
||||
func (r *stubUserRepoForQuota) Update(_ context.Context, user *User) error {
|
||||
if r.updateErr != nil {
|
||||
return r.updateErr
|
||||
}
|
||||
r.users[user.ID] = user
|
||||
return nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) Create(context.Context, *User) error { return nil }
|
||||
func (r *stubUserRepoForQuota) GetByEmail(context.Context, string) (*User, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) GetFirstAdmin(context.Context) (*User, error) { return nil, nil }
|
||||
func (r *stubUserRepoForQuota) Delete(context.Context, int64) error { return nil }
|
||||
func (r *stubUserRepoForQuota) List(context.Context, pagination.PaginationParams) ([]User, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) ListWithFilters(context.Context, pagination.PaginationParams, UserListFilters) ([]User, *pagination.PaginationResult, error) {
|
||||
return nil, nil, nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) UpdateBalance(context.Context, int64, float64) error { return nil }
|
||||
func (r *stubUserRepoForQuota) DeductBalance(context.Context, int64, float64) error { return nil }
|
||||
func (r *stubUserRepoForQuota) UpdateConcurrency(context.Context, int64, int) error { return nil }
|
||||
func (r *stubUserRepoForQuota) ExistsByEmail(context.Context, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) RemoveGroupFromUserAllowedGroups(context.Context, int64, int64) error {
|
||||
return nil
|
||||
}
|
||||
func (r *stubUserRepoForQuota) UpdateTotpSecret(context.Context, int64, *string) error { return nil }
|
||||
func (r *stubUserRepoForQuota) EnableTotp(context.Context, int64) error { return nil }
|
||||
func (r *stubUserRepoForQuota) DisableTotp(context.Context, int64) error { return nil }
|
||||
func (r *stubUserRepoForQuota) AddGroupToAllowedGroups(context.Context, int64, int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ==================== 辅助函数：构造带 CDN 缓存的 SoraS3Storage ====================
|
||||
|
||||
// newS3StorageWithCDN 创建一个预缓存了 CDN 配置的 SoraS3Storage，
|
||||
// 避免实际初始化 AWS 客户端。用于测试 GetAccessURL 的 CDN 路径。
|
||||
func newS3StorageWithCDN(cdnURL string) *SoraS3Storage {
|
||||
storage := &SoraS3Storage{}
|
||||
storage.cfg = &SoraS3Settings{
|
||||
Enabled: true,
|
||||
Bucket: "test-bucket",
|
||||
CDNURL: cdnURL,
|
||||
}
|
||||
// 需要 non-nil client 使 getClient 命中缓存
|
||||
storage.client = s3.New(s3.Options{})
|
||||
return storage
|
||||
}
|
||||
|
||||
// newS3StorageFailingDelete 创建一个 settingService=nil 的 SoraS3Storage，
|
||||
// 使 DeleteObjects 返回错误（无法获取配置）。用于测试 Delete 方法 S3 清理失败但仍继续的场景。
|
||||
func newS3StorageFailingDelete() *SoraS3Storage {
|
||||
return &SoraS3Storage{} // settingService 为 nil → getConfig 返回 error
|
||||
}
|
||||
|
||||
// ==================== CreatePending ====================
|
||||
|
||||
func TestCreatePending_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, err := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "一只猫跳舞", "video")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), gen.ID)
|
||||
require.Equal(t, int64(1), gen.UserID)
|
||||
require.Equal(t, "sora2-landscape-10s", gen.Model)
|
||||
require.Equal(t, "一只猫跳舞", gen.Prompt)
|
||||
require.Equal(t, "video", gen.MediaType)
|
||||
require.Equal(t, SoraGenStatusPending, gen.Status)
|
||||
require.Equal(t, SoraStorageTypeNone, gen.StorageType)
|
||||
require.Nil(t, gen.APIKeyID)
|
||||
}
|
||||
|
||||
func TestCreatePending_WithAPIKeyID(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
apiKeyID := int64(42)
|
||||
gen, err := svc.CreatePending(context.Background(), 1, &apiKeyID, "gpt-image", "画一朵花", "image")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, gen.APIKeyID)
|
||||
require.Equal(t, int64(42), *gen.APIKeyID)
|
||||
}
|
||||
|
||||
func TestCreatePending_RepoError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.createErr = fmt.Errorf("db write error")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, err := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
require.Error(t, err)
|
||||
require.Nil(t, gen)
|
||||
require.Contains(t, err.Error(), "create generation")
|
||||
}
|
||||
|
||||
// ==================== MarkGenerating ====================
|
||||
|
||||
func TestMarkGenerating_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusPending}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkGenerating(context.Background(), 1, "upstream-task-123")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusGenerating, repo.gens[1].Status)
|
||||
require.Equal(t, "upstream-task-123", repo.gens[1].UpstreamTaskID)
|
||||
}
|
||||
|
||||
func TestMarkGenerating_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkGenerating(context.Background(), 999, "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkGenerating_UpdateError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusPending}
|
||||
repo.updateErr = fmt.Errorf("update failed")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkGenerating(context.Background(), 1, "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== MarkCompleted ====================
|
||||
|
||||
func TestMarkCompleted_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCompleted(context.Background(), 1,
|
||||
"https://cdn.example.com/video.mp4",
|
||||
[]string{"https://cdn.example.com/video.mp4"},
|
||||
SoraStorageTypeS3,
|
||||
[]string{"sora/1/2024/01/01/uuid.mp4"},
|
||||
1048576,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
gen := repo.gens[1]
|
||||
require.Equal(t, SoraGenStatusCompleted, gen.Status)
|
||||
require.Equal(t, "https://cdn.example.com/video.mp4", gen.MediaURL)
|
||||
require.Equal(t, []string{"https://cdn.example.com/video.mp4"}, gen.MediaURLs)
|
||||
require.Equal(t, SoraStorageTypeS3, gen.StorageType)
|
||||
require.Equal(t, []string{"sora/1/2024/01/01/uuid.mp4"}, gen.S3ObjectKeys)
|
||||
require.Equal(t, int64(1048576), gen.FileSizeBytes)
|
||||
require.NotNil(t, gen.CompletedAt)
|
||||
}
|
||||
|
||||
func TestMarkCompleted_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCompleted(context.Background(), 999, "", nil, "", nil, 0)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkCompleted_UpdateError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
repo.updateErr = fmt.Errorf("update failed")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCompleted(context.Background(), 1, "url", nil, SoraStorageTypeUpstream, nil, 0)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== MarkFailed ====================
|
||||
|
||||
func TestMarkFailed_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkFailed(context.Background(), 1, "上游返回 500 错误")
|
||||
require.NoError(t, err)
|
||||
gen := repo.gens[1]
|
||||
require.Equal(t, SoraGenStatusFailed, gen.Status)
|
||||
require.Equal(t, "上游返回 500 错误", gen.ErrorMessage)
|
||||
require.NotNil(t, gen.CompletedAt)
|
||||
}
|
||||
|
||||
func TestMarkFailed_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkFailed(context.Background(), 999, "error")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkFailed_UpdateError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
repo.updateErr = fmt.Errorf("update failed")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkFailed(context.Background(), 1, "err")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== MarkCancelled ====================
|
||||
|
||||
func TestMarkCancelled_Pending(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusPending}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusCancelled, repo.gens[1].Status)
|
||||
require.NotNil(t, repo.gens[1].CompletedAt)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_Generating(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusCancelled, repo.gens[1].Status)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_Completed(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCompleted}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.Error(t, err)
|
||||
require.ErrorIs(t, err, ErrSoraGenerationNotActive)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_Failed(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusFailed}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_AlreadyCancelled(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCancelled}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 999)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestMarkCancelled_UpdateError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusPending}
|
||||
repo.updateErr = fmt.Errorf("update failed")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.MarkCancelled(context.Background(), 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== GetByID ====================
|
||||
|
||||
func TestGetByID_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCompleted, Model: "sora2-landscape-10s"}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, err := svc.GetByID(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), gen.ID)
|
||||
require.Equal(t, "sora2-landscape-10s", gen.Model)
|
||||
}
|
||||
|
||||
func TestGetByID_WrongUser(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 2, Status: SoraGenStatusCompleted}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, err := svc.GetByID(context.Background(), 1, 1)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, gen)
|
||||
require.Contains(t, err.Error(), "无权访问")
|
||||
}
|
||||
|
||||
func TestGetByID_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, err := svc.GetByID(context.Background(), 999, 1)
|
||||
require.Error(t, err)
|
||||
require.Nil(t, gen)
|
||||
}
|
||||
|
||||
// ==================== List ====================
|
||||
|
||||
func TestList_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCompleted, MediaType: "video"}
|
||||
repo.gens[2] = &SoraGeneration{ID: 2, UserID: 1, Status: SoraGenStatusPending, MediaType: "image"}
|
||||
repo.gens[3] = &SoraGeneration{ID: 3, UserID: 2, Status: SoraGenStatusCompleted, MediaType: "video"}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gens, total, err := svc.List(context.Background(), SoraGenerationListParams{UserID: 1, Page: 1, PageSize: 20})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, gens, 2) // 只有 userID=1 的
|
||||
require.Equal(t, int64(2), total)
|
||||
}
|
||||
|
||||
func TestList_DefaultPagination(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
// page=0, pageSize=0 → 应修正为 page=1, pageSize=20
|
||||
_, _, err := svc.List(context.Background(), SoraGenerationListParams{UserID: 1})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestList_MaxPageSize(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
// pageSize > 100 → 应限制为 100
|
||||
_, _, err := svc.List(context.Background(), SoraGenerationListParams{UserID: 1, Page: 1, PageSize: 200})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestList_Error(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.listErr = fmt.Errorf("db error")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
_, _, err := svc.List(context.Background(), SoraGenerationListParams{UserID: 1})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== Delete ====================
|
||||
|
||||
func TestDelete_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCompleted, StorageType: SoraStorageTypeUpstream}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDelete_WrongUser(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 2, Status: SoraGenStatusCompleted}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "无权删除")
|
||||
}
|
||||
|
||||
func TestDelete_NotFound(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 999, 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestDelete_S3Cleanup_NilS3(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, StorageType: SoraStorageTypeS3, S3ObjectKeys: []string{"key1"}}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err) // s3Storage 为 nil，跳过清理
|
||||
}
|
||||
|
||||
func TestDelete_QuotaRelease_NilQuota(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, StorageType: SoraStorageTypeS3, FileSizeBytes: 1024}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err) // quotaService 为 nil，跳过释放
|
||||
}
|
||||
|
||||
func TestDelete_NonS3NoCleanup(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, StorageType: SoraStorageTypeLocal, FileSizeBytes: 1024}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestDelete_DeleteRepoError(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, StorageType: SoraStorageTypeUpstream}
|
||||
repo.deleteErr = fmt.Errorf("delete failed")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== CountActiveByUser ====================
|
||||
|
||||
func TestCountActiveByUser_Success(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusPending}
|
||||
repo.gens[2] = &SoraGeneration{ID: 2, UserID: 1, Status: SoraGenStatusGenerating}
|
||||
repo.gens[3] = &SoraGeneration{ID: 3, UserID: 1, Status: SoraGenStatusCompleted} // 不算
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
count, err := svc.CountActiveByUser(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), count)
|
||||
}
|
||||
|
||||
func TestCountActiveByUser_NoActive(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{ID: 1, UserID: 1, Status: SoraGenStatusCompleted}
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
count, err := svc.CountActiveByUser(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), count)
|
||||
}
|
||||
|
||||
func TestCountActiveByUser_Error(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
repo.countErr = fmt.Errorf("db error")
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
_, err := svc.CountActiveByUser(context.Background(), 1)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== ResolveMediaURLs ====================
|
||||
|
||||
func TestResolveMediaURLs_NilGen(t *testing.T) {
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), nil, nil)
|
||||
require.NoError(t, svc.ResolveMediaURLs(context.Background(), nil))
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_NonS3(t *testing.T) {
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), nil, nil)
|
||||
gen := &SoraGeneration{StorageType: SoraStorageTypeUpstream, MediaURL: "https://original.com/v.mp4"}
|
||||
require.NoError(t, svc.ResolveMediaURLs(context.Background(), gen))
|
||||
require.Equal(t, "https://original.com/v.mp4", gen.MediaURL) // 不变
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_S3NilStorage(t *testing.T) {
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), nil, nil)
|
||||
gen := &SoraGeneration{StorageType: SoraStorageTypeS3, S3ObjectKeys: []string{"key1"}}
|
||||
require.NoError(t, svc.ResolveMediaURLs(context.Background(), gen))
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_Local(t *testing.T) {
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), nil, nil)
|
||||
gen := &SoraGeneration{StorageType: SoraStorageTypeLocal, MediaURL: "/video/2024/01/01/file.mp4"}
|
||||
require.NoError(t, svc.ResolveMediaURLs(context.Background(), gen))
|
||||
require.Equal(t, "/video/2024/01/01/file.mp4", gen.MediaURL) // 不变
|
||||
}
|
||||
|
||||
// ==================== 状态流转完整测试 ====================
|
||||
|
||||
func TestStatusTransition_PendingToCompletedFlow(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
// 1. 创建 pending
|
||||
gen, err := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusPending, gen.Status)
|
||||
|
||||
// 2. 标记 generating
|
||||
err = svc.MarkGenerating(context.Background(), gen.ID, "task-123")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusGenerating, repo.gens[gen.ID].Status)
|
||||
|
||||
// 3. 标记 completed
|
||||
err = svc.MarkCompleted(context.Background(), gen.ID, "https://s3.com/video.mp4", nil, SoraStorageTypeS3, []string{"key"}, 1024)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusCompleted, repo.gens[gen.ID].Status)
|
||||
}
|
||||
|
||||
func TestStatusTransition_PendingToFailedFlow(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, _ := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
_ = svc.MarkGenerating(context.Background(), gen.ID, "")
|
||||
|
||||
err := svc.MarkFailed(context.Background(), gen.ID, "上游超时")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusFailed, repo.gens[gen.ID].Status)
|
||||
require.Equal(t, "上游超时", repo.gens[gen.ID].ErrorMessage)
|
||||
}
|
||||
|
||||
func TestStatusTransition_PendingToCancelledFlow(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, _ := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
err := svc.MarkCancelled(context.Background(), gen.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusCancelled, repo.gens[gen.ID].Status)
|
||||
}
|
||||
|
||||
func TestStatusTransition_GeneratingToCancelledFlow(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, _ := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
_ = svc.MarkGenerating(context.Background(), gen.ID, "")
|
||||
err := svc.MarkCancelled(context.Background(), gen.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, SoraGenStatusCancelled, repo.gens[gen.ID].Status)
|
||||
}
|
||||
|
||||
// ==================== 权限隔离测试 ====================
|
||||
|
||||
func TestUserIsolation_CannotAccessOthersRecord(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, _ := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
|
||||
// 用户 2 尝试访问用户 1 的记录
|
||||
_, err := svc.GetByID(context.Background(), gen.ID, 2)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "无权访问")
|
||||
}
|
||||
|
||||
func TestUserIsolation_CannotDeleteOthersRecord(t *testing.T) {
|
||||
repo := newStubGenRepo()
|
||||
svc := NewSoraGenerationService(repo, nil, nil)
|
||||
|
||||
gen, _ := svc.CreatePending(context.Background(), 1, nil, "sora2-landscape-10s", "test", "video")
|
||||
|
||||
err := svc.Delete(context.Background(), gen.ID, 2)
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "无权删除")
|
||||
}
|
||||
|
||||
// ==================== Delete: S3 清理 + 配额释放路径 ====================
|
||||
|
||||
func TestDelete_S3Cleanup_WithS3Storage(t *testing.T) {
|
||||
// S3 存储存在但 deleteObjects 会失败（settingService=nil），
|
||||
// 验证 Delete 仍然成功（S3 错误只是记录日志）
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{
|
||||
ID: 1, UserID: 1,
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{"sora/1/2024/01/01/abc.mp4"},
|
||||
}
|
||||
s3Storage := newS3StorageFailingDelete()
|
||||
svc := NewSoraGenerationService(repo, s3Storage, nil)
|
||||
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err) // S3 清理失败不影响删除
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDelete_QuotaRelease_WithQuotaService(t *testing.T) {
|
||||
// 配额服务存在时，删除记录会调用 ReleaseUsage（仅记录日志，不再追踪用量）
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{
|
||||
ID: 1, UserID: 1,
|
||||
StorageType: SoraStorageTypeS3,
|
||||
FileSizeBytes: 1048576, // 1MB
|
||||
}
|
||||
|
||||
quotaService := NewSoraQuotaService(nil)
|
||||
|
||||
svc := NewSoraGenerationService(repo, nil, quotaService)
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
// 删除成功即可，用量追踪已移除
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDelete_S3Cleanup_And_QuotaRelease(t *testing.T) {
|
||||
// S3 清理 + 配额释放日志同时触发
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{
|
||||
ID: 1, UserID: 1,
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{"key1"},
|
||||
FileSizeBytes: 512,
|
||||
}
|
||||
|
||||
quotaService := NewSoraQuotaService(nil)
|
||||
s3Storage := newS3StorageFailingDelete()
|
||||
|
||||
svc := NewSoraGenerationService(repo, s3Storage, quotaService)
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDelete_QuotaRelease_LocalStorage(t *testing.T) {
|
||||
// 本地存储同样会记录释放日志
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{
|
||||
ID: 1, UserID: 1,
|
||||
StorageType: SoraStorageTypeLocal,
|
||||
FileSizeBytes: 1024,
|
||||
}
|
||||
|
||||
quotaService := NewSoraQuotaService(nil)
|
||||
|
||||
svc := NewSoraGenerationService(repo, nil, quotaService)
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
// 删除成功即可，用量追踪已移除
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
func TestDelete_QuotaRelease_ZeroFileSize(t *testing.T) {
|
||||
// FileSizeBytes=0 仍然会删除记录（用量追踪已移除）
|
||||
repo := newStubGenRepo()
|
||||
repo.gens[1] = &SoraGeneration{
|
||||
ID: 1, UserID: 1,
|
||||
StorageType: SoraStorageTypeS3,
|
||||
FileSizeBytes: 0,
|
||||
}
|
||||
|
||||
quotaService := NewSoraQuotaService(nil)
|
||||
|
||||
svc := NewSoraGenerationService(repo, nil, quotaService)
|
||||
err := svc.Delete(context.Background(), 1, 1)
|
||||
require.NoError(t, err)
|
||||
_, exists := repo.gens[1]
|
||||
require.False(t, exists)
|
||||
}
|
||||
|
||||
// ==================== ResolveMediaURLs: S3 + CDN 路径 ====================
|
||||
|
||||
func TestResolveMediaURLs_S3_CDN_SingleKey(t *testing.T) {
|
||||
s3Storage := newS3StorageWithCDN("https://cdn.example.com")
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), s3Storage, nil)
|
||||
|
||||
gen := &SoraGeneration{
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{"sora/1/2024/01/01/video.mp4"},
|
||||
MediaURL: "original",
|
||||
}
|
||||
err := svc.ResolveMediaURLs(context.Background(), gen)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/video.mp4", gen.MediaURL)
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_S3_CDN_MultipleKeys(t *testing.T) {
|
||||
s3Storage := newS3StorageWithCDN("https://cdn.example.com/")
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), s3Storage, nil)
|
||||
|
||||
gen := &SoraGeneration{
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{
|
||||
"sora/1/2024/01/01/img1.png",
|
||||
"sora/1/2024/01/01/img2.png",
|
||||
"sora/1/2024/01/01/img3.png",
|
||||
},
|
||||
MediaURL: "original",
|
||||
}
|
||||
err := svc.ResolveMediaURLs(context.Background(), gen)
|
||||
require.NoError(t, err)
|
||||
// 主 URL 更新为第一个 key 的 CDN URL
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/img1.png", gen.MediaURL)
|
||||
// 多图 URLs 全部更新
|
||||
require.Len(t, gen.MediaURLs, 3)
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/img1.png", gen.MediaURLs[0])
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/img2.png", gen.MediaURLs[1])
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/img3.png", gen.MediaURLs[2])
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_S3_EmptyKeys(t *testing.T) {
|
||||
s3Storage := newS3StorageWithCDN("https://cdn.example.com")
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), s3Storage, nil)
|
||||
|
||||
gen := &SoraGeneration{
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{},
|
||||
MediaURL: "original",
|
||||
}
|
||||
err := svc.ResolveMediaURLs(context.Background(), gen)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "original", gen.MediaURL) // 不变
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_S3_GetAccessURL_Error(t *testing.T) {
|
||||
// 使用无 settingService 的 S3 Storage，getClient 会失败
|
||||
s3Storage := newS3StorageFailingDelete() // 同样 GetAccessURL 也会失败
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), s3Storage, nil)
|
||||
|
||||
gen := &SoraGeneration{
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{"sora/1/2024/01/01/video.mp4"},
|
||||
MediaURL: "original",
|
||||
}
|
||||
err := svc.ResolveMediaURLs(context.Background(), gen)
|
||||
require.Error(t, err) // GetAccessURL 失败应传播错误
|
||||
}
|
||||
|
||||
func TestResolveMediaURLs_S3_MultiKey_ErrorOnSecond(t *testing.T) {
|
||||
// 只有一个 key 时走主 URL 路径成功，但多 key 路径的错误也需覆盖
|
||||
s3Storage := newS3StorageFailingDelete()
|
||||
svc := NewSoraGenerationService(newStubGenRepo(), s3Storage, nil)
|
||||
|
||||
gen := &SoraGeneration{
|
||||
StorageType: SoraStorageTypeS3,
|
||||
S3ObjectKeys: []string{
|
||||
"sora/1/2024/01/01/img1.png",
|
||||
"sora/1/2024/01/01/img2.png",
|
||||
},
|
||||
MediaURL: "original",
|
||||
}
|
||||
err := svc.ResolveMediaURLs(context.Background(), gen)
|
||||
require.Error(t, err) // 第一个 key 的 GetAccessURL 就会失败
|
||||
}
|
||||
@@ -1,120 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/robfig/cron/v3"
|
||||
)
|
||||
|
||||
var soraCleanupCronParser = cron.NewParser(cron.Minute | cron.Hour | cron.Dom | cron.Month | cron.Dow)
|
||||
|
||||
// SoraMediaCleanupService 定期清理本地媒体文件
|
||||
type SoraMediaCleanupService struct {
|
||||
storage *SoraMediaStorage
|
||||
cfg *config.Config
|
||||
|
||||
cron *cron.Cron
|
||||
|
||||
startOnce sync.Once
|
||||
stopOnce sync.Once
|
||||
}
|
||||
|
||||
func NewSoraMediaCleanupService(storage *SoraMediaStorage, cfg *config.Config) *SoraMediaCleanupService {
|
||||
return &SoraMediaCleanupService{
|
||||
storage: storage,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *SoraMediaCleanupService) Start() {
|
||||
if s == nil || s.cfg == nil {
|
||||
return
|
||||
}
|
||||
if !s.cfg.Sora.Storage.Cleanup.Enabled {
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] not started (disabled)")
|
||||
return
|
||||
}
|
||||
if s.storage == nil || !s.storage.Enabled() {
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] not started (storage disabled)")
|
||||
return
|
||||
}
|
||||
|
||||
s.startOnce.Do(func() {
|
||||
schedule := strings.TrimSpace(s.cfg.Sora.Storage.Cleanup.Schedule)
|
||||
if schedule == "" {
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] not started (empty schedule)")
|
||||
return
|
||||
}
|
||||
loc := time.Local
|
||||
if strings.TrimSpace(s.cfg.Timezone) != "" {
|
||||
if parsed, err := time.LoadLocation(strings.TrimSpace(s.cfg.Timezone)); err == nil && parsed != nil {
|
||||
loc = parsed
|
||||
}
|
||||
}
|
||||
c := cron.New(cron.WithParser(soraCleanupCronParser), cron.WithLocation(loc))
|
||||
if _, err := c.AddFunc(schedule, func() { s.runCleanup() }); err != nil {
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] not started (invalid schedule=%q): %v", schedule, err)
|
||||
return
|
||||
}
|
||||
s.cron = c
|
||||
s.cron.Start()
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] started (schedule=%q tz=%s)", schedule, loc.String())
|
||||
})
|
||||
}
|
||||
|
||||
func (s *SoraMediaCleanupService) Stop() {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.stopOnce.Do(func() {
|
||||
if s.cron != nil {
|
||||
ctx := s.cron.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-time.After(3 * time.Second):
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] cron stop timed out")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (s *SoraMediaCleanupService) runCleanup() {
|
||||
if s.cfg == nil || s.storage == nil {
|
||||
return
|
||||
}
|
||||
retention := s.cfg.Sora.Storage.Cleanup.RetentionDays
|
||||
if retention <= 0 {
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] skipped (retention_days=%d)", retention)
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().AddDate(0, 0, -retention)
|
||||
deleted := 0
|
||||
|
||||
roots := []string{s.storage.ImageRoot(), s.storage.VideoRoot()}
|
||||
for _, root := range roots {
|
||||
if root == "" {
|
||||
continue
|
||||
}
|
||||
_ = filepath.Walk(root, func(p string, info os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if info.IsDir() {
|
||||
return nil
|
||||
}
|
||||
if info.ModTime().Before(cutoff) {
|
||||
if rmErr := os.Remove(p); rmErr == nil {
|
||||
deleted++
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
logger.LegacyPrintf("service.sora_media_cleanup", "[SoraCleanup] cleanup finished, deleted=%d", deleted)
|
||||
}
|
||||
@@ -1,207 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSoraMediaCleanupService_RunCleanup_NilCfg(t *testing.T) {
|
||||
storage := &SoraMediaStorage{}
|
||||
svc := &SoraMediaCleanupService{storage: storage, cfg: nil}
|
||||
// 不应 panic
|
||||
svc.runCleanup()
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_RunCleanup_NilStorage(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
svc := &SoraMediaCleanupService{storage: nil, cfg: cfg}
|
||||
// 不应 panic
|
||||
svc.runCleanup()
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_RunCleanup_ZeroRetention(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
RetentionDays: 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
// retention=0 应跳过清理
|
||||
svc.runCleanup()
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_NilCfg(t *testing.T) {
|
||||
svc := NewSoraMediaCleanupService(nil, nil)
|
||||
svc.Start() // cfg == nil 时应直接返回
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_StorageDisabled(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraMediaCleanupService(nil, cfg)
|
||||
svc.Start() // storage == nil 时应直接返回
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_WithTimezone(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Timezone: "Asia/Shanghai",
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
Schedule: "0 3 * * *",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
svc.Start()
|
||||
t.Cleanup(svc.Stop)
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_Disabled(t *testing.T) {
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: false,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
svc := NewSoraMediaCleanupService(nil, cfg)
|
||||
svc.Start() // 不应 panic，也不应启动 cron
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_NilSelf(t *testing.T) {
|
||||
var svc *SoraMediaCleanupService
|
||||
svc.Start() // 不应 panic
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_EmptySchedule(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
Schedule: "",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
svc.Start() // 空 schedule 不应启动
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_InvalidSchedule(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
Schedule: "invalid-cron",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
svc.Start() // 无效 schedule 不应 panic
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Start_ValidSchedule(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
Schedule: "0 3 * * *",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
svc.Start()
|
||||
t.Cleanup(svc.Stop)
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Stop_NilSelf(t *testing.T) {
|
||||
var svc *SoraMediaCleanupService
|
||||
svc.Stop() // 不应 panic
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_Stop_WithoutStart(t *testing.T) {
|
||||
svc := NewSoraMediaCleanupService(nil, &config.Config{})
|
||||
svc.Stop() // cron 未启动时 Stop 不应 panic
|
||||
}
|
||||
|
||||
func TestSoraMediaCleanupService_RunCleanup(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
Cleanup: config.SoraStorageCleanupConfig{
|
||||
Enabled: true,
|
||||
RetentionDays: 1,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
require.NoError(t, storage.EnsureLocalDirs())
|
||||
|
||||
oldImage := filepath.Join(storage.ImageRoot(), "old.png")
|
||||
newVideo := filepath.Join(storage.VideoRoot(), "new.mp4")
|
||||
require.NoError(t, os.WriteFile(oldImage, []byte("old"), 0o644))
|
||||
require.NoError(t, os.WriteFile(newVideo, []byte("new"), 0o644))
|
||||
|
||||
oldTime := time.Now().Add(-48 * time.Hour)
|
||||
require.NoError(t, os.Chtimes(oldImage, oldTime, oldTime))
|
||||
|
||||
cleanup := NewSoraMediaCleanupService(storage, cfg)
|
||||
cleanup.runCleanup()
|
||||
|
||||
require.NoFileExists(t, oldImage)
|
||||
require.FileExists(t, newVideo)
|
||||
}
|
||||
@@ -1,48 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// SignSoraMediaURL 生成 Sora 媒体临时签名
|
||||
func SignSoraMediaURL(path string, query string, expires int64, key string) string {
|
||||
key = strings.TrimSpace(key)
|
||||
if key == "" {
|
||||
return ""
|
||||
}
|
||||
mac := hmac.New(sha256.New, []byte(key))
|
||||
if _, err := mac.Write([]byte(buildSoraMediaSignPayload(path, query))); err != nil {
|
||||
return ""
|
||||
}
|
||||
if _, err := mac.Write([]byte("|")); err != nil {
|
||||
return ""
|
||||
}
|
||||
if _, err := mac.Write([]byte(strconv.FormatInt(expires, 10))); err != nil {
|
||||
return ""
|
||||
}
|
||||
return hex.EncodeToString(mac.Sum(nil))
|
||||
}
|
||||
|
||||
// VerifySoraMediaURL 校验 Sora 媒体签名
|
||||
func VerifySoraMediaURL(path string, query string, expires int64, signature string, key string) bool {
|
||||
signature = strings.TrimSpace(signature)
|
||||
if signature == "" {
|
||||
return false
|
||||
}
|
||||
expected := SignSoraMediaURL(path, query, expires, key)
|
||||
if expected == "" {
|
||||
return false
|
||||
}
|
||||
return hmac.Equal([]byte(signature), []byte(expected))
|
||||
}
|
||||
|
||||
func buildSoraMediaSignPayload(path string, query string) string {
|
||||
if strings.TrimSpace(query) == "" {
|
||||
return path
|
||||
}
|
||||
return path + "?" + query
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
package service
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestSoraMediaSignVerify(t *testing.T) {
|
||||
key := "test-key"
|
||||
path := "/tmp/abc.png"
|
||||
query := "a=1&b=2"
|
||||
expires := int64(1700000000)
|
||||
|
||||
signature := SignSoraMediaURL(path, query, expires, key)
|
||||
if signature == "" {
|
||||
t.Fatal("签名为空")
|
||||
}
|
||||
if !VerifySoraMediaURL(path, query, expires, signature, key) {
|
||||
t.Fatal("签名校验失败")
|
||||
}
|
||||
if VerifySoraMediaURL(path, "a=1", expires, signature, key) {
|
||||
t.Fatal("签名参数不同仍然通过")
|
||||
}
|
||||
if VerifySoraMediaURL(path, query, expires+1, signature, key) {
|
||||
t.Fatal("签名过期校验未失败")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSoraMediaSignWithEmptyKey(t *testing.T) {
|
||||
signature := SignSoraMediaURL("/tmp/a.png", "a=1", 1, "")
|
||||
if signature != "" {
|
||||
t.Fatalf("空密钥不应生成签名")
|
||||
}
|
||||
if VerifySoraMediaURL("/tmp/a.png", "a=1", 1, "sig", "") {
|
||||
t.Fatalf("空密钥不应通过校验")
|
||||
}
|
||||
}
|
||||
@@ -1,381 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"mime"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
soraStorageDefaultRoot = "/app/data/sora"
|
||||
)
|
||||
|
||||
// SoraMediaStorage 负责下载并落地 Sora 媒体
|
||||
type SoraMediaStorage struct {
|
||||
cfg *config.Config
|
||||
root string
|
||||
imageRoot string
|
||||
videoRoot string
|
||||
downloadTimeout time.Duration
|
||||
maxDownloadBytes int64
|
||||
fallbackToUpstream bool
|
||||
debug bool
|
||||
sem chan struct{}
|
||||
ready bool
|
||||
}
|
||||
|
||||
func NewSoraMediaStorage(cfg *config.Config) *SoraMediaStorage {
|
||||
storage := &SoraMediaStorage{cfg: cfg}
|
||||
storage.refreshConfig()
|
||||
if storage.Enabled() {
|
||||
if err := storage.EnsureLocalDirs(); err != nil {
|
||||
log.Printf("[SoraStorage] 初始化失败: %v", err)
|
||||
}
|
||||
}
|
||||
return storage
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) Enabled() bool {
|
||||
if s == nil || s.cfg == nil {
|
||||
return false
|
||||
}
|
||||
return strings.ToLower(strings.TrimSpace(s.cfg.Sora.Storage.Type)) == "local"
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) Root() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
return s.root
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) ImageRoot() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
return s.imageRoot
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) VideoRoot() string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
return s.videoRoot
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) refreshConfig() {
|
||||
if s == nil || s.cfg == nil {
|
||||
return
|
||||
}
|
||||
root := strings.TrimSpace(s.cfg.Sora.Storage.LocalPath)
|
||||
if root == "" {
|
||||
root = soraStorageDefaultRoot
|
||||
}
|
||||
root = filepath.Clean(root)
|
||||
if !filepath.IsAbs(root) {
|
||||
if absRoot, err := filepath.Abs(root); err == nil {
|
||||
root = absRoot
|
||||
}
|
||||
}
|
||||
s.root = root
|
||||
s.imageRoot = filepath.Join(root, "image")
|
||||
s.videoRoot = filepath.Join(root, "video")
|
||||
|
||||
maxConcurrent := s.cfg.Sora.Storage.MaxConcurrentDownloads
|
||||
if maxConcurrent <= 0 {
|
||||
maxConcurrent = 4
|
||||
}
|
||||
timeoutSeconds := s.cfg.Sora.Storage.DownloadTimeoutSeconds
|
||||
if timeoutSeconds <= 0 {
|
||||
timeoutSeconds = 120
|
||||
}
|
||||
s.downloadTimeout = time.Duration(timeoutSeconds) * time.Second
|
||||
|
||||
maxBytes := s.cfg.Sora.Storage.MaxDownloadBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 200 << 20
|
||||
}
|
||||
s.maxDownloadBytes = maxBytes
|
||||
s.fallbackToUpstream = s.cfg.Sora.Storage.FallbackToUpstream
|
||||
s.debug = s.cfg.Sora.Storage.Debug
|
||||
s.sem = make(chan struct{}, maxConcurrent)
|
||||
}
|
||||
|
||||
// EnsureLocalDirs 创建并校验本地目录
|
||||
func (s *SoraMediaStorage) EnsureLocalDirs() error {
|
||||
if s == nil || !s.Enabled() {
|
||||
return nil
|
||||
}
|
||||
if err := os.MkdirAll(s.imageRoot, 0o755); err != nil {
|
||||
return fmt.Errorf("create image dir: %w", err)
|
||||
}
|
||||
if err := os.MkdirAll(s.videoRoot, 0o755); err != nil {
|
||||
return fmt.Errorf("create video dir: %w", err)
|
||||
}
|
||||
s.ready = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// StoreFromURLs 下载并存储媒体，返回相对路径或回退 URL
|
||||
func (s *SoraMediaStorage) StoreFromURLs(ctx context.Context, mediaType string, urls []string) ([]string, error) {
|
||||
if len(urls) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if s == nil || !s.Enabled() {
|
||||
return urls, nil
|
||||
}
|
||||
if !s.ready {
|
||||
if err := s.EnsureLocalDirs(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
results := make([]string, 0, len(urls))
|
||||
for _, raw := range urls {
|
||||
relative, err := s.downloadAndStore(ctx, mediaType, raw)
|
||||
if err != nil {
|
||||
if s.fallbackToUpstream {
|
||||
results = append(results, raw)
|
||||
continue
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, relative)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
// TotalSizeByRelativePaths 统计本地存储路径总大小（仅统计 /image 和 /video 路径）。
|
||||
func (s *SoraMediaStorage) TotalSizeByRelativePaths(paths []string) (int64, error) {
|
||||
if s == nil || len(paths) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
var total int64
|
||||
for _, p := range paths {
|
||||
localPath, err := s.resolveLocalPath(p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
info, err := os.Stat(localPath)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
continue
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
if info.Mode().IsRegular() {
|
||||
total += info.Size()
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
// DeleteByRelativePaths 删除本地媒体路径（仅删除 /image 和 /video 路径）。
|
||||
func (s *SoraMediaStorage) DeleteByRelativePaths(paths []string) error {
|
||||
if s == nil || len(paths) == 0 {
|
||||
return nil
|
||||
}
|
||||
var lastErr error
|
||||
for _, p := range paths {
|
||||
localPath, err := s.resolveLocalPath(p)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if err := os.Remove(localPath); err != nil && !os.IsNotExist(err) {
|
||||
lastErr = err
|
||||
}
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) resolveLocalPath(relativePath string) (string, error) {
|
||||
if s == nil || strings.TrimSpace(relativePath) == "" {
|
||||
return "", errors.New("empty path")
|
||||
}
|
||||
cleaned := path.Clean(relativePath)
|
||||
if !strings.HasPrefix(cleaned, "/image/") && !strings.HasPrefix(cleaned, "/video/") {
|
||||
return "", errors.New("not a local media path")
|
||||
}
|
||||
if strings.TrimSpace(s.root) == "" {
|
||||
return "", errors.New("storage root not configured")
|
||||
}
|
||||
relative := strings.TrimPrefix(cleaned, "/")
|
||||
return filepath.Join(s.root, filepath.FromSlash(relative)), nil
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) downloadAndStore(ctx context.Context, mediaType, rawURL string) (string, error) {
|
||||
if strings.TrimSpace(rawURL) == "" {
|
||||
return "", errors.New("empty url")
|
||||
}
|
||||
root := s.imageRoot
|
||||
if mediaType == "video" {
|
||||
root = s.videoRoot
|
||||
}
|
||||
if root == "" {
|
||||
return "", errors.New("storage root not configured")
|
||||
}
|
||||
|
||||
retries := 3
|
||||
for attempt := 1; attempt <= retries; attempt++ {
|
||||
release, err := s.acquire(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
relative, err := s.downloadOnce(ctx, root, mediaType, rawURL)
|
||||
release()
|
||||
if err == nil {
|
||||
return relative, nil
|
||||
}
|
||||
if s.debug {
|
||||
log.Printf("[SoraStorage] 下载失败(%d/%d): %s err=%v", attempt, retries, sanitizeMediaLogURL(rawURL), err)
|
||||
}
|
||||
if attempt < retries {
|
||||
time.Sleep(time.Duration(attempt*attempt) * time.Second)
|
||||
continue
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
return "", errors.New("download retries exhausted")
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) downloadOnce(ctx context.Context, root, mediaType, rawURL string) (string, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
client := &http.Client{Timeout: s.downloadTimeout}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
|
||||
return "", fmt.Errorf("download failed: %d %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
ext := normalizeSoraFileExt(fileExtFromURL(rawURL))
|
||||
if ext == "" {
|
||||
ext = normalizeSoraFileExt(fileExtFromContentType(resp.Header.Get("Content-Type")))
|
||||
}
|
||||
if ext == "" {
|
||||
ext = ".bin"
|
||||
}
|
||||
if s.maxDownloadBytes > 0 && resp.ContentLength > s.maxDownloadBytes {
|
||||
return "", fmt.Errorf("download size exceeds limit: %d", resp.ContentLength)
|
||||
}
|
||||
|
||||
storageRoot, err := os.OpenRoot(root)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { _ = storageRoot.Close() }()
|
||||
|
||||
datePath := time.Now().Format("2006/01/02")
|
||||
datePathFS := filepath.FromSlash(datePath)
|
||||
if err := storageRoot.MkdirAll(datePathFS, 0o755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
filename := uuid.NewString() + ext
|
||||
filePath := filepath.Join(datePathFS, filename)
|
||||
out, err := storageRoot.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer func() { _ = out.Close() }()
|
||||
|
||||
limited := io.LimitReader(resp.Body, s.maxDownloadBytes+1)
|
||||
written, err := io.Copy(out, limited)
|
||||
if err != nil {
|
||||
removePartialDownload(storageRoot, filePath)
|
||||
return "", err
|
||||
}
|
||||
if s.maxDownloadBytes > 0 && written > s.maxDownloadBytes {
|
||||
removePartialDownload(storageRoot, filePath)
|
||||
return "", fmt.Errorf("download size exceeds limit: %d", written)
|
||||
}
|
||||
|
||||
relative := path.Join("/", mediaType, datePath, filename)
|
||||
if s.debug {
|
||||
log.Printf("[SoraStorage] 已落地 %s -> %s", sanitizeMediaLogURL(rawURL), relative)
|
||||
}
|
||||
return relative, nil
|
||||
}
|
||||
|
||||
func (s *SoraMediaStorage) acquire(ctx context.Context) (func(), error) {
|
||||
if s.sem == nil {
|
||||
return func() {}, nil
|
||||
}
|
||||
select {
|
||||
case s.sem <- struct{}{}:
|
||||
return func() { <-s.sem }, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func fileExtFromURL(raw string) string {
|
||||
parsed, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
ext := path.Ext(parsed.Path)
|
||||
return strings.ToLower(ext)
|
||||
}
|
||||
|
||||
func fileExtFromContentType(ct string) string {
|
||||
if ct == "" {
|
||||
return ""
|
||||
}
|
||||
if exts, err := mime.ExtensionsByType(ct); err == nil && len(exts) > 0 {
|
||||
return strings.ToLower(exts[0])
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeSoraFileExt(ext string) string {
|
||||
ext = strings.ToLower(strings.TrimSpace(ext))
|
||||
switch ext {
|
||||
case ".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".svg", ".tif", ".tiff", ".heic",
|
||||
".mp4", ".mov", ".webm", ".m4v", ".avi", ".mkv", ".3gp", ".flv":
|
||||
return ext
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func removePartialDownload(root *os.Root, filePath string) {
|
||||
if root == nil || strings.TrimSpace(filePath) == "" {
|
||||
return
|
||||
}
|
||||
_ = root.Remove(filePath)
|
||||
}
|
||||
|
||||
// sanitizeMediaLogURL 脱敏 URL 用于日志记录（去除 query 参数中可能的 token 信息）
|
||||
func sanitizeMediaLogURL(rawURL string) string {
|
||||
parsed, err := url.Parse(rawURL)
|
||||
if err != nil {
|
||||
if len(rawURL) > 80 {
|
||||
return rawURL[:80] + "..."
|
||||
}
|
||||
return rawURL
|
||||
}
|
||||
safe := parsed.Scheme + "://" + parsed.Host + parsed.Path
|
||||
if len(safe) > 120 {
|
||||
return safe[:120] + "..."
|
||||
}
|
||||
return safe
|
||||
}
|
||||
@@ -1,119 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSoraMediaStorage_StoreFromURLs(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("data"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
MaxConcurrentDownloads: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
urls, err := storage.StoreFromURLs(context.Background(), "image", []string{server.URL + "/img.png"})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, urls, 1)
|
||||
require.True(t, strings.HasPrefix(urls[0], "/image/"))
|
||||
require.True(t, strings.HasSuffix(urls[0], ".png"))
|
||||
|
||||
localPath := filepath.Join(tmpDir, filepath.FromSlash(strings.TrimPrefix(urls[0], "/")))
|
||||
require.FileExists(t, localPath)
|
||||
}
|
||||
|
||||
func TestSoraMediaStorage_FallbackToUpstream(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
FallbackToUpstream: true,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
url := server.URL + "/broken.png"
|
||||
urls, err := storage.StoreFromURLs(context.Background(), "image", []string{url})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{url}, urls)
|
||||
}
|
||||
|
||||
func TestSoraMediaStorage_MaxDownloadBytes(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "image/png")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte("too-large"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
cfg := &config.Config{
|
||||
Sora: config.SoraConfig{
|
||||
Storage: config.SoraStorageConfig{
|
||||
Type: "local",
|
||||
LocalPath: tmpDir,
|
||||
MaxDownloadBytes: 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
storage := NewSoraMediaStorage(cfg)
|
||||
_, err := storage.StoreFromURLs(context.Background(), "image", []string{server.URL + "/img.png"})
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestNormalizeSoraFileExt(t *testing.T) {
|
||||
require.Equal(t, ".png", normalizeSoraFileExt(".PNG"))
|
||||
require.Equal(t, ".mp4", normalizeSoraFileExt(".mp4"))
|
||||
require.Equal(t, "", normalizeSoraFileExt("../../etc/passwd"))
|
||||
require.Equal(t, "", normalizeSoraFileExt(".php"))
|
||||
}
|
||||
|
||||
func TestRemovePartialDownload(t *testing.T) {
|
||||
tmpDir := t.TempDir()
|
||||
root, err := os.OpenRoot(tmpDir)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = root.Close() }()
|
||||
|
||||
filePath := "partial.bin"
|
||||
f, err := root.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
|
||||
require.NoError(t, err)
|
||||
_, _ = f.WriteString("partial")
|
||||
_ = f.Close()
|
||||
|
||||
removePartialDownload(root, filePath)
|
||||
_, err = root.Stat(filePath)
|
||||
require.Error(t, err)
|
||||
require.True(t, os.IsNotExist(err))
|
||||
}
|
||||
@@ -1,488 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
)
|
||||
|
||||
// SoraModelConfig Sora 模型配置
|
||||
type SoraModelConfig struct {
|
||||
Type string
|
||||
Width int
|
||||
Height int
|
||||
Orientation string
|
||||
Frames int
|
||||
Model string
|
||||
Size string
|
||||
RequirePro bool
|
||||
// Prompt-enhance 专用参数
|
||||
ExpansionLevel string
|
||||
DurationS int
|
||||
}
|
||||
|
||||
var soraModelConfigs = map[string]SoraModelConfig{
|
||||
"gpt-image": {
|
||||
Type: "image",
|
||||
Width: 360,
|
||||
Height: 360,
|
||||
},
|
||||
"gpt-image-landscape": {
|
||||
Type: "image",
|
||||
Width: 540,
|
||||
Height: 360,
|
||||
},
|
||||
"gpt-image-portrait": {
|
||||
Type: "image",
|
||||
Width: 360,
|
||||
Height: 540,
|
||||
},
|
||||
"sora2-landscape-10s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 300,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
},
|
||||
"sora2-portrait-10s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 300,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
},
|
||||
"sora2-landscape-15s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 450,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
},
|
||||
"sora2-portrait-15s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 450,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
},
|
||||
"sora2-landscape-25s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 750,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2-portrait-25s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 750,
|
||||
Model: "sy_8",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-landscape-10s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 300,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-portrait-10s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 300,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-landscape-15s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 450,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-portrait-15s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 450,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-landscape-25s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 750,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-portrait-25s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 750,
|
||||
Model: "sy_ore",
|
||||
Size: "small",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-hd-landscape-10s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 300,
|
||||
Model: "sy_ore",
|
||||
Size: "large",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-hd-portrait-10s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 300,
|
||||
Model: "sy_ore",
|
||||
Size: "large",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-hd-landscape-15s": {
|
||||
Type: "video",
|
||||
Orientation: "landscape",
|
||||
Frames: 450,
|
||||
Model: "sy_ore",
|
||||
Size: "large",
|
||||
RequirePro: true,
|
||||
},
|
||||
"sora2pro-hd-portrait-15s": {
|
||||
Type: "video",
|
||||
Orientation: "portrait",
|
||||
Frames: 450,
|
||||
Model: "sy_ore",
|
||||
Size: "large",
|
||||
RequirePro: true,
|
||||
},
|
||||
"prompt-enhance-short-10s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "short",
|
||||
DurationS: 10,
|
||||
},
|
||||
"prompt-enhance-short-15s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "short",
|
||||
DurationS: 15,
|
||||
},
|
||||
"prompt-enhance-short-20s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "short",
|
||||
DurationS: 20,
|
||||
},
|
||||
"prompt-enhance-medium-10s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "medium",
|
||||
DurationS: 10,
|
||||
},
|
||||
"prompt-enhance-medium-15s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "medium",
|
||||
DurationS: 15,
|
||||
},
|
||||
"prompt-enhance-medium-20s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "medium",
|
||||
DurationS: 20,
|
||||
},
|
||||
"prompt-enhance-long-10s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "long",
|
||||
DurationS: 10,
|
||||
},
|
||||
"prompt-enhance-long-15s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "long",
|
||||
DurationS: 15,
|
||||
},
|
||||
"prompt-enhance-long-20s": {
|
||||
Type: "prompt_enhance",
|
||||
ExpansionLevel: "long",
|
||||
DurationS: 20,
|
||||
},
|
||||
}
|
||||
|
||||
var soraModelIDs = []string{
|
||||
"gpt-image",
|
||||
"gpt-image-landscape",
|
||||
"gpt-image-portrait",
|
||||
"sora2-landscape-10s",
|
||||
"sora2-portrait-10s",
|
||||
"sora2-landscape-15s",
|
||||
"sora2-portrait-15s",
|
||||
"sora2-landscape-25s",
|
||||
"sora2-portrait-25s",
|
||||
"sora2pro-landscape-10s",
|
||||
"sora2pro-portrait-10s",
|
||||
"sora2pro-landscape-15s",
|
||||
"sora2pro-portrait-15s",
|
||||
"sora2pro-landscape-25s",
|
||||
"sora2pro-portrait-25s",
|
||||
"sora2pro-hd-landscape-10s",
|
||||
"sora2pro-hd-portrait-10s",
|
||||
"sora2pro-hd-landscape-15s",
|
||||
"sora2pro-hd-portrait-15s",
|
||||
"prompt-enhance-short-10s",
|
||||
"prompt-enhance-short-15s",
|
||||
"prompt-enhance-short-20s",
|
||||
"prompt-enhance-medium-10s",
|
||||
"prompt-enhance-medium-15s",
|
||||
"prompt-enhance-medium-20s",
|
||||
"prompt-enhance-long-10s",
|
||||
"prompt-enhance-long-15s",
|
||||
"prompt-enhance-long-20s",
|
||||
}
|
||||
|
||||
// GetSoraModelConfig 返回 Sora 模型配置
|
||||
func GetSoraModelConfig(model string) (SoraModelConfig, bool) {
|
||||
key := strings.ToLower(strings.TrimSpace(model))
|
||||
cfg, ok := soraModelConfigs[key]
|
||||
return cfg, ok
|
||||
}
|
||||
|
||||
// SoraModelFamily 模型家族（前端 Sora 客户端使用）
|
||||
type SoraModelFamily struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Type string `json:"type"`
|
||||
Orientations []string `json:"orientations"`
|
||||
Durations []int `json:"durations,omitempty"`
|
||||
}
|
||||
|
||||
var (
|
||||
videoSuffixRe = regexp.MustCompile(`-(landscape|portrait)-(\d+)s$`)
|
||||
imageSuffixRe = regexp.MustCompile(`-(landscape|portrait)$`)
|
||||
|
||||
soraFamilyNames = map[string]string{
|
||||
"sora2": "Sora 2",
|
||||
"sora2pro": "Sora 2 Pro",
|
||||
"sora2pro-hd": "Sora 2 Pro HD",
|
||||
"gpt-image": "GPT Image",
|
||||
}
|
||||
)
|
||||
|
||||
// BuildSoraModelFamilies 从 soraModelConfigs 自动聚合模型家族及其支持的方向和时长
|
||||
func BuildSoraModelFamilies() []SoraModelFamily {
|
||||
type familyData struct {
|
||||
modelType string
|
||||
orientations map[string]bool
|
||||
durations map[int]bool
|
||||
}
|
||||
families := make(map[string]*familyData)
|
||||
|
||||
for id, cfg := range soraModelConfigs {
|
||||
if cfg.Type == "prompt_enhance" {
|
||||
continue
|
||||
}
|
||||
var famID, orientation string
|
||||
var duration int
|
||||
|
||||
switch cfg.Type {
|
||||
case "video":
|
||||
if m := videoSuffixRe.FindStringSubmatch(id); m != nil {
|
||||
famID = id[:len(id)-len(m[0])]
|
||||
orientation = m[1]
|
||||
duration, _ = strconv.Atoi(m[2])
|
||||
}
|
||||
case "image":
|
||||
if m := imageSuffixRe.FindStringSubmatch(id); m != nil {
|
||||
famID = id[:len(id)-len(m[0])]
|
||||
orientation = m[1]
|
||||
} else {
|
||||
famID = id
|
||||
orientation = "square"
|
||||
}
|
||||
}
|
||||
if famID == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
fd, ok := families[famID]
|
||||
if !ok {
|
||||
fd = &familyData{
|
||||
modelType: cfg.Type,
|
||||
orientations: make(map[string]bool),
|
||||
durations: make(map[int]bool),
|
||||
}
|
||||
families[famID] = fd
|
||||
}
|
||||
if orientation != "" {
|
||||
fd.orientations[orientation] = true
|
||||
}
|
||||
if duration > 0 {
|
||||
fd.durations[duration] = true
|
||||
}
|
||||
}
|
||||
|
||||
// 排序：视频在前、图像在后，同类按名称排序
|
||||
famIDs := make([]string, 0, len(families))
|
||||
for id := range families {
|
||||
famIDs = append(famIDs, id)
|
||||
}
|
||||
sort.Slice(famIDs, func(i, j int) bool {
|
||||
fi, fj := families[famIDs[i]], families[famIDs[j]]
|
||||
if fi.modelType != fj.modelType {
|
||||
return fi.modelType == "video"
|
||||
}
|
||||
return famIDs[i] < famIDs[j]
|
||||
})
|
||||
|
||||
result := make([]SoraModelFamily, 0, len(famIDs))
|
||||
for _, famID := range famIDs {
|
||||
fd := families[famID]
|
||||
fam := SoraModelFamily{
|
||||
ID: famID,
|
||||
Name: soraFamilyNames[famID],
|
||||
Type: fd.modelType,
|
||||
}
|
||||
if fam.Name == "" {
|
||||
fam.Name = famID
|
||||
}
|
||||
for o := range fd.orientations {
|
||||
fam.Orientations = append(fam.Orientations, o)
|
||||
}
|
||||
sort.Strings(fam.Orientations)
|
||||
for d := range fd.durations {
|
||||
fam.Durations = append(fam.Durations, d)
|
||||
}
|
||||
sort.Ints(fam.Durations)
|
||||
result = append(result, fam)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// BuildSoraModelFamiliesFromIDs 从任意模型 ID 列表聚合模型家族（用于解析上游返回的模型列表）。
|
||||
// 通过命名约定自动识别视频/图像模型并分组。
|
||||
func BuildSoraModelFamiliesFromIDs(modelIDs []string) []SoraModelFamily {
|
||||
type familyData struct {
|
||||
modelType string
|
||||
orientations map[string]bool
|
||||
durations map[int]bool
|
||||
}
|
||||
families := make(map[string]*familyData)
|
||||
|
||||
for _, id := range modelIDs {
|
||||
id = strings.ToLower(strings.TrimSpace(id))
|
||||
if id == "" || strings.HasPrefix(id, "prompt-enhance") {
|
||||
continue
|
||||
}
|
||||
|
||||
var famID, orientation, modelType string
|
||||
var duration int
|
||||
|
||||
if m := videoSuffixRe.FindStringSubmatch(id); m != nil {
|
||||
// 视频模型: {family}-{orientation}-{duration}s
|
||||
famID = id[:len(id)-len(m[0])]
|
||||
orientation = m[1]
|
||||
duration, _ = strconv.Atoi(m[2])
|
||||
modelType = "video"
|
||||
} else if m := imageSuffixRe.FindStringSubmatch(id); m != nil {
|
||||
// 图像模型（带方向）: {family}-{orientation}
|
||||
famID = id[:len(id)-len(m[0])]
|
||||
orientation = m[1]
|
||||
modelType = "image"
|
||||
} else if cfg, ok := soraModelConfigs[id]; ok && cfg.Type == "image" {
|
||||
// 已知的无后缀图像模型（如 gpt-image）
|
||||
famID = id
|
||||
orientation = "square"
|
||||
modelType = "image"
|
||||
} else if strings.Contains(id, "image") {
|
||||
// 未知但名称包含 image 的模型，推断为图像模型
|
||||
famID = id
|
||||
orientation = "square"
|
||||
modelType = "image"
|
||||
} else {
|
||||
continue
|
||||
}
|
||||
|
||||
if famID == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
fd, ok := families[famID]
|
||||
if !ok {
|
||||
fd = &familyData{
|
||||
modelType: modelType,
|
||||
orientations: make(map[string]bool),
|
||||
durations: make(map[int]bool),
|
||||
}
|
||||
families[famID] = fd
|
||||
}
|
||||
if orientation != "" {
|
||||
fd.orientations[orientation] = true
|
||||
}
|
||||
if duration > 0 {
|
||||
fd.durations[duration] = true
|
||||
}
|
||||
}
|
||||
|
||||
famIDs := make([]string, 0, len(families))
|
||||
for id := range families {
|
||||
famIDs = append(famIDs, id)
|
||||
}
|
||||
sort.Slice(famIDs, func(i, j int) bool {
|
||||
fi, fj := families[famIDs[i]], families[famIDs[j]]
|
||||
if fi.modelType != fj.modelType {
|
||||
return fi.modelType == "video"
|
||||
}
|
||||
return famIDs[i] < famIDs[j]
|
||||
})
|
||||
|
||||
result := make([]SoraModelFamily, 0, len(famIDs))
|
||||
for _, famID := range famIDs {
|
||||
fd := families[famID]
|
||||
fam := SoraModelFamily{
|
||||
ID: famID,
|
||||
Name: soraFamilyNames[famID],
|
||||
Type: fd.modelType,
|
||||
}
|
||||
if fam.Name == "" {
|
||||
fam.Name = famID
|
||||
}
|
||||
for o := range fd.orientations {
|
||||
fam.Orientations = append(fam.Orientations, o)
|
||||
}
|
||||
sort.Strings(fam.Orientations)
|
||||
for d := range fd.durations {
|
||||
fam.Durations = append(fam.Durations, d)
|
||||
}
|
||||
sort.Ints(fam.Durations)
|
||||
result = append(result, fam)
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// DefaultSoraModels returns the default Sora model list.
|
||||
func DefaultSoraModels(cfg *config.Config) []openai.Model {
|
||||
models := make([]openai.Model, 0, len(soraModelIDs))
|
||||
for _, id := range soraModelIDs {
|
||||
models = append(models, openai.Model{
|
||||
ID: id,
|
||||
Object: "model",
|
||||
OwnedBy: "openai",
|
||||
Type: "model",
|
||||
DisplayName: id,
|
||||
})
|
||||
}
|
||||
if cfg != nil && cfg.Gateway.SoraModelFilters.HidePromptEnhance {
|
||||
filtered := models[:0]
|
||||
for _, model := range models {
|
||||
if strings.HasPrefix(strings.ToLower(model.ID), "prompt-enhance") {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
models = filtered
|
||||
}
|
||||
return models
|
||||
}
|
||||
@@ -1,110 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
)
|
||||
|
||||
// SoraQuotaService manages Sora storage quota using only the system default quota.
|
||||
type SoraQuotaService struct {
|
||||
settingService *SettingService
|
||||
}
|
||||
|
||||
func NewSoraQuotaService(settingService *SettingService) *SoraQuotaService {
|
||||
return &SoraQuotaService{
|
||||
settingService: settingService,
|
||||
}
|
||||
}
|
||||
|
||||
type QuotaInfo struct {
|
||||
QuotaBytes int64 `json:"quota_bytes"`
|
||||
UsedBytes int64 `json:"used_bytes"`
|
||||
AvailableBytes int64 `json:"available_bytes"`
|
||||
QuotaSource string `json:"quota_source"`
|
||||
Source string `json:"source,omitempty"`
|
||||
}
|
||||
|
||||
var ErrSoraStorageQuotaExceeded = errors.New("sora storage quota exceeded")
|
||||
|
||||
type QuotaExceededError struct {
|
||||
QuotaBytes int64
|
||||
UsedBytes int64
|
||||
}
|
||||
|
||||
func (e *QuotaExceededError) Error() string {
|
||||
if e == nil {
|
||||
return "storage quota exceeded"
|
||||
}
|
||||
return fmt.Sprintf("storage quota exceeded (used %d / quota %d bytes)", e.UsedBytes, e.QuotaBytes)
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) GetQuota(ctx context.Context, userID int64) (*QuotaInfo, error) {
|
||||
info := &QuotaInfo{
|
||||
UsedBytes: 0,
|
||||
}
|
||||
|
||||
defaultQuota := s.getSystemDefaultQuota(ctx)
|
||||
if defaultQuota > 0 {
|
||||
info.QuotaBytes = defaultQuota
|
||||
info.QuotaSource = "system"
|
||||
info.Source = info.QuotaSource
|
||||
info.AvailableBytes = defaultQuota
|
||||
return info, nil
|
||||
}
|
||||
|
||||
info.QuotaSource = "unlimited"
|
||||
info.Source = info.QuotaSource
|
||||
info.AvailableBytes = 0
|
||||
return info, nil
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) CheckQuota(ctx context.Context, userID int64, additionalBytes int64) error {
|
||||
quota, err := s.GetQuota(ctx, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if quota.QuotaBytes == 0 {
|
||||
return nil
|
||||
}
|
||||
if additionalBytes > quota.QuotaBytes {
|
||||
return &QuotaExceededError{
|
||||
QuotaBytes: quota.QuotaBytes,
|
||||
UsedBytes: 0,
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) AddUsage(ctx context.Context, userID int64, bytes int64) error {
|
||||
if bytes <= 0 {
|
||||
return nil
|
||||
}
|
||||
logger.LegacyPrintf("service.sora_quota", "debug: [SoraQuota] AddUsage user=%d +%d (usage tracking removed)", userID, bytes)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) ReleaseUsage(ctx context.Context, userID int64, bytes int64) error {
|
||||
if bytes <= 0 {
|
||||
return nil
|
||||
}
|
||||
logger.LegacyPrintf("service.sora_quota", "debug: [SoraQuota] ReleaseUsage user=%d -%d (usage tracking removed)", userID, bytes)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) getSystemDefaultQuota(ctx context.Context) int64 {
|
||||
if s.settingService == nil {
|
||||
return 0
|
||||
}
|
||||
settings, err := s.settingService.GetSoraS3Settings(ctx)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return settings.DefaultStorageQuotaBytes
|
||||
}
|
||||
|
||||
func (s *SoraQuotaService) GetQuotaFromSettings(ctx context.Context) int64 {
|
||||
return s.getSystemDefaultQuota(ctx)
|
||||
}
|
||||
@@ -1,274 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ==================== Stub: SettingRepository (用于 SettingService) ====================
|
||||
|
||||
var _ SettingRepository = (*stubSettingRepoForQuota)(nil)
|
||||
|
||||
type stubSettingRepoForQuota struct {
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func newStubSettingRepoForQuota(values map[string]string) *stubSettingRepoForQuota {
|
||||
if values == nil {
|
||||
values = make(map[string]string)
|
||||
}
|
||||
return &stubSettingRepoForQuota{values: values}
|
||||
}
|
||||
|
||||
func (r *stubSettingRepoForQuota) Get(_ context.Context, key string) (*Setting, error) {
|
||||
if v, ok := r.values[key]; ok {
|
||||
return &Setting{Key: key, Value: v}, nil
|
||||
}
|
||||
return nil, ErrSettingNotFound
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) GetValue(_ context.Context, key string) (string, error) {
|
||||
if v, ok := r.values[key]; ok {
|
||||
return v, nil
|
||||
}
|
||||
return "", ErrSettingNotFound
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) Set(_ context.Context, key, value string) error {
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
|
||||
result := make(map[string]string)
|
||||
for _, k := range keys {
|
||||
if v, ok := r.values[k]; ok {
|
||||
result[k] = v
|
||||
}
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) SetMultiple(_ context.Context, settings map[string]string) error {
|
||||
for k, v := range settings {
|
||||
r.values[k] = v
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) GetAll(_ context.Context) (map[string]string, error) {
|
||||
return r.values, nil
|
||||
}
|
||||
func (r *stubSettingRepoForQuota) Delete(_ context.Context, key string) error {
|
||||
delete(r.values, key)
|
||||
return nil
|
||||
}
|
||||
|
||||
// ==================== GetQuota ====================
|
||||
|
||||
func TestGetQuota_SystemLevel(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "104857600", // 100MB
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
quota, err := svc.GetQuota(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(104857600), quota.QuotaBytes)
|
||||
require.Equal(t, "system", quota.Source)
|
||||
}
|
||||
|
||||
func TestGetQuota_NoLimit(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
quota, err := svc.GetQuota(context.Background(), 1)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), quota.QuotaBytes)
|
||||
require.Equal(t, "unlimited", quota.Source)
|
||||
}
|
||||
|
||||
// ==================== CheckQuota ====================
|
||||
|
||||
func TestCheckQuota_Sufficient(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "104857600", // 100MB
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
err := svc.CheckQuota(context.Background(), 1, 1024)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestCheckQuota_Exceeded(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "1024", // 1KB
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
err := svc.CheckQuota(context.Background(), 1, 2048)
|
||||
require.Error(t, err)
|
||||
var qe *QuotaExceededError
|
||||
require.ErrorAs(t, err, &qe)
|
||||
require.Equal(t, int64(1024), qe.QuotaBytes)
|
||||
require.Equal(t, int64(0), qe.UsedBytes)
|
||||
}
|
||||
|
||||
func TestCheckQuota_NoLimit(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
err := svc.CheckQuota(context.Background(), 1, 999999999)
|
||||
require.NoError(t, err) // 无限制时始终通过
|
||||
}
|
||||
|
||||
// ==================== AddUsage/ReleaseUsage ====================
|
||||
// 用量追踪已移除，仅记录日志
|
||||
|
||||
func TestAddUsage_NoTracking(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
err := svc.AddUsage(context.Background(), 1, 1024)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestAddUsage_ZeroBytes(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
err := svc.AddUsage(context.Background(), 1, 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestReleaseUsage_NoTracking(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
err := svc.ReleaseUsage(context.Background(), 1, 1024)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestReleaseUsage_ZeroBytes(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
err := svc.ReleaseUsage(context.Background(), 1, 0)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// ==================== GetQuotaFromSettings ====================
|
||||
|
||||
func TestGetQuotaFromSettings_NilSettingService(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
require.Equal(t, int64(0), svc.GetQuotaFromSettings(context.Background()))
|
||||
}
|
||||
|
||||
func TestGetQuotaFromSettings_WithSettings(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "52428800", // 50MB
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
require.Equal(t, int64(52428800), svc.GetQuotaFromSettings(context.Background()))
|
||||
}
|
||||
|
||||
// ==================== SetUserSoraQuota ====================
|
||||
// 用户级配额设置已移除
|
||||
|
||||
|
||||
// ==================== GetQuota — 字段完整性 ====================
|
||||
|
||||
func TestGetQuota_SystemLevel_AllFields(t *testing.T) {
|
||||
// 验证 QuotaInfo 所有字段在系统配额模式下正确
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "104857600", // 100MB
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
quota, err := svc.GetQuota(context.Background(), 42)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(104857600), quota.QuotaBytes)
|
||||
require.Equal(t, int64(0), quota.UsedBytes, "用量追踪已移除，UsedBytes 应始终为 0")
|
||||
require.Equal(t, int64(104857600), quota.AvailableBytes, "无用量追踪时，可用 = 总配额")
|
||||
require.Equal(t, "system", quota.QuotaSource)
|
||||
require.Equal(t, "system", quota.Source)
|
||||
}
|
||||
|
||||
func TestGetQuota_Unlimited_AllFields(t *testing.T) {
|
||||
// 验证无限制模式所有字段
|
||||
svc := NewSoraQuotaService(nil)
|
||||
|
||||
quota, err := svc.GetQuota(context.Background(), 99)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), quota.QuotaBytes)
|
||||
require.Equal(t, int64(0), quota.UsedBytes)
|
||||
require.Equal(t, int64(0), quota.AvailableBytes)
|
||||
require.Equal(t, "unlimited", quota.QuotaSource)
|
||||
require.Equal(t, "unlimited", quota.Source)
|
||||
}
|
||||
|
||||
// ==================== CheckQuota — 边界条件 ====================
|
||||
|
||||
func TestCheckQuota_ExactlyAtLimit(t *testing.T) {
|
||||
// 请求大小 == 配额 → 应通过（<= 判断）
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "1024",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
err := svc.CheckQuota(context.Background(), 1, 1024)
|
||||
require.NoError(t, err) // exactly at limit is OK
|
||||
}
|
||||
|
||||
func TestCheckQuota_OneByteOverLimit(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraDefaultStorageQuotaBytes: "1024",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
svc := NewSoraQuotaService(settingService)
|
||||
|
||||
err := svc.CheckQuota(context.Background(), 1, 1025)
|
||||
require.Error(t, err)
|
||||
|
||||
var qe *QuotaExceededError
|
||||
require.ErrorAs(t, err, &qe)
|
||||
require.Equal(t, int64(1024), qe.QuotaBytes)
|
||||
require.Equal(t, int64(0), qe.UsedBytes)
|
||||
}
|
||||
|
||||
func TestCheckQuota_NegativeBytes(t *testing.T) {
|
||||
// 负数字节应视为无操作
|
||||
svc := NewSoraQuotaService(nil)
|
||||
err := svc.CheckQuota(context.Background(), 1, -100)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// ==================== AddUsage/ReleaseUsage — 边界条件 ====================
|
||||
|
||||
func TestAddUsage_NegativeBytes(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
err := svc.AddUsage(context.Background(), 1, -1)
|
||||
require.NoError(t, err) // <= 0 是 no-op
|
||||
}
|
||||
|
||||
func TestReleaseUsage_NegativeBytes(t *testing.T) {
|
||||
svc := NewSoraQuotaService(nil)
|
||||
err := svc.ReleaseUsage(context.Background(), 1, -999)
|
||||
require.NoError(t, err) // <= 0 是 no-op
|
||||
}
|
||||
|
||||
// ==================== QuotaExceededError ====================
|
||||
|
||||
func TestQuotaExceededError_NilSafe(t *testing.T) {
|
||||
var e *QuotaExceededError
|
||||
msg := e.Error()
|
||||
require.Contains(t, msg, "storage quota exceeded")
|
||||
}
|
||||
|
||||
func TestQuotaExceededError_Format(t *testing.T) {
|
||||
e := &QuotaExceededError{QuotaBytes: 1024, UsedBytes: 512}
|
||||
msg := e.Error()
|
||||
require.Contains(t, msg, "512")
|
||||
require.Contains(t, msg, "1024")
|
||||
}
|
||||
@@ -1,398 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"path"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
v4 "github.com/aws/aws-sdk-go-v2/aws/signer/v4"
|
||||
awsconfig "github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
)
|
||||
|
||||
// SoraS3Storage 负责 Sora 媒体文件的 S3 存储操作。
|
||||
// 从 Settings 表读取 S3 配置，初始化并缓存 S3 客户端。
|
||||
type SoraS3Storage struct {
|
||||
settingService *SettingService
|
||||
|
||||
mu sync.RWMutex
|
||||
client *s3.Client
|
||||
cfg *SoraS3Settings // 上次加载的配置快照
|
||||
|
||||
healthCheckedAt time.Time
|
||||
healthErr error
|
||||
healthTTL time.Duration
|
||||
}
|
||||
|
||||
const defaultSoraS3HealthTTL = 30 * time.Second
|
||||
|
||||
// UpstreamDownloadError 表示从上游下载媒体失败（包含 HTTP 状态码）。
|
||||
type UpstreamDownloadError struct {
|
||||
StatusCode int
|
||||
}
|
||||
|
||||
func (e *UpstreamDownloadError) Error() string {
|
||||
if e == nil {
|
||||
return "upstream download failed"
|
||||
}
|
||||
return fmt.Sprintf("upstream returned %d", e.StatusCode)
|
||||
}
|
||||
|
||||
// NewSoraS3Storage 创建 S3 存储服务实例。
|
||||
func NewSoraS3Storage(settingService *SettingService) *SoraS3Storage {
|
||||
return &SoraS3Storage{
|
||||
settingService: settingService,
|
||||
healthTTL: defaultSoraS3HealthTTL,
|
||||
}
|
||||
}
|
||||
|
||||
// Enabled 返回 S3 存储是否已启用且配置有效。
|
||||
func (s *SoraS3Storage) Enabled(ctx context.Context) bool {
|
||||
cfg, err := s.getConfig(ctx)
|
||||
if err != nil || cfg == nil {
|
||||
return false
|
||||
}
|
||||
return cfg.Enabled && cfg.Bucket != ""
|
||||
}
|
||||
|
||||
// getConfig 获取当前 S3 配置（从 settings 表读取）。
|
||||
func (s *SoraS3Storage) getConfig(ctx context.Context) (*SoraS3Settings, error) {
|
||||
if s.settingService == nil {
|
||||
return nil, fmt.Errorf("setting service not available")
|
||||
}
|
||||
return s.settingService.GetSoraS3Settings(ctx)
|
||||
}
|
||||
|
||||
// getClient 获取或初始化 S3 客户端（带缓存）。
|
||||
// 配置变更时调用 RefreshClient 清除缓存。
|
||||
func (s *SoraS3Storage) getClient(ctx context.Context) (*s3.Client, *SoraS3Settings, error) {
|
||||
s.mu.RLock()
|
||||
if s.client != nil && s.cfg != nil {
|
||||
client, cfg := s.client, s.cfg
|
||||
s.mu.RUnlock()
|
||||
return client, cfg, nil
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
|
||||
return s.initClient(ctx)
|
||||
}
|
||||
|
||||
func (s *SoraS3Storage) initClient(ctx context.Context) (*s3.Client, *SoraS3Settings, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
// 双重检查
|
||||
if s.client != nil && s.cfg != nil {
|
||||
return s.client, s.cfg, nil
|
||||
}
|
||||
|
||||
cfg, err := s.getConfig(ctx)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("load s3 config: %w", err)
|
||||
}
|
||||
if !cfg.Enabled {
|
||||
return nil, nil, fmt.Errorf("sora s3 storage is disabled")
|
||||
}
|
||||
if cfg.Bucket == "" || cfg.AccessKeyID == "" || cfg.SecretAccessKey == "" {
|
||||
return nil, nil, fmt.Errorf("sora s3 config incomplete: bucket, access_key_id, secret_access_key are required")
|
||||
}
|
||||
|
||||
client, region, err := buildSoraS3Client(ctx, cfg)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
s.client = client
|
||||
s.cfg = cfg
|
||||
logger.LegacyPrintf("service.sora_s3", "[SoraS3] 客户端已初始化 bucket=%s endpoint=%s region=%s", cfg.Bucket, cfg.Endpoint, region)
|
||||
return client, cfg, nil
|
||||
}
|
||||
|
||||
// RefreshClient 清除缓存的 S3 客户端，下次使用时重新初始化。
|
||||
// 应在系统设置中 S3 配置变更时调用。
|
||||
func (s *SoraS3Storage) RefreshClient() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.client = nil
|
||||
s.cfg = nil
|
||||
s.healthCheckedAt = time.Time{}
|
||||
s.healthErr = nil
|
||||
logger.LegacyPrintf("service.sora_s3", "[SoraS3] 客户端缓存已清除，下次使用将重新初始化")
|
||||
}
|
||||
|
||||
// TestConnection 测试 S3 连接（HeadBucket）。
|
||||
func (s *SoraS3Storage) TestConnection(ctx context.Context) error {
|
||||
client, cfg, err := s.getClient(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: &cfg.Bucket,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("s3 HeadBucket failed: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsHealthy 返回 S3 健康状态（带短缓存，避免每次请求都触发 HeadBucket）。
|
||||
func (s *SoraS3Storage) IsHealthy(ctx context.Context) bool {
|
||||
if s == nil {
|
||||
return false
|
||||
}
|
||||
now := time.Now()
|
||||
s.mu.RLock()
|
||||
lastCheck := s.healthCheckedAt
|
||||
lastErr := s.healthErr
|
||||
ttl := s.healthTTL
|
||||
s.mu.RUnlock()
|
||||
|
||||
if ttl <= 0 {
|
||||
ttl = defaultSoraS3HealthTTL
|
||||
}
|
||||
if !lastCheck.IsZero() && now.Sub(lastCheck) < ttl {
|
||||
return lastErr == nil
|
||||
}
|
||||
|
||||
err := s.TestConnection(ctx)
|
||||
s.mu.Lock()
|
||||
s.healthCheckedAt = time.Now()
|
||||
s.healthErr = err
|
||||
s.mu.Unlock()
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// TestConnectionWithSettings 使用临时配置测试连接，不污染缓存的客户端。
|
||||
func (s *SoraS3Storage) TestConnectionWithSettings(ctx context.Context, cfg *SoraS3Settings) error {
|
||||
if cfg == nil {
|
||||
return fmt.Errorf("s3 config is required")
|
||||
}
|
||||
if !cfg.Enabled {
|
||||
return fmt.Errorf("sora s3 storage is disabled")
|
||||
}
|
||||
if cfg.Endpoint == "" || cfg.Bucket == "" || cfg.AccessKeyID == "" || cfg.SecretAccessKey == "" {
|
||||
return fmt.Errorf("sora s3 config incomplete: endpoint, bucket, access_key_id, secret_access_key are required")
|
||||
}
|
||||
client, _, err := buildSoraS3Client(ctx, cfg)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: &cfg.Bucket,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("s3 HeadBucket failed: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// GenerateObjectKey 生成 S3 object key。
|
||||
// 格式: {prefix}sora/{userID}/{YYYY/MM/DD}/{uuid}.{ext}
|
||||
func (s *SoraS3Storage) GenerateObjectKey(prefix string, userID int64, ext string) string {
|
||||
if !strings.HasPrefix(ext, ".") {
|
||||
ext = "." + ext
|
||||
}
|
||||
datePath := time.Now().Format("2006/01/02")
|
||||
key := fmt.Sprintf("sora/%d/%s/%s%s", userID, datePath, uuid.NewString(), ext)
|
||||
if prefix != "" {
|
||||
prefix = strings.TrimRight(prefix, "/") + "/"
|
||||
key = prefix + key
|
||||
}
|
||||
return key
|
||||
}
|
||||
|
||||
// UploadFromURL 从上游 URL 下载并流式上传到 S3。
|
||||
// 返回 S3 object key。
|
||||
func (s *SoraS3Storage) UploadFromURL(ctx context.Context, userID int64, sourceURL string) (string, int64, error) {
|
||||
client, cfg, err := s.getClient(ctx)
|
||||
if err != nil {
|
||||
return "", 0, err
|
||||
}
|
||||
|
||||
// 下载源文件
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, sourceURL, nil)
|
||||
if err != nil {
|
||||
return "", 0, fmt.Errorf("create download request: %w", err)
|
||||
}
|
||||
httpClient := &http.Client{Timeout: 5 * time.Minute}
|
||||
resp, err := httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", 0, fmt.Errorf("download from upstream: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
_ = resp.Body.Close()
|
||||
}()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", 0, &UpstreamDownloadError{StatusCode: resp.StatusCode}
|
||||
}
|
||||
|
||||
// 推断文件扩展名
|
||||
ext := fileExtFromURL(sourceURL)
|
||||
if ext == "" {
|
||||
ext = fileExtFromContentType(resp.Header.Get("Content-Type"))
|
||||
}
|
||||
if ext == "" {
|
||||
ext = ".bin"
|
||||
}
|
||||
|
||||
objectKey := s.GenerateObjectKey(cfg.Prefix, userID, ext)
|
||||
|
||||
// 检测 Content-Type
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
contentType = "application/octet-stream"
|
||||
}
|
||||
|
||||
reader, writer := io.Pipe()
|
||||
uploadErrCh := make(chan error, 1)
|
||||
go func() {
|
||||
defer close(uploadErrCh)
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
logger.LegacyPrintf("service.sora_s3", "[SoraS3] Upload goroutine panic recovered: %v", r)
|
||||
uploadErrCh <- fmt.Errorf("upload panic: %v", r)
|
||||
}
|
||||
}()
|
||||
input := &s3.PutObjectInput{
|
||||
Bucket: &cfg.Bucket,
|
||||
Key: &objectKey,
|
||||
Body: reader,
|
||||
ContentType: &contentType,
|
||||
}
|
||||
if resp.ContentLength >= 0 {
|
||||
input.ContentLength = &resp.ContentLength
|
||||
}
|
||||
_, uploadErr := client.PutObject(ctx, input)
|
||||
uploadErrCh <- uploadErr
|
||||
}()
|
||||
|
||||
written, copyErr := io.CopyBuffer(writer, resp.Body, make([]byte, 1024*1024))
|
||||
_ = writer.CloseWithError(copyErr)
|
||||
uploadErr := <-uploadErrCh
|
||||
if copyErr != nil {
|
||||
return "", 0, fmt.Errorf("stream upload copy failed: %w", copyErr)
|
||||
}
|
||||
if uploadErr != nil {
|
||||
return "", 0, fmt.Errorf("s3 upload: %w", uploadErr)
|
||||
}
|
||||
|
||||
logger.LegacyPrintf("service.sora_s3", "[SoraS3] 上传完成 key=%s size=%d", objectKey, written)
|
||||
return objectKey, written, nil
|
||||
}
|
||||
|
||||
func buildSoraS3Client(ctx context.Context, cfg *SoraS3Settings) (*s3.Client, string, error) {
|
||||
if cfg == nil {
|
||||
return nil, "", fmt.Errorf("s3 config is required")
|
||||
}
|
||||
region := cfg.Region
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
}
|
||||
|
||||
awsCfg, err := awsconfig.LoadDefaultConfig(ctx,
|
||||
awsconfig.WithRegion(region),
|
||||
awsconfig.WithCredentialsProvider(
|
||||
credentials.NewStaticCredentialsProvider(cfg.AccessKeyID, cfg.SecretAccessKey, ""),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("load aws config: %w", err)
|
||||
}
|
||||
|
||||
client := s3.NewFromConfig(awsCfg, func(o *s3.Options) {
|
||||
if cfg.Endpoint != "" {
|
||||
o.BaseEndpoint = &cfg.Endpoint
|
||||
}
|
||||
if cfg.ForcePathStyle {
|
||||
o.UsePathStyle = true
|
||||
}
|
||||
o.APIOptions = append(o.APIOptions, v4.SwapComputePayloadSHA256ForUnsignedPayloadMiddleware)
|
||||
// 兼容非 TLS 连接（如 MinIO）的流式上传，避免 io.Pipe checksum 校验失败
|
||||
o.RequestChecksumCalculation = aws.RequestChecksumCalculationWhenRequired
|
||||
})
|
||||
return client, region, nil
|
||||
}
|
||||
|
||||
// DeleteObjects 删除一组 S3 object（遍历逐一删除）。
|
||||
func (s *SoraS3Storage) DeleteObjects(ctx context.Context, objectKeys []string) error {
|
||||
if len(objectKeys) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
client, cfg, err := s.getClient(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var lastErr error
|
||||
for _, key := range objectKeys {
|
||||
k := key
|
||||
_, err := client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: &cfg.Bucket,
|
||||
Key: &k,
|
||||
})
|
||||
if err != nil {
|
||||
logger.LegacyPrintf("service.sora_s3", "[SoraS3] 删除失败 key=%s err=%v", key, err)
|
||||
lastErr = err
|
||||
}
|
||||
}
|
||||
return lastErr
|
||||
}
|
||||
|
||||
// GetAccessURL 获取 S3 文件的访问 URL。
|
||||
// CDN URL 优先，否则生成 24h 预签名 URL。
|
||||
func (s *SoraS3Storage) GetAccessURL(ctx context.Context, objectKey string) (string, error) {
|
||||
_, cfg, err := s.getClient(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// CDN URL 优先
|
||||
if cfg.CDNURL != "" {
|
||||
cdnBase := strings.TrimRight(cfg.CDNURL, "/")
|
||||
return cdnBase + "/" + objectKey, nil
|
||||
}
|
||||
|
||||
// 生成 24h 预签名 URL
|
||||
return s.GeneratePresignedURL(ctx, objectKey, 24*time.Hour)
|
||||
}
|
||||
|
||||
// GeneratePresignedURL 生成预签名 URL。
|
||||
func (s *SoraS3Storage) GeneratePresignedURL(ctx context.Context, objectKey string, ttl time.Duration) (string, error) {
|
||||
client, cfg, err := s.getClient(ctx)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
presignClient := s3.NewPresignClient(client)
|
||||
result, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: &cfg.Bucket,
|
||||
Key: &objectKey,
|
||||
}, s3.WithPresignExpires(ttl))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("presign url: %w", err)
|
||||
}
|
||||
return result.URL, nil
|
||||
}
|
||||
|
||||
// GetMediaType 从 object key 推断媒体类型（image/video）。
|
||||
func GetMediaTypeFromKey(objectKey string) string {
|
||||
ext := strings.ToLower(path.Ext(objectKey))
|
||||
switch ext {
|
||||
case ".mp4", ".mov", ".webm", ".m4v", ".avi", ".mkv", ".3gp", ".flv":
|
||||
return "video"
|
||||
default:
|
||||
return "image"
|
||||
}
|
||||
}
|
||||
@@ -1,263 +0,0 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ==================== RefreshClient ====================
|
||||
|
||||
func TestRefreshClient(t *testing.T) {
|
||||
s := newS3StorageWithCDN("https://cdn.example.com")
|
||||
require.NotNil(t, s.client)
|
||||
require.NotNil(t, s.cfg)
|
||||
|
||||
s.RefreshClient()
|
||||
require.Nil(t, s.client)
|
||||
require.Nil(t, s.cfg)
|
||||
}
|
||||
|
||||
func TestRefreshClient_AlreadyNil(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
s.RefreshClient() // 不应 panic
|
||||
require.Nil(t, s.client)
|
||||
require.Nil(t, s.cfg)
|
||||
}
|
||||
|
||||
// ==================== GetMediaTypeFromKey ====================
|
||||
|
||||
func TestGetMediaTypeFromKey_VideoExtensions(t *testing.T) {
|
||||
for _, ext := range []string{".mp4", ".mov", ".webm", ".m4v", ".avi", ".mkv", ".3gp", ".flv"} {
|
||||
require.Equal(t, "video", GetMediaTypeFromKey("path/to/file"+ext), "ext=%s", ext)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetMediaTypeFromKey_VideoUpperCase(t *testing.T) {
|
||||
require.Equal(t, "video", GetMediaTypeFromKey("file.MP4"))
|
||||
require.Equal(t, "video", GetMediaTypeFromKey("file.MOV"))
|
||||
}
|
||||
|
||||
func TestGetMediaTypeFromKey_ImageExtensions(t *testing.T) {
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.png"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.jpg"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.jpeg"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.gif"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.webp"))
|
||||
}
|
||||
|
||||
func TestGetMediaTypeFromKey_NoExtension(t *testing.T) {
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("path/to/file"))
|
||||
}
|
||||
|
||||
func TestGetMediaTypeFromKey_UnknownExtension(t *testing.T) {
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.bin"))
|
||||
require.Equal(t, "image", GetMediaTypeFromKey("file.xyz"))
|
||||
}
|
||||
|
||||
// ==================== Enabled ====================
|
||||
|
||||
func TestEnabled_NilSettingService(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
require.False(t, s.Enabled(context.Background()))
|
||||
}
|
||||
|
||||
func TestEnabled_ConfigDisabled(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "false",
|
||||
SettingKeySoraS3Bucket: "test-bucket",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
require.False(t, s.Enabled(context.Background()))
|
||||
}
|
||||
|
||||
func TestEnabled_ConfigEnabledWithBucket(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "true",
|
||||
SettingKeySoraS3Bucket: "my-bucket",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
require.True(t, s.Enabled(context.Background()))
|
||||
}
|
||||
|
||||
func TestEnabled_ConfigEnabledEmptyBucket(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "true",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
require.False(t, s.Enabled(context.Background()))
|
||||
}
|
||||
|
||||
// ==================== initClient ====================
|
||||
|
||||
func TestInitClient_Disabled(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "false",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
|
||||
_, _, err := s.getClient(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "disabled")
|
||||
}
|
||||
|
||||
func TestInitClient_IncompleteConfig(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "true",
|
||||
SettingKeySoraS3Bucket: "test-bucket",
|
||||
// 缺少 access_key_id 和 secret_access_key
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
|
||||
_, _, err := s.getClient(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "incomplete")
|
||||
}
|
||||
|
||||
func TestInitClient_DefaultRegion(t *testing.T) {
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "true",
|
||||
SettingKeySoraS3Bucket: "test-bucket",
|
||||
SettingKeySoraS3AccessKeyID: "AKID",
|
||||
SettingKeySoraS3SecretAccessKey: "SECRET",
|
||||
// Region 为空 → 默认 us-east-1
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
|
||||
client, cfg, err := s.getClient(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, client)
|
||||
require.Equal(t, "test-bucket", cfg.Bucket)
|
||||
}
|
||||
|
||||
func TestInitClient_DoubleCheck(t *testing.T) {
|
||||
// 验证双重检查锁定：第二次 getClient 命中缓存
|
||||
settingRepo := newStubSettingRepoForQuota(map[string]string{
|
||||
SettingKeySoraS3Enabled: "true",
|
||||
SettingKeySoraS3Bucket: "test-bucket",
|
||||
SettingKeySoraS3AccessKeyID: "AKID",
|
||||
SettingKeySoraS3SecretAccessKey: "SECRET",
|
||||
})
|
||||
settingService := NewSettingService(settingRepo, &config.Config{})
|
||||
s := NewSoraS3Storage(settingService)
|
||||
|
||||
client1, _, err1 := s.getClient(context.Background())
|
||||
require.NoError(t, err1)
|
||||
client2, _, err2 := s.getClient(context.Background())
|
||||
require.NoError(t, err2)
|
||||
require.Equal(t, client1, client2) // 同一客户端实例
|
||||
}
|
||||
|
||||
func TestInitClient_NilSettingService(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
_, _, err := s.getClient(context.Background())
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), "setting service not available")
|
||||
}
|
||||
|
||||
// ==================== GenerateObjectKey ====================
|
||||
|
||||
func TestGenerateObjectKey_ExtWithoutDot(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
key := s.GenerateObjectKey("", 1, "mp4")
|
||||
require.Contains(t, key, ".mp4")
|
||||
require.True(t, len(key) > 0)
|
||||
}
|
||||
|
||||
func TestGenerateObjectKey_ExtWithDot(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
key := s.GenerateObjectKey("", 1, ".mp4")
|
||||
require.Contains(t, key, ".mp4")
|
||||
// 不应出现 ..mp4
|
||||
require.NotContains(t, key, "..mp4")
|
||||
}
|
||||
|
||||
func TestGenerateObjectKey_WithPrefix(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
key := s.GenerateObjectKey("uploads/", 42, ".png")
|
||||
require.True(t, len(key) > 0)
|
||||
require.Contains(t, key, "uploads/sora/42/")
|
||||
}
|
||||
|
||||
func TestGenerateObjectKey_PrefixWithoutTrailingSlash(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
key := s.GenerateObjectKey("uploads", 42, ".png")
|
||||
require.Contains(t, key, "uploads/sora/42/")
|
||||
}
|
||||
|
||||
// ==================== GeneratePresignedURL ====================
|
||||
|
||||
func TestGeneratePresignedURL_GetClientError(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil) // settingService=nil → getClient 失败
|
||||
_, err := s.GeneratePresignedURL(context.Background(), "key", 3600)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== GetAccessURL ====================
|
||||
|
||||
func TestGetAccessURL_CDN(t *testing.T) {
|
||||
s := newS3StorageWithCDN("https://cdn.example.com")
|
||||
url, err := s.GetAccessURL(context.Background(), "sora/1/2024/01/01/video.mp4")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://cdn.example.com/sora/1/2024/01/01/video.mp4", url)
|
||||
}
|
||||
|
||||
func TestGetAccessURL_CDNTrailingSlash(t *testing.T) {
|
||||
s := newS3StorageWithCDN("https://cdn.example.com/")
|
||||
url, err := s.GetAccessURL(context.Background(), "key.mp4")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://cdn.example.com/key.mp4", url)
|
||||
}
|
||||
|
||||
func TestGetAccessURL_GetClientError(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
_, err := s.GetAccessURL(context.Background(), "key")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== TestConnection ====================
|
||||
|
||||
func TestTestConnection_GetClientError(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
err := s.TestConnection(context.Background())
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== UploadFromURL ====================
|
||||
|
||||
func TestUploadFromURL_GetClientError(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
_, _, err := s.UploadFromURL(context.Background(), 1, "https://example.com/file.mp4")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// ==================== DeleteObjects ====================
|
||||
|
||||
func TestDeleteObjects_EmptyKeys(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
err := s.DeleteObjects(context.Background(), []string{})
|
||||
require.NoError(t, err) // 空列表直接返回
|
||||
}
|
||||
|
||||
func TestDeleteObjects_NilKeys(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
err := s.DeleteObjects(context.Background(), nil)
|
||||
require.NoError(t, err) // nil 列表直接返回
|
||||
}
|
||||
|
||||
func TestDeleteObjects_GetClientError(t *testing.T) {
|
||||
s := NewSoraS3Storage(nil)
|
||||
err := s.DeleteObjects(context.Background(), []string{"key1", "key2"})
|
||||
require.Error(t, err)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,149 +0,0 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// forwardToUpstream 将请求 HTTP 透传到上游 Sora 服务（用于 apikey 类型账号）。
|
||||
// 上游地址为 account.GetBaseURL() + "/sora/v1/chat/completions"，
|
||||
// 使用 account.GetCredential("api_key") 作为 Bearer Token。
|
||||
// 支持流式和非流式响应的直接透传。
|
||||
func (s *SoraGatewayService) forwardToUpstream(
|
||||
ctx context.Context,
|
||||
c *gin.Context,
|
||||
account *Account,
|
||||
body []byte,
|
||||
clientStream bool,
|
||||
startTime time.Time,
|
||||
) (*ForwardResult, error) {
|
||||
apiKey := account.GetCredential("api_key")
|
||||
if apiKey == "" {
|
||||
s.writeSoraError(c, http.StatusBadGateway, "upstream_error", "Sora apikey account missing api_key credential", clientStream)
|
||||
return nil, fmt.Errorf("sora apikey account %d missing api_key", account.ID)
|
||||
}
|
||||
|
||||
baseURL := account.GetBaseURL()
|
||||
if baseURL == "" {
|
||||
s.writeSoraError(c, http.StatusBadGateway, "upstream_error", "Sora apikey account missing base_url", clientStream)
|
||||
return nil, fmt.Errorf("sora apikey account %d missing base_url", account.ID)
|
||||
}
|
||||
// 校验 scheme 合法性（仅允许 http/https）
|
||||
if !strings.HasPrefix(baseURL, "http://") && !strings.HasPrefix(baseURL, "https://") {
|
||||
s.writeSoraError(c, http.StatusBadGateway, "upstream_error", "Sora apikey base_url must start with http:// or https://", clientStream)
|
||||
return nil, fmt.Errorf("sora apikey account %d invalid base_url scheme: %s", account.ID, baseURL)
|
||||
}
|
||||
upstreamURL := strings.TrimRight(baseURL, "/") + "/sora/v1/chat/completions"
|
||||
|
||||
// 构建上游请求
|
||||
upstreamReq, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
s.writeSoraError(c, http.StatusInternalServerError, "api_error", "Failed to create upstream request", clientStream)
|
||||
return nil, fmt.Errorf("create upstream request: %w", err)
|
||||
}
|
||||
|
||||
upstreamReq.Header.Set("Content-Type", "application/json")
|
||||
upstreamReq.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
// 透传客户端的部分请求头
|
||||
for _, header := range []string{"Accept", "Accept-Encoding"} {
|
||||
if v := c.GetHeader(header); v != "" {
|
||||
upstreamReq.Header.Set(header, v)
|
||||
}
|
||||
}
|
||||
|
||||
logger.LegacyPrintf("service.sora", "[ForwardUpstream] account=%d url=%s", account.ID, upstreamURL)
|
||||
|
||||
// 获取代理 URL
|
||||
proxyURL := ""
|
||||
if account.ProxyID != nil && account.Proxy != nil {
|
||||
proxyURL = account.Proxy.URL()
|
||||
}
|
||||
|
||||
// 发送请求
|
||||
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
s.writeSoraError(c, http.StatusBadGateway, "upstream_error", "Failed to connect to upstream Sora service", clientStream)
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadGateway,
|
||||
}
|
||||
}
|
||||
defer func() {
|
||||
_ = resp.Body.Close()
|
||||
}()
|
||||
|
||||
// 错误响应处理
|
||||
if resp.StatusCode >= 400 {
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
|
||||
|
||||
if s.shouldFailoverUpstreamError(resp.StatusCode) {
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: respBody,
|
||||
ResponseHeaders: resp.Header.Clone(),
|
||||
}
|
||||
}
|
||||
|
||||
// 非转移错误，直接透传给客户端
|
||||
c.Status(resp.StatusCode)
|
||||
for key, values := range resp.Header {
|
||||
for _, v := range values {
|
||||
c.Writer.Header().Add(key, v)
|
||||
}
|
||||
}
|
||||
if _, err := c.Writer.Write(respBody); err != nil {
|
||||
return nil, fmt.Errorf("write upstream error response: %w", err)
|
||||
}
|
||||
return nil, fmt.Errorf("upstream error: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
// 成功响应 — 直接透传
|
||||
c.Status(resp.StatusCode)
|
||||
for key, values := range resp.Header {
|
||||
lower := strings.ToLower(key)
|
||||
// 透传内容相关头部
|
||||
if lower == "content-type" || lower == "transfer-encoding" ||
|
||||
lower == "cache-control" || lower == "x-request-id" {
|
||||
for _, v := range values {
|
||||
c.Writer.Header().Add(key, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 流式复制响应体
|
||||
if flusher, ok := c.Writer.(http.Flusher); ok && clientStream {
|
||||
buf := make([]byte, 4096)
|
||||
for {
|
||||
n, readErr := resp.Body.Read(buf)
|
||||
if n > 0 {
|
||||
if _, err := c.Writer.Write(buf[:n]); err != nil {
|
||||
return nil, fmt.Errorf("stream upstream response write: %w", err)
|
||||
}
|
||||
flusher.Flush()
|
||||
}
|
||||
if readErr != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if _, err := io.Copy(c.Writer, resp.Body); err != nil {
|
||||
return nil, fmt.Errorf("copy upstream response: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
duration := time.Since(startTime)
|
||||
return &ForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Model: "", // 由调用方填充
|
||||
Stream: clientStream,
|
||||
Duration: duration,
|
||||
}, nil
|
||||
}
|
||||
@@ -471,41 +471,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewPaymentService,
|
||||
ProvidePaymentOrderExpiryService,
|
||||
ProvideHealthChecker,
|
||||
|
||||
// Sora 相关服务 (从本地版本合并)
|
||||
ProvideSoraMediaStorage,
|
||||
ProvideSoraMediaCleanupService,
|
||||
ProvideSoraSDKClient,
|
||||
wire.Bind(new(SoraClient), new(*SoraSDKClient)),
|
||||
NewSoraGatewayService,
|
||||
NewSoraGenerationService,
|
||||
NewSoraS3Storage,
|
||||
NewSoraQuotaService,
|
||||
)
|
||||
|
||||
// ProvideSoraMediaStorage 初始化 Sora 媒体存储 (从本地版本合并)
|
||||
func ProvideSoraMediaStorage(cfg *config.Config) *SoraMediaStorage {
|
||||
return NewSoraMediaStorage(cfg)
|
||||
}
|
||||
|
||||
func ProvideSoraSDKClient(
|
||||
cfg *config.Config,
|
||||
httpUpstream HTTPUpstream,
|
||||
tokenProvider *OpenAITokenProvider,
|
||||
accountRepo AccountRepository,
|
||||
soraAccountRepo SoraAccountRepository,
|
||||
) *SoraSDKClient {
|
||||
client := NewSoraSDKClient(cfg, httpUpstream, tokenProvider)
|
||||
client.SetAccountRepositories(accountRepo, soraAccountRepo)
|
||||
return client
|
||||
}
|
||||
|
||||
// ProvideSoraMediaCleanupService 创建并启动 Sora 媒体清理服务 (从本地版本合并)
|
||||
func ProvideSoraMediaCleanupService(storage *SoraMediaStorage, cfg *config.Config) *SoraMediaCleanupService {
|
||||
svc := NewSoraMediaCleanupService(storage, cfg)
|
||||
svc.Start()
|
||||
return svc
|
||||
}
|
||||
)
|
||||
|
||||
// ProvideOAuthRefreshAPI creates OAuthRefreshAPI with default lock TTL.
|
||||
// Wire cannot handle variadic parameters, so we wrap the constructor.
|
||||
|
||||
Reference in New Issue
Block a user