From 160f9645646350db07ee72ae2cff6eddcea4400c Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Mon, 18 Nov 2019 13:26:32 +0800
Subject: [PATCH] Feat: use goroutine to detect upload-canceling action /
 Object name validate

---
 pkg/filesystem/errors.go        |  1 +
 pkg/filesystem/filesystem.go    | 22 ++++++++++++++++++++++
 pkg/filesystem/hook.go          |  5 +++++
 pkg/filesystem/local/handler.go |  2 ++
 pkg/filesystem/validator.go     | 13 +++++++++++++
 routers/controllers/file.go     | 12 +++++++++++-
 6 files changed, 54 insertions(+), 1 deletion(-)

diff --git a/pkg/filesystem/errors.go b/pkg/filesystem/errors.go
index 211ff72..023def3 100644
--- a/pkg/filesystem/errors.go
+++ b/pkg/filesystem/errors.go
@@ -7,4 +7,5 @@ var (
 	FileSizeTooBigError          = errors.New("单个文件尺寸太大")
 	FileExtensionNotAllowedError = errors.New("不允许上传此类型的文件")
 	InsufficientCapacityError    = errors.New("容量空间不足")
+	IlegalObjectNameError        = errors.New("目标名称非法")
 )
diff --git a/pkg/filesystem/filesystem.go b/pkg/filesystem/filesystem.go
index 91fcfc3..a6e3f7e 100644
--- a/pkg/filesystem/filesystem.go
+++ b/pkg/filesystem/filesystem.go
@@ -2,8 +2,10 @@ package filesystem
 
 import (
 	"context"
+	"fmt"
 	"github.com/HFO4/cloudreve/models"
 	"github.com/HFO4/cloudreve/pkg/filesystem/local"
+	"github.com/gin-gonic/gin"
 	"io"
 	"path/filepath"
 )
@@ -64,6 +66,10 @@ func NewFileSystem(user *model.User) (*FileSystem, error) {
 	}, nil
 }
 
+/*
+	上传处理相关
+*/
+
 // Upload 上传文件
 func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) {
 	// 上传前的钩子
@@ -75,6 +81,9 @@ func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) {
 	// 生成文件名和路径
 	savePath := fs.GenerateSavePath(file)
 
+	// 处理客户端未完成上传时，关闭连接
+	go fs.CancelUpload(ctx, savePath, file)
+
 	// 保存文件
 	err = fs.Handler.Put(ctx, file, savePath)
 	if err != nil {
@@ -91,3 +100,16 @@ func (fs *FileSystem) GenerateSavePath(file FileData) string {
 		fs.User.Policy.GenerateFileName(fs.User.Model.ID, file.GetFileName()),
 	)
 }
+
+// CancelUpload 监测客户端取消上传
+func (fs *FileSystem) CancelUpload(ctx context.Context, path string, file FileData) {
+	ginCtx := ctx.Value("ginCtx").(*gin.Context)
+	select {
+	case <-ctx.Done():
+		// 客户端正常关闭，不执行操作
+	case <-ginCtx.Request.Context().Done():
+		// 客户端取消了上传,删除保存的文件
+		fmt.Println("取消上传")
+		// 归还空间
+	}
+}
diff --git a/pkg/filesystem/hook.go b/pkg/filesystem/hook.go
index a0d606e..854bc48 100644
--- a/pkg/filesystem/hook.go
+++ b/pkg/filesystem/hook.go
@@ -11,6 +11,11 @@ func GenericBeforeUpload(ctx context.Context, fs *FileSystem, file FileData) err
 		return FileSizeTooBigError
 	}
 
+	// 验证文件名
+	if !fs.ValidateLegalName(ctx, file.GetFileName()) {
+		return IlegalObjectNameError
+	}
+
 	// 验证扩展名
 	if !fs.ValidateExtension(ctx, file.GetFileName()) {
 		return FileExtensionNotAllowedError
diff --git a/pkg/filesystem/local/handler.go b/pkg/filesystem/local/handler.go
index 148a16c..d2ce283 100644
--- a/pkg/filesystem/local/handler.go
+++ b/pkg/filesystem/local/handler.go
@@ -24,12 +24,14 @@ func (handler Handler) Put(ctx context.Context, file io.ReadCloser, dst string)
 		}
 	}
 
+	// 创建目标文件
 	out, err := os.Create(dst)
 	if err != nil {
 		return err
 	}
 	defer out.Close()
 
+	// 写入文件内容
 	_, err = io.Copy(out, file)
 	return err
 }
diff --git a/pkg/filesystem/validator.go b/pkg/filesystem/validator.go
index 94a735b..2b0d129 100644
--- a/pkg/filesystem/validator.go
+++ b/pkg/filesystem/validator.go
@@ -7,6 +7,19 @@ import (
 	"strings"
 )
 
+// 文件/路径名保留字符
+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
+		}
+	}
+	return true
+}
+
 // ValidateFileSize 验证上传的文件大小是否超出限制
 func (fs *FileSystem) ValidateFileSize(ctx context.Context, size uint64) bool {
 	return size <= fs.User.Policy.MaxSize
diff --git a/routers/controllers/file.go b/routers/controllers/file.go
index 82a7549..d28ede0 100644
--- a/routers/controllers/file.go
+++ b/routers/controllers/file.go
@@ -34,10 +34,19 @@ func FileUpload(c *gin.Context) {
 	}
 }
 
+// FileUploadStream 本地策略流式上传
 func FileUploadStream(c *gin.Context) {
+	// 创建上下文
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
+	// 非本地策略时拒绝上传
+	if user, ok := c.Get("user"); ok && user.(*model.User).Policy.Type != "local" {
+		c.JSON(200, serializer.Err(serializer.CodePolicyNotAllowed, "当前上传策略无法使用", nil))
+		return
+	}
+
+	// 取得文件大小
 	fileSize, err := strconv.ParseUint(c.Request.Header.Get("Content-Length"), 10, 64)
 	if err != nil {
 		c.JSON(200, ErrorResponse(err))
@@ -63,7 +72,8 @@ func FileUploadStream(c *gin.Context) {
 	fs.BeforeUpload = filesystem.GenericBeforeUpload
 
 	// 执行上传
-	err = fs.Upload(ctx, fileData)
+	uploadCtx := context.WithValue(ctx, "ginCtx", c)
+	err = fs.Upload(uploadCtx, fileData)
 	if err != nil {
 		c.JSON(200, serializer.Err(serializer.CodeUploadFailed, err.Error(), err))
 		return