github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/service/aria2/add.go (about)

     1  package aria2
     2  
     3  import (
     4  	model "github.com/cloudreve/Cloudreve/v3/models"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
     6  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
     7  	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/monitor"
     8  	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
     9  	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
    10  	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
    11  	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
    12  	"github.com/cloudreve/Cloudreve/v3/pkg/util"
    13  	"github.com/gin-gonic/gin"
    14  )
    15  
    16  // AddURLService 添加URL离线下载服务
    17  type BatchAddURLService struct {
    18  	URLs []string `json:"url" binding:"required"`
    19  	Dst  string   `json:"dst" binding:"required,min=1"`
    20  }
    21  
    22  // Add 主机批量创建新的链接离线下载任务
    23  func (service *BatchAddURLService) Add(c *gin.Context, taskType int) serializer.Response {
    24  	// 创建文件系统
    25  	fs, err := filesystem.NewFileSystemFromContext(c)
    26  	if err != nil {
    27  		return serializer.Err(serializer.CodeCreateFSError, "", err)
    28  	}
    29  	defer fs.Recycle()
    30  
    31  	// 检查用户组权限
    32  	if !fs.User.Group.OptionsSerialized.Aria2 {
    33  		return serializer.Err(serializer.CodeGroupNotAllowed, "", nil)
    34  	}
    35  
    36  	// 存放目录是否存在
    37  	if exist, _ := fs.IsPathExist(service.Dst); !exist {
    38  		return serializer.Err(serializer.CodeParentNotExist, "", nil)
    39  	}
    40  
    41  	// 检查批量任务数量
    42  	limit := fs.User.Group.OptionsSerialized.Aria2BatchSize
    43  	if limit > 0 && len(service.URLs) > limit {
    44  		return serializer.Err(serializer.CodeBatchAria2Size, "", nil)
    45  	}
    46  
    47  	res := make([]serializer.Response, 0, len(service.URLs))
    48  	for _, target := range service.URLs {
    49  		subService := &AddURLService{
    50  			URL: target,
    51  			Dst: service.Dst,
    52  		}
    53  
    54  		addRes := subService.Add(c, fs, taskType)
    55  		res = append(res, addRes)
    56  	}
    57  
    58  	return serializer.Response{Data: res}
    59  }
    60  
    61  // AddURLService 添加URL离线下载服务
    62  type AddURLService struct {
    63  	URL string `json:"url" binding:"required"`
    64  	Dst string `json:"dst" binding:"required,min=1"`
    65  }
    66  
    67  // Add 主机创建新的链接离线下载任务
    68  func (service *AddURLService) Add(c *gin.Context, fs *filesystem.FileSystem, taskType int) serializer.Response {
    69  	if fs == nil {
    70  		var err error
    71  		// 创建文件系统
    72  		fs, err = filesystem.NewFileSystemFromContext(c)
    73  		if err != nil {
    74  			return serializer.Err(serializer.CodeCreateFSError, "", err)
    75  		}
    76  		defer fs.Recycle()
    77  
    78  		// 检查用户组权限
    79  		if !fs.User.Group.OptionsSerialized.Aria2 {
    80  			return serializer.Err(serializer.CodeGroupNotAllowed, "", nil)
    81  		}
    82  
    83  		// 存放目录是否存在
    84  		if exist, _ := fs.IsPathExist(service.Dst); !exist {
    85  			return serializer.Err(serializer.CodeParentNotExist, "", nil)
    86  		}
    87  	}
    88  
    89  	downloads := model.GetDownloadsByStatusAndUser(0, fs.User.ID, common.Downloading, common.Paused, common.Ready)
    90  	limit := fs.User.Group.OptionsSerialized.Aria2BatchSize
    91  	if limit > 0 && len(downloads)+1 > limit {
    92  		return serializer.Err(serializer.CodeBatchAria2Size, "", nil)
    93  	}
    94  
    95  	// 创建任务
    96  	task := &model.Download{
    97  		Status: common.Ready,
    98  		Type:   taskType,
    99  		Dst:    service.Dst,
   100  		UserID: fs.User.ID,
   101  		Source: service.URL,
   102  	}
   103  
   104  	// 获取 Aria2 负载均衡器
   105  	lb := aria2.GetLoadBalancer()
   106  
   107  	// 获取 Aria2 实例
   108  	err, node := cluster.Default.BalanceNodeByFeature("aria2", lb)
   109  	if err != nil {
   110  		return serializer.Err(serializer.CodeInternalSetting, "Failed to get Aria2 instance", err)
   111  	}
   112  
   113  	// 创建任务
   114  	gid, err := node.GetAria2Instance().CreateTask(task, fs.User.Group.OptionsSerialized.Aria2Options)
   115  	if err != nil {
   116  		return serializer.Err(serializer.CodeCreateTaskError, "", err)
   117  	}
   118  
   119  	task.GID = gid
   120  	task.NodeID = node.ID()
   121  	_, err = task.Create()
   122  	if err != nil {
   123  		return serializer.DBErr("Failed to create task record", err)
   124  	}
   125  
   126  	// 创建任务监控
   127  	monitor.NewMonitor(task, cluster.Default, mq.GlobalMQ)
   128  
   129  	return serializer.Response{}
   130  }
   131  
   132  // Add 从机创建新的链接离线下载任务
   133  func Add(c *gin.Context, service *serializer.SlaveAria2Call) serializer.Response {
   134  	caller, _ := c.Get("MasterAria2Instance")
   135  
   136  	// 创建任务
   137  	gid, err := caller.(common.Aria2).CreateTask(service.Task, service.GroupOptions)
   138  	if err != nil {
   139  		return serializer.Err(serializer.CodeInternalSetting, "Failed to create aria2 task", err)
   140  	}
   141  
   142  	// 创建事件通知回调
   143  	siteID, _ := c.Get("MasterSiteID")
   144  	mq.GlobalMQ.SubscribeCallback(gid, func(message mq.Message) {
   145  		if err := cluster.DefaultController.SendNotification(siteID.(string), message.TriggeredBy, message); err != nil {
   146  			util.Log().Warning("Failed to send remote download task status change notifications: %s", err)
   147  		}
   148  	})
   149  
   150  	return serializer.Response{Data: gid}
   151  }