Add: context

pkg/filesystem/filesystem.go

@@ -1,6 +1,7 @@
 package filesystem
 
 import (
+	"context"
 	"github.com/HFO4/cloudreve/models"
 	"io"
 )
@@ -25,11 +26,11 @@ type FileSystem struct {
 	  钩子函数
 	*/
 	// 上传文件前
-	BeforeUpload func(fs *FileSystem, file FileData) error
+	BeforeUpload func(ctx context.Context, fs *FileSystem, file FileData) error
 	// 上传文件后
-	AfterUpload func(fs *FileSystem) error
+	AfterUpload func(ctx context.Context, fs *FileSystem) error
 	// 文件验证失败后
-	ValidateFailed func(fs *FileSystem) error
+	ValidateFailed func(ctx context.Context, fs *FileSystem) error
 
 	/*
 	  文件系统处理适配器
@@ -45,8 +46,8 @@ func NewFileSystem(user *model.User) *FileSystem {
 }
 
 // Upload 上传文件
-func (fs *FileSystem) Upload(file FileData) (err error) {
-	err = fs.BeforeUpload(fs, file)
+func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) {
+	err = fs.BeforeUpload(ctx, fs, file)
 	if err != nil {
 		return err
 	}

pkg/filesystem/hook.go

@@ -1,21 +1,24 @@
 package filesystem
 
-import "errors"
+import (
+	"context"
+	"errors"
+)
 
 // GenericBeforeUpload 通用上传前处理钩子，包含数据库操作
-func GenericBeforeUpload(fs *FileSystem, file FileData) error {
+func GenericBeforeUpload(ctx context.Context, fs *FileSystem, file FileData) error {
 	// 验证单文件尺寸
-	if !fs.ValidateFileSize(file.GetSize()) {
+	if !fs.ValidateFileSize(ctx, file.GetSize()) {
 		return errors.New("单个文件尺寸太大")
 	}
 
 	// 验证并扣除容量
-	if !fs.ValidateCapacity(file.GetSize()) {
+	if !fs.ValidateCapacity(ctx, file.GetSize()) {
 		return errors.New("容量空间不足")
 	}
 
 	// 验证扩展名
-	if !fs.ValidateExtension(file.GetFileName()) {
+	if !fs.ValidateExtension(ctx, file.GetFileName()) {
 		return errors.New("不允许上传此类型的文件")
 	}
 	return nil

pkg/filesystem/validator.go

@@ -1,17 +1,18 @@
 package filesystem
 
 import (
+	"context"
 	"github.com/HFO4/cloudreve/pkg/util"
 	"path/filepath"
 )
 
 // ValidateFileSize 验证上传的文件大小是否超出限制
-func (fs *FileSystem) ValidateFileSize(size uint64) bool {
+func (fs *FileSystem) ValidateFileSize(ctx context.Context, size uint64) bool {
 	return size <= fs.User.Policy.MaxSize
 }
 
 // ValidateCapacity 验证并扣除用户容量
-func (fs *FileSystem) ValidateCapacity(size uint64) bool {
+func (fs *FileSystem) ValidateCapacity(ctx context.Context, size uint64) bool {
 	if fs.User.DeductionCapacity(size) {
 		return true
 	}
@@ -19,7 +20,7 @@ func (fs *FileSystem) ValidateCapacity(size uint64) bool {
 }
 
 // ValidateExtension 验证文件扩展名
-func (fs *FileSystem) ValidateExtension(fileName string) bool {
+func (fs *FileSystem) ValidateExtension(ctx context.Context, fileName string) bool {
 	// 不需要验证
 	if len(fs.User.Policy.OptionsSerialized.FileType) == 0 {
 		return true

routers/controllers/file.go

@@ -1,6 +1,7 @@
 package controllers
 
 import (
+	"context"
 	"github.com/HFO4/cloudreve/models"
 	"github.com/HFO4/cloudreve/pkg/serializer"
 	"github.com/HFO4/cloudreve/service/file"
@@ -16,10 +17,17 @@ func FileUpload(c *gin.Context) {
 		return
 	}
 
+	var (
+		ctx    context.Context
+		cancel context.CancelFunc
+	)
+	ctx, cancel = context.WithCancel(context.Background())
+
 	var service file.UploadService
+	defer cancel()
 
 	if err := c.ShouldBind(&service); err == nil {
-		res := service.Upload(c)
+		res := service.Upload(ctx, c)
 		c.JSON(200, res)
 	} else {
 		c.JSON(200, ErrorResponse(err))

service/file/upload.go

@@ -1,6 +1,7 @@
 package file
 
 import (
+	"context"
 	"github.com/HFO4/cloudreve/models"
 	"github.com/HFO4/cloudreve/pkg/filesystem"
 	"github.com/HFO4/cloudreve/pkg/filesystem/local"
@@ -17,7 +18,7 @@ type UploadService struct {
 }
 
 // Upload 处理本地策略小文件上传
-func (service *UploadService) Upload(c *gin.Context) serializer.Response {
+func (service *UploadService) Upload(ctx context.Context, c *gin.Context) serializer.Response {
 	// TODO 检查文件大小是否小于分片最大大小
 	file, err := service.File.Open()
 	if err != nil {
@@ -36,7 +37,7 @@ func (service *UploadService) Upload(c *gin.Context) serializer.Response {
 		User:         user.(*model.User),
 	}
 
-	err = fs.Upload(fileData)
+	err = fs.Upload(ctx, fileData)
 	if err != nil {
 		return serializer.Err(serializer.CodeUploadFailed, err.Error(), err)
 	}