From 4a5782b4e5e59d4e575835906dc3eaae4f1ce3a5 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Mon, 30 Dec 2019 19:37:21 +0800
Subject: [PATCH] Feat: remote upload callback / validation

---
 models/init.go              |  2 +-
 pkg/filesystem/file.go      | 13 +++++++++--
 pkg/filesystem/hooks.go     |  2 +-
 pkg/filesystem/validator.go |  5 +---
 pkg/serializer/upload.go    |  1 +
 service/callback/upload.go  | 46 +++++++++++++++++++++++++++++++++++++
 6 files changed, 61 insertions(+), 8 deletions(-)

diff --git a/models/init.go b/models/init.go
index bafc977..c782745 100644
--- a/models/init.go
+++ b/models/init.go
@@ -47,7 +47,7 @@ func Init() {
 
 	// Debug模式下，输出所有 SQL 日志
 	if conf.SystemConfig.Debug {
-		db.LogMode(false)
+		db.LogMode(true)
 	}
 
 	//db.SetLogger(util.Log())
diff --git a/pkg/filesystem/file.go b/pkg/filesystem/file.go
index 06d6b39..81967ad 100644
--- a/pkg/filesystem/file.go
+++ b/pkg/filesystem/file.go
@@ -43,6 +43,12 @@ func (fs *FileSystem) withSpeedLimit(rs response.RSCloser) response.RSCloser {
 
 // AddFile 新增文件记录
 func (fs *FileSystem) AddFile(ctx context.Context, parent *model.Folder) (*model.File, error) {
+	// 添加文件记录前的钩子
+	err := fs.Trigger(ctx, "BeforeAddFile")
+	if err != nil {
+		return nil, err
+	}
+
 	file := ctx.Value(fsctx.FileHeaderCtx).(FileHeader)
 	filePath := ctx.Value(fsctx.SavePathCtx).(string)
 
@@ -54,10 +60,13 @@ func (fs *FileSystem) AddFile(ctx context.Context, parent *model.Folder) (*model
 		FolderID:   parent.ID,
 		PolicyID:   fs.User.Policy.ID,
 	}
-	_, err := newFile.Create()
+	_, err = newFile.Create()
 
 	if err != nil {
-		return nil, err
+		if err := fs.Trigger(ctx, "AfterValidateFailed"); err != nil {
+			util.Log().Debug("AfterValidateFailed 钩子执行失败，%s", err)
+		}
+		return nil, ErrFileExisted.WithError(err)
 	}
 
 	return &newFile, nil
diff --git a/pkg/filesystem/hooks.go b/pkg/filesystem/hooks.go
index 2acfbc4..fdde788 100644
--- a/pkg/filesystem/hooks.go
+++ b/pkg/filesystem/hooks.go
@@ -208,7 +208,6 @@ func GenericAfterUpdate(ctx context.Context, fs *FileSystem) error {
 }
 
 // SlaveAfterUpload Slave模式下上传完成钩子
-// TODO 测试
 func SlaveAfterUpload(ctx context.Context, fs *FileSystem) error {
 	fileHeader := ctx.Value(fsctx.FileHeaderCtx).(FileHeader)
 	policy := ctx.Value(fsctx.UploadPolicyCtx).(serializer.UploadPolicy)
@@ -225,6 +224,7 @@ func SlaveAfterUpload(ctx context.Context, fs *FileSystem) error {
 		Name:       file.Name,
 		SourceName: file.SourceName,
 		PicInfo:    file.PicInfo,
+		Size:       fileHeader.GetSize(),
 	}
 	return request.RemoteCallback(policy.CallbackURL, callbackBody)
 }
diff --git a/pkg/filesystem/validator.go b/pkg/filesystem/validator.go
index c87492e..30e7453 100644
--- a/pkg/filesystem/validator.go
+++ b/pkg/filesystem/validator.go
@@ -44,10 +44,7 @@ func (fs *FileSystem) ValidateFileSize(ctx context.Context, size uint64) bool {
 
 // ValidateCapacity 验证并扣除用户容量
 func (fs *FileSystem) ValidateCapacity(ctx context.Context, size uint64) bool {
-	if fs.User.IncreaseStorage(size) {
-		return true
-	}
-	return false
+	return fs.User.IncreaseStorage(size)
 }
 
 // ValidateExtension 验证文件扩展名
diff --git a/pkg/serializer/upload.go b/pkg/serializer/upload.go
index 0d430b1..aeb1ddc 100644
--- a/pkg/serializer/upload.go
+++ b/pkg/serializer/upload.go
@@ -34,6 +34,7 @@ type RemoteUploadCallback struct {
 	Name       string `json:"name"`
 	SourceName string `json:"source_name"`
 	PicInfo    string `json:"pic_info"`
+	Size       uint64 `json:"size"`
 }
 
 func init() {
diff --git a/service/callback/upload.go b/service/callback/upload.go
index afbbe6c..d37f72e 100644
--- a/service/callback/upload.go
+++ b/service/callback/upload.go
@@ -1,8 +1,12 @@
 package callback
 
 import (
+	"context"
 	"github.com/HFO4/cloudreve/pkg/filesystem"
+	"github.com/HFO4/cloudreve/pkg/filesystem/fsctx"
+	"github.com/HFO4/cloudreve/pkg/filesystem/local"
 	"github.com/HFO4/cloudreve/pkg/serializer"
+	"github.com/HFO4/cloudreve/pkg/util"
 	"github.com/gin-gonic/gin"
 )
 
@@ -20,6 +24,48 @@ func (service *RemoteUploadCallbackService) Process(c *gin.Context) serializer.R
 	}
 	defer fs.Recycle()
 
+	// 获取回调会话
+	callbackSessionRaw, ok := c.Get("callbackSession")
+	if !ok {
+		return serializer.Err(serializer.CodeInternalSetting, "找不到回调会话", nil)
+	}
+	callbackSession := callbackSessionRaw.(*serializer.UploadSession)
+
+	// 获取父目录
+	exist, parentFolder := fs.IsPathExist(callbackSession.VirtualPath)
+	if !exist {
+		return serializer.Err(serializer.CodeParamErr, "指定目录不存在", nil)
+	}
+
+	// 创建文件头
+	fileHeader := local.FileStream{
+		Size:        service.Data.Size,
+		VirtualPath: callbackSession.VirtualPath,
+		Name:        service.Data.Name,
+	}
+
+	// 生成上下文
+	ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, fileHeader)
+	ctx = context.WithValue(ctx, fsctx.SavePathCtx, service.Data.SourceName)
+
+	// 添加钩子
+	fs.Use("BeforeAddFile", filesystem.HookValidateFile)
+	fs.Use("BeforeAddFile", filesystem.HookValidateCapacity)
+	fs.Use("AfterValidateFailed", filesystem.HookGiveBackCapacity)
+
+	// 向数据库中添加文件
+	file, err := fs.AddFile(ctx, parentFolder)
+	if err != nil {
+		return serializer.Err(serializer.CodeUploadFailed, err.Error(), err)
+	}
+
+	// 如果是图片，则更新图片信息
+	if service.Data.PicInfo != "" {
+		if err := file.UpdatePicInfo(service.Data.PicInfo); err != nil {
+			util.Log().Debug("无法更新回调文件的图片信息：%s", err)
+		}
+	}
+
 	return serializer.Response{
 		Code: 0,
 	}