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

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

backend/internal/handler/chunk_quota_security_test.go

@@ -0,0 +1,108 @@
+package handler
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"infinite-canvas/backend/internal/auth"
+	"infinite-canvas/backend/internal/model"
+	"infinite-canvas/backend/internal/platform"
+	"infinite-canvas/backend/internal/repository"
+	"infinite-canvas/backend/internal/service"
+)
+
+func TestChunkAdmissionReservesBeforeBody(t *testing.T) {
+	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	sqlDB, _ := db.DB()
+	sqlDB.SetMaxOpenConns(1)
+	t.Cleanup(func() { _ = sqlDB.Close() })
+	if err := db.AutoMigrate(&model.User{}, &model.AuthSession{}, &model.SystemSetting{}, &model.Resource{}, &model.UploadReservation{}, &model.UserDailyUploadUsage{}, &model.StorageLocation{}, &model.UserOSSSetting{}, &model.UserDailyActivity{}); err != nil {
+		t.Fatal(err)
+	}
+	if err := db.Create(&model.User{ID: "chunk-user", Username: "chunk-user", Status: model.UserStatusActive}).Error; err != nil {
+		t.Fatal(err)
+	}
+	if err := db.Create(&model.AuthSession{ID: "chunk-session", UserID: "chunk-user", TokenHash: auth.HashToken("test-token"), ExpiresAt: time.Now().Add(time.Hour)}).Error; err != nil {
+		t.Fatal(err)
+	}
+	policy := platform.DefaultRuntimePolicy()
+	policy.Resource.ResourceUploadMB = 1
+	data, _ := json.Marshal(policy)
+	if err := db.Create(&model.SystemSetting{Key: "runtime_policy", ValueJSON: string(data)}).Error; err != nil {
+		t.Fatal(err)
+	}
+	repo := repository.New(db)
+	svc := service.New(repo, t.TempDir())
+	previous := runtimeService
+	ConfigureRuntime(svc)
+	t.Cleanup(func() { runtimeService = previous })
+	gin.SetMode(gin.TestMode)
+	router := gin.New()
+	RegisterChunkedUploadRoutes(router.Group("/api"), svc)
+	call := func(method, path, body string) *httptest.ResponseRecorder {
+		req := httptest.NewRequest(method, path, strings.NewReader(body))
+		req.Header.Set("Content-Type", "application/json")
+		req.AddCookie(&http.Cookie{Name: service.SessionCookieName, Value: "chunk-session.test-token"})
+		w := httptest.NewRecorder()
+		router.ServeHTTP(w, req)
+		return w
+	}
+	w := call("POST", "/api/resources/uploads", `{"fileName":"large.txt","size":1048576}`)
+	if w.Code != 400 {
+		t.Fatalf("oversize admission status=%d body=%s", w.Code, w.Body.String())
+	}
+	w = call("POST", "/api/resources/uploads", `{"fileName":"test.txt","size":7}`)
+	if w.Code != 200 {
+		t.Fatalf("admission status=%d body=%s", w.Code, w.Body.String())
+	}
+	var response struct {
+		Data struct {
+			UploadID string `json:"uploadId"`
+		} `json:"data"`
+	}
+	if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
+		t.Fatal(err)
+	}
+	id := response.Data.UploadID
+	t.Cleanup(func() {
+		if sess := takeChunkSession(id); sess != nil {
+			sess.mu.Lock()
+			dropChunkSession(id)
+			sess.mu.Unlock()
+		}
+	})
+	if used, _ := repo.DailyUploadBytes("chunk-user", time.Now().UTC().Format("2006-01-02")); used != 7 {
+		t.Fatalf("before first chunk daily bytes=%d want 7", used)
+	}
+	var count int64
+	if err := db.Model(&model.UploadReservation{}).Where("id = ?", id).Count(&count).Error; err != nil || count != 1 {
+		t.Fatalf("admission reservation count=%d err=%v", count, err)
+	}
+	w = call("PUT", fmt.Sprintf("/api/resources/uploads/%s/chunks/0", id), "payload")
+	if w.Code != 200 {
+		t.Fatalf("chunk status=%d body=%s", w.Code, w.Body.String())
+	}
+	w = call("POST", fmt.Sprintf("/api/resources/uploads/%s/complete", id), "")
+	if w.Code != 200 || !bytes.Contains(w.Body.Bytes(), []byte(`"resource"`)) {
+		t.Fatalf("complete status=%d body=%s", w.Code, w.Body.String())
+	}
+	if err := db.Model(&model.UploadReservation{}).Count(&count).Error; err != nil || count != 0 {
+		t.Fatalf("remaining reservations=%d err=%v", count, err)
+	}
+	w = call("POST", fmt.Sprintf("/api/resources/uploads/%s/complete", id), "")
+	if w.Code != 404 {
+		t.Fatalf("duplicate completion status=%d", w.Code)
+	}
+}

backend/internal/handler/resource_upload_session.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"os"
 	"path/filepath"
@@ -18,8 +19,8 @@ import (
 )
 
 // 分片上传会话：把“导入本地媒体”拆成 开始→逐片→合并 三段，单片上限 8MB，
-// 文件整体不再受 multipart 单请求大小限制（对齐 Concat 桌面端“任意大小直接入库”的体验）。
-// 会话状态只存在内存（重启即失效 → 前端整传重试），磁盘暂存在系统临时目录，随会话清理。
+// 整体文件仍遵守资源单文件上限；收片前在数据库预留日额度和存储容量。
+// 分片路由需要粘性路由，内存会话失效后重传；配额预留跨实例共享并按时间到期。
 const (
 	chunkUploadChunkSize   = 8 << 20
 	chunkUploadSlackBytes  = 64 << 10 // MaxBytesReader 允许的超片余量
@@ -29,6 +30,9 @@ const (
 )
 
 type chunkedUploadSession struct {
+	mu             sync.Mutex
+	closed         bool
+	release        func() error
 	ID             string
 	UserID         string
 	FileName       string
@@ -87,9 +91,16 @@ func removeExpiredChunkSessions() {
 	chunkUploadSessions.Lock()
 	defer chunkUploadSessions.Unlock()
 	for id, sess := range chunkUploadSessions.m {
-		if now.Sub(sess.CreatedAt) > chunkUploadTTL {
+		if now.Sub(sess.CreatedAt) > chunkUploadTTL && sess.mu.TryLock() {
+			sess.closed = true
+			if sess.release != nil {
+				if err := sess.release(); err != nil {
+					log.Printf("release expired upload session failed: %v", err)
+				}
+			}
 			_ = os.RemoveAll(sess.Dir)
 			delete(chunkUploadSessions.m, id)
+			sess.mu.Unlock()
 		}
 	}
 }
@@ -105,6 +116,13 @@ func dropChunkSession(id string) {
 	chunkUploadSessions.Lock()
 	defer chunkUploadSessions.Unlock()
 	if sess := chunkUploadSessions.m[id]; sess != nil {
+		// Caller holds sess.mu. Expiry uses TryLock to avoid lock inversion.
+		sess.closed = true
+		if sess.release != nil {
+			if err := sess.release(); err != nil {
+				log.Printf("release upload session failed: %v", err)
+			}
+		}
 		_ = os.RemoveAll(sess.Dir)
 		delete(chunkUploadSessions.m, id)
 	}
@@ -144,31 +162,24 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			fail(c, http.StatusBadRequest, fmt.Errorf("文件大小必须大于 0"))
 			return
 		}
-		// 超账号存储总量的文件无论如何都会失败，提前给出明确提示。
-		if policy.Resource.StoredFileGB > 0 && req.Size > int64(policy.Resource.StoredFileGB)<<30 {
-			fail(c, http.StatusBadRequest, fmt.Errorf("文件超过账号存储总量上限 %dGB", policy.Resource.StoredFileGB))
-			return
-		}
-		// 同一用户并发会话数兜底，防内存占用失控。
-		active := 0
-		chunkUploadSessions.Lock()
-		for _, sess := range chunkUploadSessions.m {
-			if sess.UserID == user.ID {
-				active++
-			}
-		}
-		chunkUploadSessions.Unlock()
-		if active >= chunkUploadMaxPerUser {
-			fail(c, http.StatusTooManyRequests, fmt.Errorf("同时进行中的上传过多，请稍后重试"))
+		removeExpiredChunkSessions()
+		id, createdAt := newUploadSessionID(), time.Now()
+		if err := svc.ReserveChunkUploadSession(user.ID, id, req.Size, createdAt.Add(chunkUploadTTL), chunkUploadMaxPerUser); err != nil {
+			failService(c, err)
 			return
 		}
+		release := func() error { return svc.ReleaseChunkUploadSession(user.ID, id) }
 		dir, err := os.MkdirTemp("", "canvas-chunk-upload-*")
 		if err != nil {
+			if releaseErr := release(); releaseErr != nil {
+				log.Printf("release upload session failed: %v", releaseErr)
+			}
 			failService(c, err)
 			return
 		}
 		session := &chunkedUploadSession{
-			ID:             newUploadSessionID(),
+			ID:             id,
+			release:        release,
 			UserID:         user.ID,
 			FileName:       req.FileName,
 			Kind:           req.Kind,
@@ -179,7 +190,7 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			IdempotencyKey: req.IdempotencyKey,
 			ChunkCount:     int((req.Size + chunkUploadChunkSize - 1) / chunkUploadChunkSize),
 			Dir:            dir,
-			CreatedAt:      time.Now(),
+			CreatedAt:      createdAt,
 		}
 		chunkUploadSessions.Lock()
 		chunkUploadSessions.m[session.ID] = session
@@ -199,6 +210,12 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			return
 		}
 		index, err := strconv.Atoi(c.Param("index"))
+		session.mu.Lock()
+		defer session.mu.Unlock()
+		if session.closed || time.Since(session.CreatedAt) >= chunkUploadTTL {
+			fail(c, http.StatusNotFound, fmt.Errorf("上传会话不存在或已过期，请重新导入"))
+			return
+		}
 		if err != nil || index < 0 || index >= session.ChunkCount {
 			fail(c, http.StatusBadRequest, fmt.Errorf("非法的分片序号"))
 			return
@@ -224,7 +241,7 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			fail(c, http.StatusBadRequest, fmt.Errorf("分片 %d 上传不完整，请重试", index))
 			return
 		}
-		if readErr == nil || extra > 0 {
+		if readErr != io.EOF || extra > 0 {
 			_ = os.Remove(session.chunkPath(index))
 			fail(c, http.StatusBadRequest, fmt.Errorf("分片 %d 超过大小限制", index))
 			return
@@ -245,6 +262,12 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			fail(c, http.StatusNotFound, fmt.Errorf("上传会话不存在或已过期，请重新导入"))
 			return
 		}
+		session.mu.Lock()
+		defer session.mu.Unlock()
+		if session.closed || time.Since(session.CreatedAt) >= chunkUploadTTL {
+			fail(c, http.StatusNotFound, fmt.Errorf("上传会话不存在或已过期，请重新导入"))
+			return
+		}
 		if !session.hasAllChunks() {
 			fail(c, http.StatusBadRequest, fmt.Errorf("上传文件不完整，请重新导入"))
 			return
@@ -287,7 +310,7 @@ func RegisterChunkedUploadRoutes(r *gin.RouterGroup, svc *service.Service) {
 			return
 		}
 		defer fh.Close()
-		resource, svcErr := svc.UploadResourceFile(user.ID, session.FileName, session.Size, session.Kind, session.Width, session.Height, session.DurationMs, fh, session.IdempotencyKey)
+		resource, svcErr := svc.UploadReservedResourceFile(user.ID, session.ID, session.FileName, session.Size, session.Kind, session.Width, session.Height, session.DurationMs, fh, session.IdempotencyKey)
 		// 无论成败都结束会话：失败时前端会整传重试，不需要保留残片。
 		dropChunkSession(id)
 		if svcErr != nil {