mirror of
				https://github.com/cloudreve/cloudreve.git
				synced 2025-10-31 08:39:10 +08:00 
			
		
		
		
	Feat: after uploading hooks and checks
This commit is contained in:
		| @ -15,3 +15,10 @@ type Folder struct { | |||||||
| 	// 关联模型 | 	// 关联模型 | ||||||
| 	OptionsSerialized PolicyOption `gorm:"-"` | 	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 | ||||||
|  | } | ||||||
|  | |||||||
							
								
								
									
										20
									
								
								models/folder_test.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										20
									
								
								models/folder_test.go
									
									
									
									
									
										Normal file
									
								
							| @ -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.) | ||||||
|  | } | ||||||
| @ -85,6 +85,11 @@ func (policy *Policy) GeneratePath(uid uint) string { | |||||||
|  |  | ||||||
| // GenerateFileName 生成存储文件名 | // GenerateFileName 生成存储文件名 | ||||||
| func (policy *Policy) GenerateFileName(uid uint, origin string) string { | func (policy *Policy) GenerateFileName(uid uint, origin string) string { | ||||||
|  | 	// 未开启自动重命名时,直接返回原始文件名 | ||||||
|  | 	if !policy.AutoRename { | ||||||
|  | 		return origin | ||||||
|  | 	} | ||||||
|  |  | ||||||
| 	fileRule := policy.FileNameRule | 	fileRule := policy.FileNameRule | ||||||
|  |  | ||||||
| 	replaceTable := map[string]string{ | 	replaceTable := map[string]string{ | ||||||
|  | |||||||
| @ -40,7 +40,7 @@ type User struct { | |||||||
| 	Options       string `json:"-",gorm:"size:4096"` | 	Options       string `json:"-",gorm:"size:4096"` | ||||||
|  |  | ||||||
| 	// 关联模型 | 	// 关联模型 | ||||||
| 	Group  Group | 	Group  Group  `gorm:"association_autoupdate:false"` | ||||||
| 	Policy Policy `gorm:"PRELOAD:false,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 { | func (user *User) DeductionStorage(size uint64) bool { | ||||||
| 	if size <= user.Storage { | 	if size <= user.Storage { | ||||||
| 		user.Storage -= size | 		user.Storage -= size | ||||||
| 		DB.Save(user) | 		DB.Model(user).UpdateColumn("storage", gorm.Expr("storage - ?", size)) | ||||||
| 		return true | 		return true | ||||||
| 	} | 	} | ||||||
| 	return false | 	return false | ||||||
| @ -68,7 +68,7 @@ func (user *User) DeductionStorage(size uint64) bool { | |||||||
| func (user *User) IncreaseStorage(size uint64) bool { | func (user *User) IncreaseStorage(size uint64) bool { | ||||||
| 	if size <= user.GetRemainingCapacity() { | 	if size <= user.GetRemainingCapacity() { | ||||||
| 		user.Storage += size | 		user.Storage += size | ||||||
| 		DB.Save(user) | 		DB.Model(user).UpdateColumn("storage", gorm.Expr("storage + ?", size)) | ||||||
| 		return true | 		return true | ||||||
| 	} | 	} | ||||||
| 	return false | 	return false | ||||||
|  | |||||||
							
								
								
									
										12
									
								
								pkg/filesystem/context.go
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										12
									
								
								pkg/filesystem/context.go
									
									
									
									
									
										Normal file
									
								
							| @ -0,0 +1,12 @@ | |||||||
|  | package filesystem | ||||||
|  |  | ||||||
|  | type key int | ||||||
|  |  | ||||||
|  | const ( | ||||||
|  | 	// GinCtx Gin的上下文 | ||||||
|  | 	GinCtx key = iota | ||||||
|  | 	// SavePathCtx 文件物理路径 | ||||||
|  | 	SavePathCtx | ||||||
|  | 	// FileCtx 上传的文件 | ||||||
|  | 	FileCtx | ||||||
|  | ) | ||||||
| @ -2,7 +2,6 @@ package filesystem | |||||||
|  |  | ||||||
| import ( | import ( | ||||||
| 	"context" | 	"context" | ||||||
| 	"fmt" |  | ||||||
| 	"github.com/HFO4/cloudreve/models" | 	"github.com/HFO4/cloudreve/models" | ||||||
| 	"github.com/HFO4/cloudreve/pkg/filesystem/local" | 	"github.com/HFO4/cloudreve/pkg/filesystem/local" | ||||||
| 	"github.com/HFO4/cloudreve/pkg/util" | 	"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 | 	AfterUpload func(ctx context.Context, fs *FileSystem) error | ||||||
| 	// 文件保存成功,插入数据库验证失败后 | 	// 文件保存成功,插入数据库验证失败后 | ||||||
| 	AfterValidateFailed 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 | 	}, nil | ||||||
| } | } | ||||||
|  |  | ||||||
| /* | /* ================ | ||||||
| 	 上传处理相关 | 	 上传处理相关 | ||||||
|  |    ================ | ||||||
| */ | */ | ||||||
|  |  | ||||||
| // Upload 上传文件 | // Upload 上传文件 | ||||||
| func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) { | func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) { | ||||||
|  | 	ctx = context.WithValue(ctx, FileCtx, file) | ||||||
|  |  | ||||||
| 	// 上传前的钩子 | 	// 上传前的钩子 | ||||||
| 	if fs.BeforeUpload != nil { | 	if fs.BeforeUpload != nil { | ||||||
| 		err = fs.BeforeUpload(ctx, fs, file) | 		err = fs.BeforeUpload(ctx, fs) | ||||||
| 		if err != nil { | 		if err != nil { | ||||||
| 			return err | 			return err | ||||||
| 		} | 		} | ||||||
| @ -98,6 +100,15 @@ func (fs *FileSystem) Upload(ctx context.Context, file FileData) (err error) { | |||||||
| 		return err | 		return err | ||||||
| 	} | 	} | ||||||
|  |  | ||||||
|  | 	// 上传完成后的钩子 | ||||||
|  | 	if fs.AfterUpload != nil { | ||||||
|  | 		ctx = context.WithValue(ctx, SavePathCtx, savePath) | ||||||
|  | 		err = fs.AfterUpload(ctx, fs) | ||||||
|  | 		if err != nil { | ||||||
|  | 			return err | ||||||
|  | 		} | ||||||
|  | 	} | ||||||
|  |  | ||||||
| 	return nil | 	return nil | ||||||
| } | } | ||||||
|  |  | ||||||
| @ -111,20 +122,30 @@ func (fs *FileSystem) GenerateSavePath(file FileData) string { | |||||||
|  |  | ||||||
| // CancelUpload 监测客户端取消上传 | // CancelUpload 监测客户端取消上传 | ||||||
| func (fs *FileSystem) CancelUpload(ctx context.Context, path string, file FileData) { | func (fs *FileSystem) CancelUpload(ctx context.Context, path string, file FileData) { | ||||||
| 	ginCtx := ctx.Value("ginCtx").(*gin.Context) | 	ginCtx := ctx.Value(GinCtx).(*gin.Context) | ||||||
| 	select { | 	select { | ||||||
| 	case <-ctx.Done(): | 	case <-ctx.Done(): | ||||||
| 		fmt.Println("正常关闭") |  | ||||||
| 		// 客户端正常关闭,不执行操作 | 		// 客户端正常关闭,不执行操作 | ||||||
| 	case <-ginCtx.Request.Context().Done(): | 	case <-ginCtx.Request.Context().Done(): | ||||||
| 		// 客户端取消了上传 | 		// 客户端取消了上传 | ||||||
| 		if fs.AfterUploadCanceled == nil { | 		if fs.AfterUploadCanceled == nil { | ||||||
| 			return | 			return | ||||||
| 		} | 		} | ||||||
| 		ctx = context.WithValue(ctx, "path", path) | 		ctx = context.WithValue(ctx, SavePathCtx, path) | ||||||
| 		err := fs.AfterUploadCanceled(ctx, fs, file) | 		err := fs.AfterUploadCanceled(ctx, fs) | ||||||
| 		if err != nil { | 		if err != nil { | ||||||
| 			util.Log().Warning("执行 AfterUploadCanceled 钩子出错,%s", err) | 			util.Log().Warning("执行 AfterUploadCanceled 钩子出错,%s", err) | ||||||
| 		} | 		} | ||||||
| 	} | 	} | ||||||
| } | } | ||||||
|  |  | ||||||
|  | /* ================= | ||||||
|  | 	 路径/目录相关 | ||||||
|  |    ================= | ||||||
|  | */ | ||||||
|  |  | ||||||
|  | // IsPathExist 返回给定目录是否存在 | ||||||
|  | func (fs *FileSystem) IsPathExist(path string) bool { | ||||||
|  | 	_, err := model.GetFolderByPath(path, fs.User.ID) | ||||||
|  | 	return err == nil | ||||||
|  | } | ||||||
|  | |||||||
| @ -7,7 +7,9 @@ import ( | |||||||
| ) | ) | ||||||
| 
 | 
 | ||||||
| // GenericBeforeUpload 通用上传前处理钩子,包含数据库操作 | // 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()) { | 	if !fs.ValidateFileSize(ctx, file.GetSize()) { | ||||||
| 		return FileSizeTooBigError | 		return FileSizeTooBigError | ||||||
| @ -31,7 +33,9 @@ func GenericBeforeUpload(ctx context.Context, fs *FileSystem, file FileData) err | |||||||
| } | } | ||||||
| 
 | 
 | ||||||
| // GenericAfterUploadCanceled 通用上传取消处理钩子,包含数据库操作 | // 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) | 	filePath := ctx.Value("path").(string) | ||||||
| 	// 删除临时文件 | 	// 删除临时文件 | ||||||
| 	if util.Exists(filePath) { | 	if util.Exists(filePath) { | ||||||
| @ -47,3 +51,17 @@ func GenericAfterUploadCanceled(ctx context.Context, fs *FileSystem, file FileDa | |||||||
| 	} | 	} | ||||||
| 	return nil | 	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 | ||||||
|  | } | ||||||
| @ -12,11 +12,17 @@ var reservedCharacter = []string{"\\", "?", "*", "<", "\"", ":", ">", "/"} | |||||||
|  |  | ||||||
| // ValidateLegalName 验证文件名/文件夹名是否合法 | // ValidateLegalName 验证文件名/文件夹名是否合法 | ||||||
| func (fs *FileSystem) ValidateLegalName(ctx context.Context, name string) bool { | func (fs *FileSystem) ValidateLegalName(ctx context.Context, name string) bool { | ||||||
|  | 	// 是否包含保留字符 | ||||||
| 	for _, value := range reservedCharacter { | 	for _, value := range reservedCharacter { | ||||||
| 		if strings.Contains(name, value) { | 		if strings.Contains(name, value) { | ||||||
| 			return false | 			return false | ||||||
| 		} | 		} | ||||||
| 	} | 	} | ||||||
|  |  | ||||||
|  | 	// 是否超出长度限制 | ||||||
|  | 	if len(name) >= 256 { | ||||||
|  | 		return false | ||||||
|  | 	} | ||||||
| 	return true | 	return true | ||||||
| } | } | ||||||
|  |  | ||||||
|  | |||||||
| @ -73,7 +73,7 @@ func FileUploadStream(c *gin.Context) { | |||||||
| 	fs.AfterUploadCanceled = filesystem.GenericAfterUploadCanceled | 	fs.AfterUploadCanceled = filesystem.GenericAfterUploadCanceled | ||||||
|  |  | ||||||
| 	// 执行上传 | 	// 执行上传 | ||||||
| 	uploadCtx := context.WithValue(ctx, "ginCtx", c) | 	uploadCtx := context.WithValue(ctx, filesystem.GinCtx, c) | ||||||
| 	err = fs.Upload(uploadCtx, fileData) | 	err = fs.Upload(uploadCtx, fileData) | ||||||
| 	if err != nil { | 	if err != nil { | ||||||
| 		c.JSON(200, serializer.Err(serializer.CodeUploadFailed, err.Error(), err)) | 		c.JSON(200, serializer.Err(serializer.CodeUploadFailed, err.Error(), err)) | ||||||
|  | |||||||
		Reference in New Issue
	
	Block a user
	 HFO4
					HFO4