From aa17aa8e6a3b80584c269dd4d930c677c5ca626a Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Mon, 18 Nov 2019 19:09:56 +0800
Subject: [PATCH] Feat: after uploading hooks and checks

---
 models/folder.go                              |  7 ++++
 models/folder_test.go                         | 20 +++++++++
 models/policy.go                              |  5 +++
 models/user.go                                |  6 +--
 pkg/filesystem/context.go                     | 12 ++++++
 pkg/filesystem/filesystem.go                  | 41 ++++++++++++++-----
 pkg/filesystem/{hook.go => hooks.go}          | 22 +++++++++-
 .../{hook_test.go => hooks_test.go}           |  0
 pkg/filesystem/validator.go                   |  6 +++
 routers/controllers/file.go                   |  2 +-
 10 files changed, 105 insertions(+), 16 deletions(-)
 create mode 100644 models/folder_test.go
 create mode 100644 pkg/filesystem/context.go
 rename pkg/filesystem/{hook.go => hooks.go} (66%)
 rename pkg/filesystem/{hook_test.go => hooks_test.go} (100%)

diff --git a/models/folder.go b/models/folder.go
index f7f9481..f2828ef 100644
--- a/models/folder.go
+++ b/models/folder.go
@@ -15,3 +15,10 @@ type Folder struct {
 	// 关联模型
 	OptionsSerialized PolicyOption `gorm:"-"`
 }
+
+// GetFolderByPath 根据绝对路径和UID查找目录
+func GetFolderByPath(path string, uid uint) (Folder, error) {
+	var folder Folder
+	result := DB.Where("owner = ? AND position_absolute = ?", uid, path).Find(&folder)
+	return folder, result.Error
+}
diff --git a/models/folder_test.go b/models/folder_test.go
new file mode 100644
index 0000000..3119339
--- /dev/null
+++ b/models/folder_test.go
@@ -0,0 +1,20 @@
+package model
+
+import (
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestGetFolderByPath(t *testing.T) {
+	asserts := assert.New(t)
+
+	//policyRows := sqlmock.NewRows([]string{"id", "name"}).
+	//	AddRow(1, "默认上传策略")
+	//mock.ExpectQuery("^SELECT (.+)").WillReturnRows(policyRows)
+
+	folder,_ := GetFolderByPath("/测试/test",1)
+	fmt.Println(folder)
+	asserts.NoError(mock.ExpectationsWereMet())
+	asserts.NoError(mock.)
+}
diff --git a/models/policy.go b/models/policy.go
index 9fd6401..40de2db 100644
--- a/models/policy.go
+++ b/models/policy.go
@@ -85,6 +85,11 @@ func (policy *Policy) GeneratePath(uid uint) string {
 
 // GenerateFileName 生成存储文件名
 func (policy *Policy) GenerateFileName(uid uint, origin string) string {
+	// 未开启自动重命名时，直接返回原始文件名
+	if !policy.AutoRename {
+		return origin
+	}
+
 	fileRule := policy.FileNameRule
 
 	replaceTable := map[string]string{
diff --git a/models/user.go b/models/user.go
index 2b704fe..6cb0ea7 100644
--- a/models/user.go
+++ b/models/user.go
@@ -40,7 +40,7 @@ type User struct {
 	Options       string `json:"-",gorm:"size:4096"`
 
 	// 关联模型
-	Group  Group
+	Group  Group  `gorm:"association_autoupdate:false"`
 	Policy Policy `gorm:"PRELOAD:false,association_autoupdate:false"`
 
 	// 数据库忽略字段
@@ -58,7 +58,7 @@ type UserOption struct {
 func (user *User) DeductionStorage(size uint64) bool {
 	if size <= user.Storage {
 		user.Storage -= size
-		DB.Save(user)
+		DB.Model(user).UpdateColumn("storage", gorm.Expr("storage - ?", size))
 		return true
 	}
 	return false
@@ -68,7 +68,7 @@ func (user *User) DeductionStorage(size uint64) bool {
 func (user *User) IncreaseStorage(size uint64) bool {
 	if size <= user.GetRemainingCapacity() {
 		user.Storage += size
-		DB.Save(user)
+		DB.Model(user).UpdateColumn("storage", gorm.Expr("storage + ?", size))
 		return true
 	}
 	return false
diff --git a/pkg/filesystem/context.go b/pkg/filesystem/context.go
new file mode 100644
index 0000000..59ddc2b
--- /dev/null
+++ b/pkg/filesystem/context.go
@@ -0,0 +1,12 @@
+package filesystem
+
+type key int
+
+const (
+	// GinCtx Gin的上下文
+	GinCtx key = iota
+	// SavePathCtx 文件物理路径
+	SavePathCtx
+	// FileCtx 上传的文件
+	FileCtx
+)
diff --git a/pkg/filesystem/filesystem.go b/pkg/filesystem/filesystem.go
index eef5e53..466826a 100644
--- a/pkg/filesystem/filesystem.go
+++ b/pkg/filesystem/filesystem.go
@@ -2,7 +2,6 @@ package filesystem
 
 import (
 	"context"
-	"fmt"
 	"github.com/HFO4/cloudreve/models"
 	"github.com/HFO4/cloudreve/pkg/filesystem/local"
 	"github.com/HFO4/cloudreve/pkg/util"
@@ -39,13 +38,13 @@ type FileSystem struct {
 	  钩子函数
 	*/
 	// 上传文件前
-	BeforeUpload func(ctx context.Context, fs *FileSystem, file FileData) error
+	BeforeUpload func(ctx context.Context, fs *FileSystem) error
 	// 上传文件后
 	AfterUpload func(ctx context.Context, fs *FileSystem) error
 	// 文件保存成功，插入数据库验证失败后
 	AfterValidateFailed func(ctx context.Context, fs *FileSystem) error
 	// 用户取消上传后
-	AfterUploadCanceled func(ctx context.Context, fs *FileSystem, file FileData) error
+	AfterUploadCanceled func(ctx context.Context, fs *FileSystem) error
 
 	/*
 	  文件系统处理适配器
@@ -72,15 +71,18 @@ func NewFileSystem(user *model.User) (*FileSystem, error) {
 	}, nil
 }
 
-/*
-	上传处理相关
+/* ================
+	 上传处理相关
+   ================
 */
 
 // Upload 上传文件
 func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) {
+	ctx = context.WithValue(ctx, FileCtx, file)
+
 	// 上传前的钩子
 	if fs.BeforeUpload != nil {
-		err = fs.BeforeUpload(ctx, fs, file)
+		err = fs.BeforeUpload(ctx, fs)
 		if err != nil {
 			return err
 		}
@@ -98,6 +100,15 @@ func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) {
 		return err
 	}
 
+	// 上传完成后的钩子
+	if fs.AfterUpload != nil {
+		ctx = context.WithValue(ctx, SavePathCtx, savePath)
+		err = fs.AfterUpload(ctx, fs)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -111,20 +122,30 @@ func (fs *FileSystem) GenerateSavePath(file FileData) string {
 
 // CancelUpload 监测客户端取消上传
 func (fs *FileSystem) CancelUpload(ctx context.Context, path string, file FileData) {
-	ginCtx := ctx.Value("ginCtx").(*gin.Context)
+	ginCtx := ctx.Value(GinCtx).(*gin.Context)
 	select {
 	case <-ctx.Done():
-		fmt.Println("正常关闭")
 		// 客户端正常关闭，不执行操作
 	case <-ginCtx.Request.Context().Done():
 		// 客户端取消了上传
 		if fs.AfterUploadCanceled == nil {
 			return
 		}
-		ctx = context.WithValue(ctx, "path", path)
-		err := fs.AfterUploadCanceled(ctx, fs, file)
+		ctx = context.WithValue(ctx, SavePathCtx, path)
+		err := fs.AfterUploadCanceled(ctx, fs)
 		if err != nil {
 			util.Log().Warning("执行 AfterUploadCanceled 钩子出错，%s", err)
 		}
 	}
 }
+
+/* =================
+	 路径/目录相关
+   =================
+*/
+
+// IsPathExist 返回给定目录是否存在
+func (fs *FileSystem) IsPathExist(path string) bool {
+	_, err := model.GetFolderByPath(path, fs.User.ID)
+	return err == nil
+}
diff --git a/pkg/filesystem/hook.go b/pkg/filesystem/hooks.go
similarity index 66%
rename from pkg/filesystem/hook.go
rename to pkg/filesystem/hooks.go
index 283bf08..0a42140 100644
--- a/pkg/filesystem/hook.go
+++ b/pkg/filesystem/hooks.go
@@ -7,7 +7,9 @@ import (
 )
 
 // GenericBeforeUpload 通用上传前处理钩子，包含数据库操作
-func GenericBeforeUpload(ctx context.Context, fs *FileSystem, file FileData) error {
+func GenericBeforeUpload(ctx context.Context, fs *FileSystem) error {
+	file := ctx.Value(FileCtx).(FileData)
+
 	// 验证单文件尺寸
 	if !fs.ValidateFileSize(ctx, file.GetSize()) {
 		return FileSizeTooBigError
@@ -31,7 +33,9 @@ func GenericBeforeUpload(ctx context.Context, fs *FileSystem, file FileData) err
 }
 
 // GenericAfterUploadCanceled 通用上传取消处理钩子，包含数据库操作
-func GenericAfterUploadCanceled(ctx context.Context, fs *FileSystem, file FileData) error {
+func GenericAfterUploadCanceled(ctx context.Context, fs *FileSystem) error {
+	file := ctx.Value(FileCtx).(FileData)
+
 	filePath := ctx.Value("path").(string)
 	// 删除临时文件
 	if util.Exists(filePath) {
@@ -47,3 +51,17 @@ func GenericAfterUploadCanceled(ctx context.Context, fs *FileSystem, file FileDa
 	}
 	return nil
 }
+
+// GenericAfterUpload 文件上传完成后，包含数据库操作
+func GenericAfterUpload(ctx context.Context, fs *FileSystem) error {
+	// 获取Gin的上下文
+	//ginCtx := ctx.Value(GinCtx).(*gin.Context)
+	// 文件存放的虚拟路径
+	//virtualPath := ginCtx.GetHeader("X-Path")
+
+	// 检查路径是否存在
+	if !fs.IsPathExist("/") {
+		return errors.New("sss")
+	}
+	return nil
+}
diff --git a/pkg/filesystem/hook_test.go b/pkg/filesystem/hooks_test.go
similarity index 100%
rename from pkg/filesystem/hook_test.go
rename to pkg/filesystem/hooks_test.go
diff --git a/pkg/filesystem/validator.go b/pkg/filesystem/validator.go
index 96fb525..f749f6c 100644
--- a/pkg/filesystem/validator.go
+++ b/pkg/filesystem/validator.go
@@ -12,11 +12,17 @@ var reservedCharacter = []string{"\\", "?", "*", "<", "\"", ":", ">", "/"}
 
 // ValidateLegalName 验证文件名/文件夹名是否合法
 func (fs *FileSystem) ValidateLegalName(ctx context.Context, name string) bool {
+	// 是否包含保留字符
 	for _, value := range reservedCharacter {
 		if strings.Contains(name, value) {
 			return false
 		}
 	}
+
+	// 是否超出长度限制
+	if len(name) >= 256 {
+		return false
+	}
 	return true
 }
 
diff --git a/routers/controllers/file.go b/routers/controllers/file.go
index aa19c6d..c1531a4 100644
--- a/routers/controllers/file.go
+++ b/routers/controllers/file.go
@@ -73,7 +73,7 @@ func FileUploadStream(c *gin.Context) {
 	fs.AfterUploadCanceled = filesystem.GenericAfterUploadCanceled
 
 	// 执行上传
-	uploadCtx := context.WithValue(ctx, "ginCtx", c)
+	uploadCtx := context.WithValue(ctx, filesystem.GinCtx, c)
 	err = fs.Upload(uploadCtx, fileData)
 	if err != nil {
 		c.JSON(200, serializer.Err(serializer.CodeUploadFailed, err.Error(), err))