From 841832bb6585f5bdf8234641ccfdf6d9880d6ce5 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Sat, 16 Nov 2019 20:31:34 +0800
Subject: [PATCH] Add: context

---
 pkg/filesystem/filesystem.go | 11 ++++++-----
 pkg/filesystem/hook.go       | 13 ++++++++-----
 pkg/filesystem/validator.go  |  7 ++++---
 routers/controllers/file.go  | 10 +++++++++-
 service/file/upload.go       |  5 +++--
 5 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/pkg/filesystem/filesystem.go b/pkg/filesystem/filesystem.go
index 1041d05..4b8d687 100644
--- a/pkg/filesystem/filesystem.go
+++ b/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
 	}
diff --git a/pkg/filesystem/hook.go b/pkg/filesystem/hook.go
index e4e1616..c577c53 100644
--- a/pkg/filesystem/hook.go
+++ b/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
diff --git a/pkg/filesystem/validator.go b/pkg/filesystem/validator.go
index 1629bca..8d58652 100644
--- a/pkg/filesystem/validator.go
+++ b/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
diff --git a/routers/controllers/file.go b/routers/controllers/file.go
index 15ac017..2567e02 100644
--- a/routers/controllers/file.go
+++ b/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))
diff --git a/service/file/upload.go b/service/file/upload.go
index 9478c5e..498925b 100644
--- a/service/file/upload.go
+++ b/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)
 	}