Skip to content

Commit 44f3358

Browse files
james-6-23james-6-23
committed
feat(images): support external image storage
1 parent 6c02cb9 commit 44f3358

19 files changed

Lines changed: 1495 additions & 22 deletions

File tree

‎admin/handler.go‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"github.com/codex2api/auth"
2525
"github.com/codex2api/cache"
2626
"github.com/codex2api/database"
27+
"github.com/codex2api/internal/imagestore"
2728
"github.com/codex2api/proxy"
2829
"github.com/codex2api/security"
2930
"github.com/codex2api/security/promptfilter"
@@ -140,6 +141,7 @@ func (h *Handler) RegisterRoutes(r *gin.Engine) {
140141
api.GET("/ops/overview", h.GetOpsOverview)
141142
api.GET("/settings", h.GetSettings)
142143
api.PUT("/settings", h.UpdateSettings)
144+
api.POST("/settings/image-storage/test", h.TestImageStorageConnection)
143145
api.GET("/prompt-filter/logs", h.ListPromptFilterLogs)
144146
api.DELETE("/prompt-filter/logs", h.ClearPromptFilterLogs)
145147
api.POST("/prompt-filter/test", h.TestPromptFilter)
@@ -2202,6 +2204,14 @@ type settingsResponse struct {
22022204
UsageLogFlushIntervalSeconds int `json:"usage_log_flush_interval_seconds"`
22032205
StreamFlushPolicy string `json:"stream_flush_policy"`
22042206
StreamFlushIntervalMS int `json:"stream_flush_interval_ms"`
2207+
ImageStorageBackend string `json:"image_storage_backend"`
2208+
ImageS3Endpoint string `json:"image_s3_endpoint"`
2209+
ImageS3Region string `json:"image_s3_region"`
2210+
ImageS3Bucket string `json:"image_s3_bucket"`
2211+
ImageS3AccessKey string `json:"image_s3_access_key"`
2212+
ImageS3SecretKey string `json:"image_s3_secret_key"`
2213+
ImageS3Prefix string `json:"image_s3_prefix"`
2214+
ImageS3ForcePathStyle bool `json:"image_s3_force_path_style"`
22052215
}
22062216

22072217
type updateSettingsReq struct {
@@ -2245,6 +2255,14 @@ type updateSettingsReq struct {
22452255
UsageLogFlushIntervalSeconds *int `json:"usage_log_flush_interval_seconds"`
22462256
StreamFlushPolicy *string `json:"stream_flush_policy"`
22472257
StreamFlushIntervalMS *int `json:"stream_flush_interval_ms"`
2258+
ImageStorageBackend *string `json:"image_storage_backend"`
2259+
ImageS3Endpoint *string `json:"image_s3_endpoint"`
2260+
ImageS3Region *string `json:"image_s3_region"`
2261+
ImageS3Bucket *string `json:"image_s3_bucket"`
2262+
ImageS3AccessKey *string `json:"image_s3_access_key"`
2263+
ImageS3SecretKey *string `json:"image_s3_secret_key"`
2264+
ImageS3Prefix *string `json:"image_s3_prefix"`
2265+
ImageS3ForcePathStyle *bool `json:"image_s3_force_path_style"`
22482266
}
22492267

22502268
// GetSettings 获取当前系统设置
@@ -2264,6 +2282,8 @@ func (h *Handler) GetSettings(c *gin.Context) {
22642282
}
22652283
promptFilterCfg := h.store.GetPromptFilterConfig()
22662284
runtimeCfg := proxy.CurrentRuntimeSettings()
2285+
imgCfg := imagestore.CurrentConfig()
2286+
imgPrefix := strings.TrimSuffix(imgCfg.Prefix, "/")
22672287
c.JSON(http.StatusOK, settingsResponse{
22682288
MaxConcurrency: h.store.GetMaxConcurrency(),
22692289
GlobalRPM: h.rateLimiter.GetRPM(),
@@ -2310,6 +2330,14 @@ func (h *Handler) GetSettings(c *gin.Context) {
23102330
UsageLogFlushIntervalSeconds: h.db.GetUsageLogFlushIntervalSeconds(),
23112331
StreamFlushPolicy: runtimeCfg.StreamFlushPolicy,
23122332
StreamFlushIntervalMS: runtimeCfg.StreamFlushIntervalMS,
2333+
ImageStorageBackend: imgCfg.Backend,
2334+
ImageS3Endpoint: imgCfg.Endpoint,
2335+
ImageS3Region: imgCfg.Region,
2336+
ImageS3Bucket: imgCfg.Bucket,
2337+
ImageS3AccessKey: imgCfg.AccessKey,
2338+
ImageS3SecretKey: imgCfg.SecretKey,
2339+
ImageS3Prefix: imgPrefix,
2340+
ImageS3ForcePathStyle: imgCfg.ForcePathStyle,
23132341
})
23142342
}
23152343

@@ -2655,6 +2683,57 @@ func (h *Handler) UpdateSettings(c *gin.Context) {
26552683
}
26562684
}
26572685

2686+
// 图片存储后端配置
2687+
imgCfg := imagestore.CurrentConfig()
2688+
imgChanged := false
2689+
if req.ImageStorageBackend != nil {
2690+
imgCfg.Backend = *req.ImageStorageBackend
2691+
imgChanged = true
2692+
}
2693+
if req.ImageS3Endpoint != nil {
2694+
imgCfg.Endpoint = *req.ImageS3Endpoint
2695+
imgChanged = true
2696+
}
2697+
if req.ImageS3Region != nil {
2698+
imgCfg.Region = *req.ImageS3Region
2699+
imgChanged = true
2700+
}
2701+
if req.ImageS3Bucket != nil {
2702+
imgCfg.Bucket = *req.ImageS3Bucket
2703+
imgChanged = true
2704+
}
2705+
if req.ImageS3AccessKey != nil {
2706+
imgCfg.AccessKey = *req.ImageS3AccessKey
2707+
imgChanged = true
2708+
}
2709+
if req.ImageS3SecretKey != nil {
2710+
imgCfg.SecretKey = *req.ImageS3SecretKey
2711+
imgChanged = true
2712+
}
2713+
if req.ImageS3Prefix != nil {
2714+
imgCfg.Prefix = *req.ImageS3Prefix
2715+
imgChanged = true
2716+
}
2717+
if req.ImageS3ForcePathStyle != nil {
2718+
imgCfg.ForcePathStyle = *req.ImageS3ForcePathStyle
2719+
imgChanged = true
2720+
}
2721+
imgCfg.LocalDir = imageAssetDir()
2722+
if imgChanged {
2723+
if err := imagestore.Configure(imgCfg); err != nil {
2724+
writeError(c, http.StatusBadRequest, "图片存储配置无效: "+err.Error())
2725+
return
2726+
}
2727+
// Configure 内部 Normalize 过，重新读出来用于持久化
2728+
imgCfg = imagestore.CurrentConfig()
2729+
log.Printf("设置已更新: image_storage_backend = %s", imgCfg.Backend)
2730+
}
2731+
imgConfigJSON, encodeErr := imagestore.EncodeConfigJSON(imgCfg)
2732+
if encodeErr != nil {
2733+
log.Printf("图片存储配置序列化失败: %v", encodeErr)
2734+
imgConfigJSON = "{}"
2735+
}
2736+
26582737
// 持久化保存到数据库
26592738
err := h.db.UpdateSystemSettings(c.Request.Context(), &database.SystemSettings{
26602739
MaxConcurrency: h.store.GetMaxConcurrency(),
@@ -2697,6 +2776,7 @@ func (h *Handler) UpdateSettings(c *gin.Context) {
26972776
UsageLogFlushIntervalSeconds: usageLogFlushIntervalSeconds,
26982777
StreamFlushPolicy: runtimeCfg.StreamFlushPolicy,
26992778
StreamFlushIntervalMS: runtimeCfg.StreamFlushIntervalMS,
2779+
ImageStorageConfig: imgConfigJSON,
27002780
})
27012781
if err != nil {
27022782
log.Printf("无法持久化保存设置: %v", err)
@@ -2762,9 +2842,63 @@ func (h *Handler) UpdateSettings(c *gin.Context) {
27622842
UsageLogFlushIntervalSeconds: usageLogFlushIntervalSeconds,
27632843
StreamFlushPolicy: runtimeCfg.StreamFlushPolicy,
27642844
StreamFlushIntervalMS: runtimeCfg.StreamFlushIntervalMS,
2845+
ImageStorageBackend: imgCfg.Backend,
2846+
ImageS3Endpoint: imgCfg.Endpoint,
2847+
ImageS3Region: imgCfg.Region,
2848+
ImageS3Bucket: imgCfg.Bucket,
2849+
ImageS3AccessKey: imgCfg.AccessKey,
2850+
ImageS3SecretKey: imgCfg.SecretKey,
2851+
ImageS3Prefix: strings.TrimSuffix(imgCfg.Prefix, "/"),
2852+
ImageS3ForcePathStyle: imgCfg.ForcePathStyle,
27652853
})
27662854
}
27672855

2856+
type testImageStorageReq struct {
2857+
Endpoint string `json:"endpoint"`
2858+
Region string `json:"region"`
2859+
Bucket string `json:"bucket"`
2860+
AccessKey string `json:"access_key"`
2861+
SecretKey string `json:"secret_key"`
2862+
Prefix string `json:"prefix"`
2863+
ForcePathStyle bool `json:"force_path_style"`
2864+
}
2865+
2866+
// TestImageStorageConnection 用提交的字段临时构造一次 S3Backend，调用 HeadBucket 验证可达性。
2867+
// 不修改任何持久化状态，便于"保存前先点测试连接"。
2868+
func (h *Handler) TestImageStorageConnection(c *gin.Context) {
2869+
var req testImageStorageReq
2870+
if err := c.ShouldBindJSON(&req); err != nil {
2871+
writeError(c, http.StatusBadRequest, "请求格式错误")
2872+
return
2873+
}
2874+
cfg := imagestore.Config{
2875+
Backend: imagestore.BackendS3,
2876+
Endpoint: req.Endpoint,
2877+
Region: req.Region,
2878+
Bucket: req.Bucket,
2879+
AccessKey: req.AccessKey,
2880+
SecretKey: req.SecretKey,
2881+
Prefix: req.Prefix,
2882+
ForcePathStyle: req.ForcePathStyle,
2883+
}.Normalize()
2884+
if err := cfg.Validate(); err != nil {
2885+
writeError(c, http.StatusBadRequest, err.Error())
2886+
return
2887+
}
2888+
backend, err := imagestore.NewS3Backend(cfg)
2889+
if err != nil {
2890+
writeError(c, http.StatusBadRequest, err.Error())
2891+
return
2892+
}
2893+
ctx, cancel := context.WithTimeout(c.Request.Context(), 10*time.Second)
2894+
defer cancel()
2895+
if err := backend.HeadBucket(ctx); err != nil {
2896+
writeError(c, http.StatusBadRequest, err.Error())
2897+
return
2898+
}
2899+
c.JSON(http.StatusOK, gin.H{"ok": true, "bucket": cfg.Bucket})
2900+
}
2901+
27682902
// ==================== 导出 & 迁移 ====================
27692903

27702904
type cpaExportEntry struct {

‎admin/image_studio.go‎

Lines changed: 72 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
_ "image/gif"
1212
_ "image/jpeg"
1313
_ "image/png"
14+
"io"
1415
"log"
1516
"mime"
1617
"net/http"
@@ -22,6 +23,7 @@ import (
2223

2324
"github.com/codex2api/database"
2425
"github.com/codex2api/internal/imageproc"
26+
"github.com/codex2api/internal/imagestore"
2527
"github.com/codex2api/internal/signedasset"
2628
"github.com/codex2api/proxy"
2729
"github.com/codex2api/security"
@@ -34,6 +36,9 @@ const defaultImageAssetDir = "/data/images"
3436
const maxInlineImageAssetCacheBytes = 64 * 1024 * 1024
3537
const defaultSignedImageThumbKB = 32
3638

39+
// thumbCache 跨请求复用缩略图。S3 后端读源图代价高，这里用 LRU 兜一下。
40+
var thumbCache = imagestore.NewThumbnailCache(0)
41+
3742
type imagePromptTemplatePayload struct {
3843
Name *string `json:"name"`
3944
Prompt *string `json:"prompt"`
@@ -362,12 +367,24 @@ func (h *Handler) attachImageJobAssetCachePayload(job *database.ImageGenerationJ
362367
if storagePath == "" {
363368
continue
364369
}
365-
info, err := os.Stat(storagePath)
366-
if err != nil || info.Size() <= 0 || info.Size() > maxInlineImageAssetCacheBytes {
370+
backend, err := imagestore.Resolve(storagePath)
371+
if err != nil {
367372
continue
368373
}
369-
data, err := os.ReadFile(storagePath)
370-
if err != nil {
374+
// 本地后端走 Stat 快速判断尺寸；S3 后端没有便宜的 Stat，直接尝试 Read 并按字节数过滤。
375+
if storedBytes := job.Assets[idx].Bytes; storedBytes > 0 && storedBytes > maxInlineImageAssetCacheBytes {
376+
continue
377+
}
378+
if !imagestore.IsS3Ref(storagePath) {
379+
info, statErr := os.Stat(storagePath)
380+
if statErr != nil || info.Size() <= 0 || info.Size() > maxInlineImageAssetCacheBytes {
381+
continue
382+
}
383+
}
384+
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
385+
data, err := backend.Read(ctx, storagePath)
386+
cancel()
387+
if err != nil || len(data) == 0 || len(data) > maxInlineImageAssetCacheBytes {
371388
continue
372389
}
373390
job.Assets[idx].CacheB64JSON = base64.StdEncoding.EncodeToString(data)
@@ -455,8 +472,8 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset
455472
writeError(c, http.StatusNotFound, "图片文件不存在")
456473
return
457474
}
458-
info, err := os.Stat(asset.StoragePath)
459-
if err != nil || info.IsDir() {
475+
backend, err := imagestore.Resolve(asset.StoragePath)
476+
if err != nil {
460477
writeError(c, http.StatusNotFound, "图片文件不存在")
461478
return
462479
}
@@ -477,8 +494,15 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset
477494
}
478495
filename := sanitizeDownloadFilename(asset.Filename)
479496
if opts.thumbKB > 0 && !opts.download {
480-
if data, err := os.ReadFile(asset.StoragePath); err == nil {
497+
cacheKey := imagestore.ThumbKey(asset.ID, opts.thumbKB)
498+
if data, contentType, ok := thumbCache.Get(cacheKey); ok {
499+
c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s"`, thumbnailFilename(filename)))
500+
c.Data(http.StatusOK, contentType, data)
501+
return
502+
}
503+
if data, err := backend.Read(c.Request.Context(), asset.StoragePath); err == nil {
481504
if thumb, contentType, ok := imageproc.MakeThumbnail(data, opts.thumbKB); ok {
505+
thumbCache.Put(cacheKey, contentType, thumb)
482506
c.Header("Content-Disposition", fmt.Sprintf(`inline; filename="%s"`, thumbnailFilename(filename)))
483507
c.Data(http.StatusOK, contentType, thumb)
484508
return
@@ -490,11 +514,34 @@ func (h *Handler) serveImageAssetFile(c *gin.Context, asset *database.ImageAsset
490514
c.Header("Content-Type", asset.MimeType)
491515
}
492516
c.Header("Content-Disposition", fmt.Sprintf(`%s; filename="%s"`, disposition, filename))
493-
c.File(asset.StoragePath)
517+
518+
rc, size, err := backend.Open(c.Request.Context(), asset.StoragePath)
519+
if err != nil {
520+
writeError(c, http.StatusNotFound, "图片文件不存在")
521+
return
522+
}
523+
defer rc.Close()
524+
contentType := strings.TrimSpace(asset.MimeType)
525+
if contentType == "" {
526+
contentType = "application/octet-stream"
527+
}
528+
if size < 0 {
529+
// S3 偶尔不带 Content-Length，退化为 chunked。
530+
c.Status(http.StatusOK)
531+
_, _ = c.Writer.Write(nil)
532+
_, _ = io.Copy(c.Writer, rc)
533+
return
534+
}
535+
c.DataFromReader(http.StatusOK, size, contentType, rc, nil)
494536
}
495537

496538
func imageAssetPathAllowed(storagePath string) bool {
497-
assetPath, err := filepath.Abs(strings.TrimSpace(storagePath))
539+
storagePath = strings.TrimSpace(storagePath)
540+
if imagestore.IsS3Ref(storagePath) {
541+
// S3 ref 不在本地目录约束内，由后端自身保证作用域。
542+
return true
543+
}
544+
assetPath, err := filepath.Abs(storagePath)
498545
if err != nil {
499546
return false
500547
}
@@ -569,7 +616,10 @@ func (h *Handler) DeleteImageAsset(c *gin.Context) {
569616
return
570617
}
571618
if asset.StoragePath != "" {
572-
_ = os.Remove(asset.StoragePath)
619+
if backend, err := imagestore.Resolve(asset.StoragePath); err == nil {
620+
_ = backend.Delete(ctx, asset.StoragePath)
621+
}
622+
thumbCache.Invalidate(asset.ID)
573623
}
574624
writeMessage(c, http.StatusOK, "已删除")
575625
}
@@ -774,9 +824,15 @@ func jpegFallbackImageJobRequest(req imageGenerationJobPayload) imageGenerationJ
774824
}
775825

776826
func (h *Handler) saveImageJobAssets(ctx context.Context, jobID int64, req imageGenerationJobPayload, responseJSON []byte) ([]database.ImageAsset, error) {
777-
dir := imageAssetDir()
778-
if err := os.MkdirAll(dir, 0o755); err != nil {
779-
return nil, fmt.Errorf("创建图库目录失败: %w", err)
827+
backend, err := imagestore.Primary()
828+
if err != nil {
829+
return nil, fmt.Errorf("图片存储未初始化: %w", err)
830+
}
831+
if backend.Name() == imagestore.BackendLocal {
832+
// LocalBackend 已在 Configure 时 mkdir，这里再保险一次以兼容 dir 在运行期被外部清理的情况。
833+
if err := os.MkdirAll(imageAssetDir(), 0o755); err != nil {
834+
return nil, fmt.Errorf("创建图库目录失败: %w", err)
835+
}
780836
}
781837
data := gjson.GetBytes(responseJSON, "data")
782838
if !data.IsArray() {
@@ -819,8 +875,8 @@ func (h *Handler) saveImageJobAssets(ctx context.Context, jobID int64, req image
819875
mimeType = "application/octet-stream"
820876
}
821877
filename := fmt.Sprintf("%d-%02d-%s.%s", jobID, idx+1, uuid.NewString()[:8], safeImageExtension(format, mimeType))
822-
storagePath := filepath.Join(dir, filename)
823-
if err := os.WriteFile(storagePath, imageBytes, 0o644); err != nil {
878+
storagePath, err := backend.Save(ctx, filename, imageBytes, mimeType)
879+
if err != nil {
824880
return saved, fmt.Errorf("保存图片失败: %w", err)
825881
}
826882
input := database.ImageAssetInput{
@@ -841,7 +897,7 @@ func (h *Handler) saveImageJobAssets(ctx context.Context, jobID int64, req image
841897
}
842898
assetID, err := h.db.InsertImageAsset(ctx, input)
843899
if err != nil {
844-
_ = os.Remove(storagePath)
900+
_ = backend.Delete(ctx, storagePath)
845901
return saved, err
846902
}
847903
asset, err := h.db.GetImageAsset(ctx, assetID)

‎admin/image_studio_test.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import (
1717
"github.com/codex2api/auth"
1818
"github.com/codex2api/cache"
1919
"github.com/codex2api/database"
20+
"github.com/codex2api/internal/imagestore"
2021
"github.com/gin-gonic/gin"
2122
)
2223

@@ -84,6 +85,9 @@ func TestSaveImageJobAssetsPersistsFilesAndMetadata(t *testing.T) {
8485
db := newTestAdminDB(t)
8586
dir := t.TempDir()
8687
t.Setenv("IMAGE_ASSET_DIR", dir)
88+
if err := imagestore.Configure(imagestore.Config{Backend: imagestore.BackendLocal, LocalDir: dir}); err != nil {
89+
t.Fatalf("imagestore.Configure: %v", err)
90+
}
8791
handler := &Handler{db: db}
8892

8993
jobID, err := db.InsertImageGenerationJob(context.Background(), database.ImageGenerationJobInput{Prompt: "a blue square"})
@@ -152,6 +156,9 @@ func TestImageAssetFileRouteRequiresAdminAuth(t *testing.T) {
152156
t.Fatalf("InsertImageGenerationJob 返回错误: %v", err)
153157
}
154158
dir := t.TempDir()
159+
if err := imagestore.Configure(imagestore.Config{Backend: imagestore.BackendLocal, LocalDir: dir}); err != nil {
160+
t.Fatalf("imagestore.Configure: %v", err)
161+
}
155162
pngBytes := tinyPNG(t)
156163
path := filepath.Join(dir, "asset.png")
157164
if err := os.WriteFile(path, pngBytes, 0o644); err != nil {

0 commit comments

Comments
 (0)