github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/pkg/filesystem/hooks.go (about) 1 package filesystem 2 3 import ( 4 "context" 5 model "github.com/cloudreve/Cloudreve/v3/models" 6 "github.com/cloudreve/Cloudreve/v3/pkg/cache" 7 "github.com/cloudreve/Cloudreve/v3/pkg/cluster" 8 "github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/local" 9 "github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx" 10 "github.com/cloudreve/Cloudreve/v3/pkg/serializer" 11 "github.com/cloudreve/Cloudreve/v3/pkg/util" 12 "io/ioutil" 13 "net/http" 14 "strconv" 15 "strings" 16 "time" 17 ) 18 19 // Hook 钩子函数 20 type Hook func(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error 21 22 // Use 注入钩子 23 func (fs *FileSystem) Use(name string, hook Hook) { 24 if fs.Hooks == nil { 25 fs.Hooks = make(map[string][]Hook) 26 } 27 if _, ok := fs.Hooks[name]; ok { 28 fs.Hooks[name] = append(fs.Hooks[name], hook) 29 return 30 } 31 fs.Hooks[name] = []Hook{hook} 32 } 33 34 // CleanHooks 清空钩子,name为空表示全部清空 35 func (fs *FileSystem) CleanHooks(name string) { 36 if name == "" { 37 fs.Hooks = nil 38 } else { 39 delete(fs.Hooks, name) 40 } 41 } 42 43 // Trigger 触发钩子,遇到第一个错误时 44 // 返回错误,后续钩子不会继续执行 45 func (fs *FileSystem) Trigger(ctx context.Context, name string, file fsctx.FileHeader) error { 46 if hooks, ok := fs.Hooks[name]; ok { 47 for _, hook := range hooks { 48 err := hook(ctx, fs, file) 49 if err != nil { 50 util.Log().Warning("Failed to execute hook:%s", err) 51 return err 52 } 53 } 54 } 55 return nil 56 } 57 58 // HookValidateFile 一系列对文件检验的集合 59 func HookValidateFile(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 60 fileInfo := file.Info() 61 62 // 验证单文件尺寸 63 if !fs.ValidateFileSize(ctx, fileInfo.Size) { 64 return ErrFileSizeTooBig 65 } 66 67 // 验证文件名 68 if !fs.ValidateLegalName(ctx, fileInfo.FileName) { 69 return ErrIllegalObjectName 70 } 71 72 // 验证扩展名 73 if !fs.ValidateExtension(ctx, fileInfo.FileName) { 74 return ErrFileExtensionNotAllowed 75 } 76 77 return nil 78 79 } 80 81 // HookResetPolicy 重设存储策略为上下文已有文件 82 func HookResetPolicy(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 83 originFile, ok := ctx.Value(fsctx.FileModelCtx).(model.File) 84 if !ok { 85 return ErrObjectNotExist 86 } 87 88 fs.Policy = originFile.GetPolicy() 89 return fs.DispatchHandler() 90 } 91 92 // HookValidateCapacity 验证用户容量 93 func HookValidateCapacity(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 94 // 验证并扣除容量 95 if fs.User.GetRemainingCapacity() < file.Info().Size { 96 return ErrInsufficientCapacity 97 } 98 return nil 99 } 100 101 // HookValidateCapacityDiff 根据原有文件和新文件的大小验证用户容量 102 func HookValidateCapacityDiff(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error { 103 originFile := ctx.Value(fsctx.FileModelCtx).(model.File) 104 newFileSize := newFile.Info().Size 105 106 if newFileSize > originFile.Size { 107 return HookValidateCapacity(ctx, fs, newFile) 108 } 109 110 return nil 111 } 112 113 // HookDeleteTempFile 删除已保存的临时文件 114 func HookDeleteTempFile(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 115 // 删除临时文件 116 _, err := fs.Handler.Delete(ctx, []string{file.Info().SavePath}) 117 if err != nil { 118 util.Log().Warning("Failed to clean-up temp files: %s", err) 119 } 120 121 return nil 122 } 123 124 // HookCleanFileContent 清空文件内容 125 func HookCleanFileContent(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 126 // 清空内容 127 return fs.Handler.Put(ctx, &fsctx.FileStream{ 128 File: ioutil.NopCloser(strings.NewReader("")), 129 SavePath: file.Info().SavePath, 130 Size: 0, 131 Mode: fsctx.Overwrite, 132 }) 133 } 134 135 // HookClearFileSize 将原始文件的尺寸设为0 136 func HookClearFileSize(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 137 originFile, ok := ctx.Value(fsctx.FileModelCtx).(model.File) 138 if !ok { 139 return ErrObjectNotExist 140 } 141 return originFile.UpdateSize(0) 142 } 143 144 // HookCancelContext 取消上下文 145 func HookCancelContext(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 146 cancelFunc, ok := ctx.Value(fsctx.CancelFuncCtx).(context.CancelFunc) 147 if ok { 148 cancelFunc() 149 } 150 return nil 151 } 152 153 // HookUpdateSourceName 更新文件SourceName 154 func HookUpdateSourceName(ctx context.Context, fs *FileSystem, file fsctx.FileHeader) error { 155 originFile, ok := ctx.Value(fsctx.FileModelCtx).(model.File) 156 if !ok { 157 return ErrObjectNotExist 158 } 159 return originFile.UpdateSourceName(originFile.SourceName) 160 } 161 162 // GenericAfterUpdate 文件内容更新后 163 func GenericAfterUpdate(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error { 164 // 更新文件尺寸 165 originFile, ok := ctx.Value(fsctx.FileModelCtx).(model.File) 166 if !ok { 167 return ErrObjectNotExist 168 } 169 170 newFile.SetModel(&originFile) 171 172 err := originFile.UpdateSize(newFile.Info().Size) 173 if err != nil { 174 return err 175 } 176 177 return nil 178 } 179 180 // SlaveAfterUpload Slave模式下上传完成钩子 181 func SlaveAfterUpload(session *serializer.UploadSession) Hook { 182 return func(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 183 if session.Callback == "" { 184 return nil 185 } 186 187 // 发送回调请求 188 callbackBody := serializer.UploadCallback{} 189 return cluster.RemoteCallback(session.Callback, callbackBody) 190 } 191 } 192 193 // GenericAfterUpload 文件上传完成后,包含数据库操作 194 func GenericAfterUpload(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 195 fileInfo := fileHeader.Info() 196 197 // 创建或查找根目录 198 folder, err := fs.CreateDirectory(ctx, fileInfo.VirtualPath) 199 if err != nil { 200 return err 201 } 202 203 // 检查文件是否存在 204 if ok, file := fs.IsChildFileExist( 205 folder, 206 fileInfo.FileName, 207 ); ok { 208 if file.UploadSessionID != nil { 209 return ErrFileUploadSessionExisted 210 } 211 212 return ErrFileExisted 213 } 214 215 // 向数据库中插入记录 216 file, err := fs.AddFile(ctx, folder, fileHeader) 217 if err != nil { 218 return ErrInsertFileRecord 219 } 220 fileHeader.SetModel(file) 221 222 return nil 223 } 224 225 // HookClearFileHeaderSize 将FileHeader大小设定为0 226 func HookClearFileHeaderSize(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 227 fileHeader.SetSize(0) 228 return nil 229 } 230 231 // HookTruncateFileTo 将物理文件截断至 size 232 func HookTruncateFileTo(size uint64) Hook { 233 return func(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 234 if handler, ok := fs.Handler.(local.Driver); ok { 235 return handler.Truncate(ctx, fileHeader.Info().SavePath, size) 236 } 237 238 return nil 239 } 240 } 241 242 // HookChunkUploadFinished 单个分片上传结束后 243 func HookChunkUploaded(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 244 fileInfo := fileHeader.Info() 245 246 // 更新文件大小 247 return fileInfo.Model.(*model.File).UpdateSize(fileInfo.AppendStart + fileInfo.Size) 248 } 249 250 // HookChunkUploadFailed 单个分片上传失败后 251 func HookChunkUploadFailed(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 252 fileInfo := fileHeader.Info() 253 254 // 更新文件大小 255 return fileInfo.Model.(*model.File).UpdateSize(fileInfo.AppendStart) 256 } 257 258 // HookPopPlaceholderToFile 将占位文件提升为正式文件 259 func HookPopPlaceholderToFile(picInfo string) Hook { 260 return func(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 261 fileInfo := fileHeader.Info() 262 fileModel := fileInfo.Model.(*model.File) 263 return fileModel.PopChunkToFile(fileInfo.LastModified, picInfo) 264 } 265 } 266 267 // HookChunkUploadFinished 分片上传结束后处理文件 268 func HookDeleteUploadSession(id string) Hook { 269 return func(ctx context.Context, fs *FileSystem, fileHeader fsctx.FileHeader) error { 270 cache.Deletes([]string{id}, UploadSessionCachePrefix) 271 return nil 272 } 273 } 274 275 // NewWebdavAfterUploadHook 每次创建一个新的钩子函数 rclone 在 PUT 请求里有 OC-Checksum 字符串 276 // 和 X-OC-Mtime 277 func NewWebdavAfterUploadHook(request *http.Request) func(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error { 278 var modtime time.Time 279 if timeVal := request.Header.Get("X-OC-Mtime"); timeVal != "" { 280 timeUnix, err := strconv.ParseInt(timeVal, 10, 64) 281 if err == nil { 282 modtime = time.Unix(timeUnix, 0) 283 } 284 } 285 checksum := request.Header.Get("OC-Checksum") 286 287 return func(ctx context.Context, fs *FileSystem, newFile fsctx.FileHeader) error { 288 file := newFile.Info().Model.(*model.File) 289 if !modtime.IsZero() { 290 err := model.DB.Model(file).UpdateColumn("updated_at", modtime).Error 291 if err != nil { 292 return err 293 } 294 } 295 296 if checksum != "" { 297 return file.UpdateMetadata(map[string]string{ 298 model.ChecksumMetadataKey: checksum, 299 }) 300 } 301 302 return nil 303 } 304 }