fix(upload): 分片上传 - 先预留配额并原子提交容量 (#679)

Co-authored-by: ddcat666 <26564566+ddcat666@users.noreply.github.com>

backend/internal/app/chunk_quota_security_test.go

@@ -0,0 +1,58 @@
+package app
+
+import (
+	"bytes"
+	"testing"
+	"time"
+)
+
+func TestChunkQuotaRejectsSingleFileLimit(t *testing.T) {
+	svc := newResourceTestService(t)
+	if _, err := svc.reserveChunkedUploadQuota("user-1", megabytes(defaultRuntimePolicy().Resource.ResourceUploadMB)); err == nil {
+		t.Fatal("chunked upload accepted a file at the exclusive single-file limit")
+	}
+}
+
+func TestChunkQuotaReservedUploadConsumesOnce(t *testing.T) {
+	svc := newResourceTestService(t)
+	if err := svc.ReserveChunkUploadSession("user", "session", 7, time.Now().Add(time.Hour), 32); err != nil {
+		t.Fatal(err)
+	}
+	day := time.Now().UTC().Format("2006-01-02")
+	if used, _ := svc.repo.DailyUploadBytes("user", day); used != 7 {
+		t.Fatalf("quota before receiving bytes=%d", used)
+	}
+	resource, err := svc.UploadReservedResourceFile("user", "session", "test.txt", 7, "file", 0, 0, 0, bytes.NewReader([]byte("payload")), "logical-key")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if resource.Size != 7 {
+		t.Fatalf("size=%d", resource.Size)
+	}
+	if used, _ := svc.repo.DailyUploadBytes("user", day); used != 7 {
+		t.Fatalf("quota after completion=%d", used)
+	}
+	if err := svc.ReserveChunkUploadSession("user", "retry", 7, time.Now().Add(time.Hour), 32); err != nil {
+		t.Fatal(err)
+	}
+	reused, err := svc.UploadReservedResourceFile("user", "retry", "test.txt", 7, "file", 0, 0, 0, bytes.NewReader([]byte("payload")), "logical-key")
+	if err != nil || reused.ID != resource.ID {
+		t.Fatalf("retry=%v err=%v", reused, err)
+	}
+	if used, _ := svc.repo.DailyUploadBytes("user", day); used != 7 {
+		t.Fatalf("quota after replay=%d", used)
+	}
+}
+
+func TestChunkQuotaFailedCompletionReleases(t *testing.T) {
+	svc := newResourceTestService(t)
+	if err := svc.ReserveChunkUploadSession("user", "failed", 7, time.Now().Add(time.Hour), 32); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := svc.UploadReservedResourceFile("user", "failed", "test.txt", 7, "file", 0, 0, 0, nil); err == nil {
+		t.Fatal("nil file accepted")
+	}
+	if used, _ := svc.repo.DailyUploadBytes("user", time.Now().UTC().Format("2006-01-02")); used != 0 {
+		t.Fatalf("failed upload retained quota=%d", used)
+	}
+}

backend/internal/app/resource_oss_fallback_test.go

@@ -21,6 +21,7 @@ func newResourceFallbackTestService(t *testing.T) (*Service, *gorm.DB) {
 	}
 	if err := db.AutoMigrate(
 		&model.Resource{},
+		&model.UploadReservation{},
 		&model.UserOSSSetting{},
 		&model.SystemSetting{},
 		&model.StorageLocation{},

backend/internal/app/resource_store.go

@@ -11,6 +11,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log"
 	"mime/multipart"
 	"net/http"
 	"os"
@@ -79,9 +80,24 @@ func (s *Service) uploadResource(userID string, header *multipart.FileHeader, ki
 }
 
 // UploadResourceFile 接收已完整落盘的本地文件（分片上传合并后调用）。
-// 它与 UploadResource 共享资源幂等、媒体探测、配额和持久化语义，唯一差异是分片会话已在 handler 校验单文件上限，
-// 因而此处不再重复该上限检查；uploadIdentity 用于跨请求重试时复用同一逻辑资源，避免重复对象。
+// 它与 UploadResource 共享资源幂等、媒体探测、配额和持久化语义。
 func (s *Service) UploadResourceFile(userID string, fileName string, size int64, kind string, width int, height int, durationMs int64, file io.ReadSeeker, uploadIdentity ...string) (*model.Resource, error) {
+	return s.uploadResourceFile(userID, "", fileName, size, kind, width, height, durationMs, file, uploadIdentity...)
+}
+
+func (s *Service) UploadReservedResourceFile(userID, reservationID, fileName string, size int64, kind string, width, height int, durationMs int64, file io.ReadSeeker, uploadIdentity ...string) (*model.Resource, error) {
+	if reservationID == "" {
+		return nil, BadAuthRequest("上传会话不存在")
+	}
+	defer func() {
+		if err := s.ReleaseChunkUploadSession(userID, reservationID); err != nil {
+			log.Printf("release upload session failed: %v", err)
+		}
+	}()
+	return s.uploadResourceFile(userID, reservationID, fileName, size, kind, width, height, durationMs, file, uploadIdentity...)
+}
+
+func (s *Service) uploadResourceFile(userID, reservationID, fileName string, size int64, kind string, width int, height int, durationMs int64, file io.ReadSeeker, uploadIdentity ...string) (*model.Resource, error) {
 	if file == nil || size <= 0 {
 		return nil, BadAuthRequest("请选择要上传的文件")
 	}
@@ -99,8 +115,17 @@ func (s *Service) UploadResourceFile(userID string, fileName string, size int64,
 	}
 	kind = uploadedResourceKind(mimeType)
 	if existing != nil {
+		if reservationID != "" {
+			if err := s.ReleaseChunkUploadSession(userID, reservationID); err != nil {
+				return nil, err
+			}
+		}
 		return s.retryStoredResource(userID, existing, kind, mimeType, size, file)
 	}
+	if reservationID != "" {
+		resource, _, err := s.storeResource(userID, kind, fileName, mimeType, size, width, height, durationMs, file, uploadKey, false, reservationID)
+		return resource, err
+	}
 	day, err := s.reserveChunkedUploadQuota(userID, size)
 	if err != nil {
 		return nil, err
@@ -253,11 +278,15 @@ func (s *Service) openResourceRange(userID string, resource *model.Resource, ran
 	return &ResourceStream{Resource: resource, Body: stream.Body, StatusCode: stream.StatusCode, ContentLength: stream.ContentLength, ContentRange: stream.ContentRange, AcceptRanges: stream.AcceptRanges}, nil
 }
 
-func (s *Service) storeResource(userID string, kind string, fileName string, mimeType string, size int64, width int, height int, durationMs int64, body io.Reader, uploadKey *string, forceLocal bool) (*model.Resource, bool, error) {
-	return s.storeResourceWithWriter(userID, kind, fileName, mimeType, size, width, height, durationMs, body, uploadKey, forceLocal, s.storeResourceObject)
+func (s *Service) storeResource(userID string, kind string, fileName string, mimeType string, size int64, width int, height int, durationMs int64, body io.Reader, uploadKey *string, forceLocal bool, reservationIDs ...string) (*model.Resource, bool, error) {
+	return s.storeResourceWithWriter(userID, kind, fileName, mimeType, size, width, height, durationMs, body, uploadKey, forceLocal, s.storeResourceObject, reservationIDs...)
 }
 
-func (s *Service) storeResourceWithWriter(userID string, kind string, fileName string, mimeType string, size int64, width int, height int, durationMs int64, body io.Reader, uploadKey *string, forceLocal bool, writeObject func(*model.Resource, string, io.Reader) (string, error)) (*model.Resource, bool, error) {
+func (s *Service) storeResourceWithWriter(userID string, kind string, fileName string, mimeType string, size int64, width int, height int, durationMs int64, body io.Reader, uploadKey *string, forceLocal bool, writeObject func(*model.Resource, string, io.Reader) (string, error), reservationIDs ...string) (*model.Resource, bool, error) {
+	reservationID := ""
+	if len(reservationIDs) > 0 {
+		reservationID = reservationIDs[0]
+	}
 	// Live2D is reserved for the administrator's validated local import path.
 	if !(forceLocal && kind == "live2d") {
 		kind = normalizeResourceKind(kind, mimeType)
@@ -321,7 +350,7 @@ func (s *Service) storeResourceWithWriter(userID string, kind string, fileName s
 	}
 	resource.Status = model.ResourceStatusReady
 	resource.ETag = etag
-	if err := s.repo.SaveResource(&resource); err != nil {
+	if err := s.saveResourceWithinStorageLimit(&resource, reservationID); err != nil {
 		cleanupErr := s.deleteStoredResourceObject(userID, &resource)
 		if cleanupErr == nil {
 			if deleteErr := s.repo.DeleteResource(userID, resource.ID); deleteErr != nil {
@@ -435,7 +464,7 @@ func (s *Service) retryStoredResource(userID string, resource *model.Resource, k
 	}
 	resource.Status = model.ResourceStatusReady
 	resource.ETag = etag
-	if err := s.repo.SaveResource(resource); err != nil {
+	if err := s.saveResourceWithinStorageLimit(resource, ""); err != nil {
 		s.releaseRetryUploadQuota(userID, day, size)
 		resource.Status = model.ResourceStatusFailed
 		resource.Error = "保存资源重试就绪状态失败"

backend/internal/app/resource_test.go

@@ -934,7 +934,7 @@ func newResourceTestService(t *testing.T) *Service {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if err := db.AutoMigrate(&model.SystemSetting{}, &model.UserOSSSetting{}, &model.StorageLocation{}, &model.UserDailyUploadUsage{}, &model.Resource{}); err != nil {
+	if err := db.AutoMigrate(&model.SystemSetting{}, &model.UserOSSSetting{}, &model.StorageLocation{}, &model.UserDailyUploadUsage{}, &model.Resource{}, &model.UploadReservation{}); err != nil {
 		t.Fatal(err)
 	}
 	return &Service{repo: repository.New(db), dataDir: t.TempDir()}

backend/internal/app/task_media_storage.go

@@ -266,7 +266,7 @@ func (s *Service) storeTaskMediaFile(task *model.Task, index int, path, mimeType
 		resource.ETag, err = s.storeTaskMediaObject(resource, "generated."+extensionFromMimeType(mimeType), file)
 		if err == nil {
 			resource.Status, resource.Error, resource.UpdatedAt = model.ResourceStatusReady, "", time.Now()
-			err = s.repo.SaveResource(resource)
+			err = s.saveResourceWithinStorageLimit(resource, "")
 		}
 	}
 	if err != nil {

backend/internal/app/upload_quota.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"time"
 
+	"infinite-canvas/backend/internal/model"
 	"infinite-canvas/backend/internal/repository"
 )
 
@@ -18,14 +19,49 @@ func (s *Service) reserveUserUploadQuota(userID string, size int64) (string, err
 	return s.reserveUserStoredFileQuota(userID, size, megabytes(policy.Resource.ResourceUploadMB), megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), fmt.Sprintf("单个上传文件必须小于 %dMB", policy.Resource.ResourceUploadMB))
 }
 
-// reserveChunkedUploadQuota 用于分片上传完成时预留额度：单文件上限对分片会话不适用（每片已独立校验），
-// 仅受“今日上传”与“账号存储总量”约束；singleFileLimit 传 size+1 使单文件上限永不命中。
+// 分片和普通上传使用相同的单文件上限。
 func (s *Service) reserveChunkedUploadQuota(userID string, size int64) (string, error) {
+	return s.reserveUserUploadQuota(userID, size)
+}
+
+func (s *Service) ReserveChunkUploadSession(userID, id string, size int64, expires time.Time, maxSessions int) error {
 	policy, err := s.RuntimePolicy()
 	if err != nil {
-		return "", err
+		return err
 	}
-	return s.reserveUserStoredFileQuota(userID, size, size+1, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), "")
+	if size <= 0 || size >= megabytes(policy.Resource.ResourceUploadMB) {
+		return BadAuthRequest(fmt.Sprintf("单个上传文件必须小于 %dMB", policy.Resource.ResourceUploadMB))
+	}
+	if !expires.After(time.Now()) || maxSessions <= 0 {
+		return BadAuthRequest("上传会话参数无效")
+	}
+	err = s.repo.ReserveUploadSession(&model.UploadReservation{ID: id, UserID: userID, Size: size, Day: time.Now().UTC().Format("2006-01-02"), ExpiresAt: expires}, megabytes(policy.Resource.DailyUploadMB), gigabytes(policy.Resource.StoredFileGB), maxSessions)
+	return uploadReservationError(err)
+}
+
+func (s *Service) ReleaseChunkUploadSession(userID, id string) error {
+	return s.repo.ReleaseUploadSession(userID, id)
+}
+
+func uploadReservationError(err error) error {
+	if errors.Is(err, repository.ErrUploadSessionLimit) {
+		return NewAppError(429, err.Error())
+	}
+	if errors.Is(err, repository.ErrUploadStorageLimit) || errors.Is(err, repository.ErrDailyUploadLimitExceeded) {
+		return QuotaExceeded(err.Error())
+	}
+	if errors.Is(err, repository.ErrUploadReservationExpired) {
+		return BadAuthRequest(err.Error())
+	}
+	return err
+}
+
+func (s *Service) saveResourceWithinStorageLimit(resource *model.Resource, reservationID string) error {
+	policy, err := s.RuntimePolicy()
+	if err != nil {
+		return err
+	}
+	return uploadReservationError(s.repo.SaveResourceWithinStorageLimit(resource, gigabytes(policy.Resource.StoredFileGB), megabytes(policy.Resource.DailyUploadMB), reservationID))
 }
 
 func (s *Service) reserveGeneratedResourceQuota(userID string, size int64) (string, error) {