Feat: validate / cancel task while downloading file in aria2

models/download.go

@@ -22,6 +22,9 @@ type Download struct {
 	Dst            string `gorm:"type:text"` // 用户文件系统存储父目录路径
 	UserID         uint   // 发起者UID
 	TaskID         uint   // 对应的转存任务ID
+
+	// 关联模型
+	User *User `gorm:"PRELOAD:false,association_autoupdate:false"`
 }
 
 // Create 创建离线下载记录
@@ -48,3 +51,13 @@ func GetDownloadsByStatus(status ...int) []Download {
 	DB.Where("status in (?)", status).Find(&tasks)
 	return tasks
 }
+
+// GetOwner 获取下载任务所属用户
+func (task *Download) GetOwner() *User {
+	if task.User == nil {
+		if user, err := GetUserByID(task.UserID); err == nil {
+			return &user
+		}
+	}
+	return task.User
+}

models/task.go

@@ -12,7 +12,7 @@ type Task struct {
 	Type     int    // 任务类型
 	UserID   uint   // 发起者UID，0表示为系统发起
 	Progress int    // 进度
-	Error    string // 错误信息
+	Error    string `gorm:"type:text"` // 错误信息
 	Props    string `gorm:"type:text"` // 任务属性
 }
 

pkg/aria2/Monitor.go

@@ -1,9 +1,13 @@
 package aria2
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	model "github.com/HFO4/cloudreve/models"
+	"github.com/HFO4/cloudreve/pkg/filesystem"
+	"github.com/HFO4/cloudreve/pkg/filesystem/driver/local"
+	"github.com/HFO4/cloudreve/pkg/filesystem/fsctx"
 	"github.com/HFO4/cloudreve/pkg/task"
 	"github.com/HFO4/cloudreve/pkg/util"
 	"github.com/zyxar/argo/rpc"
@@ -71,9 +75,18 @@ func (monitor *Monitor) Update() bool {
 		return true
 	}
 
+	// 磁力链下载需要跟随
+	if len(status.FollowedBy) > 0 {
+		util.Log().Debug("离线下载[%s]重定向至[%s]", monitor.Task.GID, status.FollowedBy[0])
+		monitor.Task.GID = status.FollowedBy[0]
+		monitor.Task.Save()
+		return false
+	}
+
 	// 更新任务信息
 	if err := monitor.UpdateTaskInfo(status); err != nil {
 		util.Log().Warning("无法更新下载任务[%s]的任务信息[%s]，", monitor.Task.GID, err)
+		monitor.setErrorStatus(err)
 		return true
 	}
 
@@ -96,6 +109,9 @@ func (monitor *Monitor) Update() bool {
 
 // UpdateTaskInfo 更新数据库中的任务信息
 func (monitor *Monitor) UpdateTaskInfo(status rpc.StatusInfo) error {
+	originSize := monitor.Task.TotalSize
+	originPath := monitor.Task.Path
+
 	monitor.Task.GID = status.Gid
 	monitor.Task.Status = getStatus(status.Status)
 
@@ -126,7 +142,68 @@ func (monitor *Monitor) UpdateTaskInfo(status rpc.StatusInfo) error {
 	attrs, _ := json.Marshal(status)
 	monitor.Task.Attrs = string(attrs)
 
-	return monitor.Task.Save()
+	if err := monitor.Task.Save(); err != nil {
+		return nil
+	}
+
+	if originSize != monitor.Task.TotalSize || originPath != monitor.Task.Path {
+		// 大小、文件名更新后，对文件限制等进行校验
+		if err := monitor.ValidateFile(); err != nil {
+			// 验证失败时取消任务
+			monitor.Cancel()
+			return err
+		}
+	}
+
+	return nil
+}
+
+// Cancel 取消上传并尝试删除临时文件
+func (monitor *Monitor) Cancel() {
+	if err := Instance.Cancel(monitor.Task); err != nil {
+		util.Log().Warning("无法取消离线下载任务[%s], %s", monitor.Task.GID, err)
+	}
+	util.Log().Debug("离线下载任务[%s]已取消，1 分钟后删除临时文件", monitor.Task.GID)
+	go func(monitor *Monitor) {
+		select {
+		case <-time.After(time.Duration(60) * time.Second):
+			monitor.RemoveTempFolder()
+		}
+	}(monitor)
+}
+
+// ValidateFile 上传过程中校验文件大小、文件名
+func (monitor *Monitor) ValidateFile() error {
+	// 找到任务创建者
+	user := monitor.Task.GetOwner()
+	if user == nil {
+		return ErrUserNotFound
+	}
+
+	// 创建文件系统
+	fs, err := filesystem.NewFileSystem(user)
+	if err != nil {
+		return err
+	}
+	defer fs.Recycle()
+
+	// 创建上下文环境
+	ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, local.FileStream{
+		Size: monitor.Task.TotalSize,
+		Name: filepath.Base(monitor.Task.Path),
+	})
+
+	// 验证文件
+	if err := filesystem.HookValidateFile(ctx, fs); err != nil {
+		return err
+	}
+
+	// 验证用户容量
+	if err := filesystem.HookValidateCapacityWithoutIncrease(ctx, fs); err != nil {
+		return err
+	}
+
+	return nil
 }
 
 // Error 任务下载出错处理，返回是否中断监控

pkg/aria2/aria2.go

@@ -20,6 +20,8 @@ type Aria2 interface {
 	CreateTask(task *model.Download) error
 	// 返回状态信息
 	Status(task *model.Download) (rpc.StatusInfo, error)
+	// 取消任务
+	Cancel(task *model.Download) error
 }
 
 const (
@@ -49,6 +51,7 @@ const (
 var (
 	// ErrNotEnabled 功能未开启错误
 	ErrNotEnabled   = serializer.NewError(serializer.CodeNoPermissionErr, "离线下载功能未开启", nil)
+	ErrUserNotFound = serializer.NewError(serializer.CodeNotFound, "无法找到任务创建者", nil)
 )
 
 // DummyAria2 未开启Aria2功能时使用的默认处理器
@@ -65,6 +68,11 @@ func (instance *DummyAria2) Status(task *model.Download) (rpc.StatusInfo, error)
 	return rpc.StatusInfo{}, ErrNotEnabled
 }
 
+// Cancel 返回未开启错误
+func (instance *DummyAria2) Cancel(task *model.Download) error {
+	return ErrNotEnabled
+}
+
 // Init 初始化
 func Init() {
 	options := model.GetSettingByNames("aria2_rpcurl", "aria2_token", "aria2_options")

pkg/aria2/caller.go

@@ -40,6 +40,12 @@ func (client *RPCService) Status(task *model.Download) (rpc.StatusInfo, error) {
 	return client.caller.TellStatus(task.GID)
 }
 
+// Cancel 取消下载
+func (client *RPCService) Cancel(task *model.Download) error {
+	_, err := client.caller.Remove(task.GID)
+	return err
+}
+
 // CreateTask 创建新任务
 func (client *RPCService) CreateTask(task *model.Download) error {
 	// 生成存储路径

pkg/filesystem/driver/onedrive/api.go

@@ -153,7 +153,7 @@ func (client *Client) UploadChunk(ctx context.Context, uploadURL string, chunk *
 		// 如果重试次数小于限制，5秒后重试
 		if chunk.Retried < model.GetIntSetting("onedrive_chunk_retries", 1) {
 			chunk.Retried++
-			util.Log().Debug("分片偏移%d上传失败，5秒钟后重试", chunk.Offset)
+			util.Log().Debug("分片偏移%d上传失败[%s]，5秒钟后重试", chunk.Offset, err)
 			time.Sleep(time.Duration(5) * time.Second)
 			return client.UploadChunk(ctx, uploadURL, chunk)
 		}
@@ -518,6 +518,7 @@ func (client *Client) request(ctx context.Context, method string, url string, bo
 	if res.Response.StatusCode < 200 || res.Response.StatusCode >= 300 {
 		decodeErr = json.Unmarshal([]byte(respBody), &errResp)
 		if decodeErr != nil {
+			util.Log().Debug("Onedrive返回未知响应[%s]", respBody)
 			return "", sysError(decodeErr)
 		}
 		return "", &errResp

pkg/filesystem/hooks.go

@@ -128,6 +128,16 @@ func HookValidateCapacity(ctx context.Context, fs *FileSystem) error {
 	return nil
 }
 
+// HookValidateCapacityWithoutIncrease 验证用户容量，不扣除
+func HookValidateCapacityWithoutIncrease(ctx context.Context, fs *FileSystem) error {
+	file := ctx.Value(fsctx.FileHeaderCtx).(FileHeader)
+	// 验证并扣除容量
+	if fs.User.GetRemainingCapacity() < file.GetSize() {
+		return ErrInsufficientCapacity
+	}
+	return nil
+}
+
 // HookChangeCapacity 根据原有文件和新文件的大小更新用户容量
 func HookChangeCapacity(ctx context.Context, fs *FileSystem) error {
 	newFile := ctx.Value(fsctx.FileHeaderCtx).(FileHeader)