From c6110e9e75146fe555b63068fc224084629dd1b4 Mon Sep 17 00:00:00 2001
From: HFO4 <912394456@qq.com>
Date: Mon, 23 Nov 2020 18:44:13 +0800
Subject: [PATCH] Feat: keep folder structure in aria2 transferring

---
 pkg/aria2/monitor.go      |  1 +
 pkg/task/tranfer.go       | 25 +++++++++++++++++++------
 pkg/task/transfer_test.go | 29 +++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/pkg/aria2/monitor.go b/pkg/aria2/monitor.go
index bd27f74989..1667f47066 100644
--- a/pkg/aria2/monitor.go
+++ b/pkg/aria2/monitor.go
@@ -251,6 +251,7 @@ func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
 		file,
 		monitor.Task.Dst,
 		monitor.Task.Parent,
+		true,
 	)
 	if err != nil {
 		monitor.setErrorStatus(err)
diff --git a/pkg/task/tranfer.go b/pkg/task/tranfer.go
index 80aa12414c..79e84ac8be 100644
--- a/pkg/task/tranfer.go
+++ b/pkg/task/tranfer.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+	"strings"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
@@ -24,9 +25,11 @@ type TransferTask struct {
 
 // TransferProps 中转任务属性
 type TransferProps struct {
-	Src    []string `json:"src"`    // 原始目录
+	Src    []string `json:"src"`    // 原始文件
 	Parent string   `json:"parent"` // 父目录
 	Dst    string   `json:"dst"`    // 目的目录ID
+	// 将会保留原始文件的目录结构，Src 除去 Parent 开头作为最终路径
+	TrimPath bool `json:"trim_path"`
 }
 
 // Props 获取任务属性
@@ -90,7 +93,16 @@ func (job *TransferTask) Do() {
 
 	for index, file := range job.TaskProps.Src {
 		job.TaskModel.SetProgress(index)
-		err = fs.UploadFromPath(context.Background(), file, path.Join(job.TaskProps.Dst, filepath.Base(file)))
+
+		dst := path.Join(job.TaskProps.Dst, filepath.Base(file))
+		if job.TaskProps.TrimPath {
+			// 保留原始目录
+			trim := util.FormSlash(job.TaskProps.Parent)
+			src := util.FormSlash(file)
+			dst = path.Join(job.TaskProps.Dst, strings.TrimPrefix(src, trim))
+		}
+
+		err = fs.UploadFromPath(context.Background(), file, dst)
 		if err != nil {
 			job.SetErrorMsg("文件转存失败", err)
 		}
@@ -108,7 +120,7 @@ func (job *TransferTask) Recycle() {
 }
 
 // NewTransferTask 新建中转任务
-func NewTransferTask(user uint, src []string, dst, parent string) (Job, error) {
+func NewTransferTask(user uint, src []string, dst, parent string, trim bool) (Job, error) {
 	creator, err := model.GetActiveUserByID(user)
 	if err != nil {
 		return nil, err
@@ -117,9 +129,10 @@ func NewTransferTask(user uint, src []string, dst, parent string) (Job, error) {
 	newTask := &TransferTask{
 		User: &creator,
 		TaskProps: TransferProps{
-			Src:    src,
-			Parent: parent,
-			Dst:    dst,
+			Src:      src,
+			Parent:   parent,
+			Dst:      dst,
+			TrimPath: trim,
 		},
 	}
 
diff --git a/pkg/task/transfer_test.go b/pkg/task/transfer_test.go
index a29043aee4..5d2c9c5ddc 100644
--- a/pkg/task/transfer_test.go
+++ b/pkg/task/transfer_test.go
@@ -102,6 +102,31 @@ func TestTransferTask_Do(t *testing.T) {
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.NotEmpty(task.GetError().Msg)
 	}
+
+	// 替换目录前缀
+	{
+		task.User = &model.User{
+			Policy: model.Policy{
+				Type: "mock",
+			},
+		}
+		task.TaskProps.Src = []string{"test/not_exist"}
+		task.TaskProps.Parent = "test/not_exist"
+		task.TaskProps.TrimPath = true
+		// 更新进度
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1,
+			1))
+		mock.ExpectCommit()
+		// 更新错误
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1,
+			1))
+		mock.ExpectCommit()
+		task.Do()
+		asserts.NoError(mock.ExpectationsWereMet())
+		asserts.NotEmpty(task.GetError().Msg)
+	}
 }
 
 func TestNewTransferTask(t *testing.T) {
@@ -113,7 +138,7 @@ func TestNewTransferTask(t *testing.T) {
 		mock.ExpectBegin()
 		mock.ExpectExec("INSERT(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
 		mock.ExpectCommit()
-		job, err := NewTransferTask(1, []string{}, "/", "/")
+		job, err := NewTransferTask(1, []string{}, "/", "/", false)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.NotNil(job)
 		asserts.NoError(err)
@@ -125,7 +150,7 @@ func TestNewTransferTask(t *testing.T) {
 		mock.ExpectBegin()
 		mock.ExpectExec("INSERT(.+)").WillReturnError(errors.New("error"))
 		mock.ExpectRollback()
-		job, err := NewTransferTask(1, []string{}, "/", "/")
+		job, err := NewTransferTask(1, []string{}, "/", "/", false)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.Nil(job)
 		asserts.Error(err)