Fix: use X-Cr- as custom header prefix

Fix js heap out of memory
Fix can't find cloudreve for cloudreve has been renamed to Cloudreve

Dockerfile

@@ -5,6 +5,9 @@ COPY ./assets /assets
 
 WORKDIR /assets
 
+# If encountered problems like JavaScript heap out of memory, please uncomment the following options
+ENV NODE_OPTIONS --max_old_space_size=4096
+
 # yarn repo connection is unstable, adjust the network timeout to 10 min.
 RUN set -ex \
     && yarn install --network-timeout 600000 \
@@ -42,7 +45,7 @@ ARG TZ="Asia/Shanghai"
 
 ENV TZ ${TZ}
 
-COPY --from=be-builder /go/bin/cloudreve /cloudreve/cloudreve
+COPY --from=be-builder /go/bin/Cloudreve /cloudreve/cloudreve
 
 RUN apk upgrade \
     && apk add bash tzdata \

bootstrap/init.go

@@ -2,6 +2,7 @@ package bootstrap
 
 import (
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/models/scripts"
 	"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
 	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
 	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
@@ -9,7 +10,7 @@ import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/crontab"
 	"github.com/cloudreve/Cloudreve/v3/pkg/email"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/cloudreve/Cloudreve/v3/pkg/task"
 	"github.com/gin-gonic/gin"
 )
@@ -27,6 +28,12 @@ func Init(path string) {
 		mode    string
 		factory func()
 	}{
+		{
+			"both",
+			func() {
+				scripts.Init()
+			},
+		},
 		{
 			"both",
 			func() {
@@ -54,7 +61,7 @@ func Init(path string) {
 		{
 			"master",
 			func() {
-				aria2.Init(false)
+				aria2.Init(false, cluster.Default, mq.GlobalMQ)
 			},
 		},
 		{
@@ -78,7 +85,7 @@ func Init(path string) {
 		{
 			"slave",
 			func() {
-				slave.Init()
+				cluster.InitController()
 			},
 		},
 		{

bootstrap/script.go

@@ -2,14 +2,14 @@ package bootstrap
 
 import (
 	"context"
-	"github.com/cloudreve/Cloudreve/v3/models/scripts"
+	"github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 )
 
 func RunScript(name string) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	if err := scripts.RunDBScript(name, ctx); err != nil {
+	if err := invoker.RunDBScript(name, ctx); err != nil {
 		util.Log().Error("数据库脚本执行失败: %s", err)
 		return
 	}

go.mod

@@ -38,7 +38,8 @@ require (
 	github.com/tencentcloud/tencentcloud-sdk-go v3.0.125+incompatible
 	github.com/tencentyun/cos-go-sdk-v5 v0.0.0-20200120023323-87ff3bc489ac
 	github.com/upyun/go-sdk v2.1.0+incompatible
-	golang.org/x/text v0.3.2
+	golang.org/x/image v0.0.0-20211028202545-6944b10bf410
+	golang.org/x/text v0.3.6
 	gopkg.in/alexcesaro/quotedprintable.v3 v3.0.0-20150716171945-2caba252f4dc // indirect
 	gopkg.in/go-playground/validator.v9 v9.29.1
 	gopkg.in/ini.v1 v1.51.0 // indirect

go.sum

@@ -12,6 +12,7 @@ github.com/agl/ed25519 v0.0.0-20170116200512-5312a6153412 h1:w1UutsfOrms1J05zt7I
 github.com/agl/ed25519 v0.0.0-20170116200512-5312a6153412/go.mod h1:WPjqKcmVOxf0XSf3YxCJs6N6AOSrOx3obionmG7T0y0=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
+github.com/aliyun/aliyun-oss-go-sdk v2.0.0/go.mod h1:T/Aws4fEfogEE9v+HPhhw+CntffsBHJ8nXQCwKr0/g8=
 github.com/aliyun/aliyun-oss-go-sdk v2.0.5+incompatible h1:A3oZlWPD/Poa19FvNbw+Zu4yKAurDBTjlRDilYGBiS4=
 github.com/aliyun/aliyun-oss-go-sdk v2.0.5+incompatible/go.mod h1:T/Aws4fEfogEE9v+HPhhw+CntffsBHJ8nXQCwKr0/g8=
 github.com/apache/thrift v0.12.0/go.mod h1:cp2SuWMxlEZw2r+iP2GNCdIi4C1qmUzdZFSVb+bacwQ=
@@ -248,6 +249,8 @@ golang.org/x/crypto v0.0.0-20190701094942-4def268fd1a4/go.mod h1:yigFU9vqHzYiE8U
 golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA=
 golang.org/x/image v0.0.0-20190501045829-6d32002ffd75 h1:TbGuee8sSq15Iguxu4deQ7+Bqq/d2rsQejGcEtADAMQ=
 golang.org/x/image v0.0.0-20190501045829-6d32002ffd75/go.mod h1:kZ7UVZpmo3dzQBMxlp+ypCbDeSB+sBbTgSJuh5dn5js=
+golang.org/x/image v0.0.0-20211028202545-6944b10bf410 h1:hTftEOvwiOq2+O8k2D5/Q7COC7k5Qcrgc2TFURJYnvQ=
+golang.org/x/image v0.0.0-20211028202545-6944b10bf410/go.mod h1:023OzeP/+EPmXeapQh35lcL3II3LrY8Ic+EFFKVhULM=
 golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
 golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU=
 golang.org/x/lint v0.0.0-20190301231843-5614ed5bae6f/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE=
@@ -288,6 +291,8 @@ golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.2 h1:tW2bmiBqwgJj/UpqtC8EpXEZVYOwU0yG4iWbprSVAcs=
 golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
+golang.org/x/text v0.3.6 h1:aRYxNxv6iGQlyVaZmk6ZgYEDa+Jg18DxebPSrd6bg1M=
+golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c h1:fqgJT0MGcGpPgpWU7VRdRjuArfcOvC4AoJmILihzhDg=
 golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180828015842-6cd1fcedba52/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

middleware/auth.go

@@ -37,6 +37,7 @@ func SignRequired(authInstance auth.Auth) gin.HandlerFunc {
 			c.Abort()
 			return
 		}
+
 		c.Next()
 	}
 }

middleware/auth_test.go

@@ -90,15 +90,27 @@ func TestSignRequired(t *testing.T) {
 	rec := httptest.NewRecorder()
 	c, _ := gin.CreateTestContext(rec)
 	c.Request, _ = http.NewRequest("GET", "/test", nil)
-	SignRequiredFunc := SignRequired(auth.HMACAuth{SecretKey: []byte(util.RandStringRunes(256))})
+	authInstance := auth.HMACAuth{SecretKey: []byte(util.RandStringRunes(256))}
+	SignRequiredFunc := SignRequired(authInstance)
 
 	// 鉴权失败
 	SignRequiredFunc(c)
 	asserts.NotNil(c)
+	asserts.True(c.IsAborted())
 
+	c, _ = gin.CreateTestContext(rec)
 	c.Request, _ = http.NewRequest("PUT", "/test", nil)
 	SignRequiredFunc(c)
 	asserts.NotNil(c)
+	asserts.True(c.IsAborted())
+
+	// Sign verify success
+	c, _ = gin.CreateTestContext(rec)
+	c.Request, _ = http.NewRequest("PUT", "/test", nil)
+	c.Request = auth.SignRequest(authInstance, c.Request, 0)
+	SignRequiredFunc(c)
+	asserts.NotNil(c)
+	asserts.False(c.IsAborted())
 }
 
 func TestWebDAVAuth(t *testing.T) {
@@ -780,8 +792,6 @@ func TestS3CallbackAuth(t *testing.T) {
 			WillReturnRows(sqlmock.NewRows([]string{"id", "group_id"}).AddRow(1, 1))
 		mock.ExpectQuery("SELECT(.+)groups(.+)").
 			WillReturnRows(sqlmock.NewRows([]string{"id", "policies"}).AddRow(1, "[702]"))
-		mock.ExpectQuery("SELECT(.+)policies(.+)").
-			WillReturnRows(sqlmock.NewRows([]string{"id", "access_key", "secret_key"}).AddRow(2, "123", "123"))
 		c, _ := gin.CreateTestContext(rec)
 		c.Params = []gin.Param{
 			{"key", "testCallBackUpyun"},
@@ -789,5 +799,6 @@ func TestS3CallbackAuth(t *testing.T) {
 		c.Request, _ = http.NewRequest("POST", "/api/v3/callback/upyun/testCallBackUpyun", ioutil.NopCloser(strings.NewReader("1")))
 		AuthFunc(c)
 		asserts.False(c.IsAborted())
+		asserts.NoError(mock.ExpectationsWereMet())
 	}
 }

middleware/cluster.go

@@ -3,7 +3,6 @@ package middleware
 import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/gin-gonic/gin"
 	"strconv"
 )
@@ -11,19 +10,19 @@ import (
 // MasterMetadata 解析主机节点发来请求的包含主机节点信息的元数据
 func MasterMetadata() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		c.Set("MasterSiteID", c.GetHeader("X-Site-Id"))
-		c.Set("MasterSiteURL", c.GetHeader("X-Site-Url"))
-		c.Set("MasterVersion", c.GetHeader("X-Cloudreve-Version"))
+		c.Set("MasterSiteID", c.GetHeader("X-Cr-Site-Id"))
+		c.Set("MasterSiteURL", c.GetHeader("X-Cr-Site-Url"))
+		c.Set("MasterVersion", c.GetHeader("X-Cr-Cloudreve-Version"))
 		c.Next()
 	}
 }
 
 // UseSlaveAria2Instance 从机用于获取对应主机节点的Aria2实例
-func UseSlaveAria2Instance() gin.HandlerFunc {
+func UseSlaveAria2Instance(clusterController cluster.Controller) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		if siteID, exist := c.Get("MasterSiteID"); exist {
 			// 获取对应主机节点的从机Aria2实例
-			caller, err := slave.DefaultController.GetAria2Instance(siteID.(string))
+			caller, err := clusterController.GetAria2Instance(siteID.(string))
 			if err != nil {
 				c.JSON(200, serializer.Err(serializer.CodeNotSet, "无法获取 Aria2 实例", err))
 				c.Abort()
@@ -40,16 +39,16 @@ func UseSlaveAria2Instance() gin.HandlerFunc {
 	}
 }
 
-func SlaveRPCSignRequired() gin.HandlerFunc {
+func SlaveRPCSignRequired(nodePool cluster.Pool) gin.HandlerFunc {
 	return func(c *gin.Context) {
-		nodeID, err := strconv.ParseUint(c.GetHeader("X-Node-Id"), 10, 64)
+		nodeID, err := strconv.ParseUint(c.GetHeader("X-Cr-Node-Id"), 10, 64)
 		if err != nil {
 			c.JSON(200, serializer.ParamErr("未知的主机节点ID", err))
 			c.Abort()
 			return
 		}
 
-		slaveNode := cluster.Default.GetNodeByID(uint(nodeID))
+		slaveNode := nodePool.GetNodeByID(uint(nodeID))
 		if slaveNode == nil {
 			c.JSON(200, serializer.ParamErr("未知的主机节点ID", err))
 			c.Abort()

middleware/cluster_test.go

@@ -0,0 +1,120 @@
+package middleware
+
+import (
+	"errors"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
+	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mocks/controllermock"
+	"github.com/gin-gonic/gin"
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestMasterMetadata(t *testing.T) {
+	a := assert.New(t)
+	masterMetaDataFunc := MasterMetadata()
+	rec := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(rec)
+	c.Request = httptest.NewRequest("GET", "/", nil)
+
+	c.Request.Header = map[string][]string{
+		"X-Cr-Site-Id":           {"expectedSiteID"},
+		"X-Cr-Site-Url":          {"expectedSiteURL"},
+		"X-Cr-Cloudreve-Version": {"expectedMasterVersion"},
+	}
+	masterMetaDataFunc(c)
+	siteID, _ := c.Get("MasterSiteID")
+	siteURL, _ := c.Get("MasterSiteURL")
+	siteVersion, _ := c.Get("MasterVersion")
+
+	a.Equal("expectedSiteID", siteID.(string))
+	a.Equal("expectedSiteURL", siteURL.(string))
+	a.Equal("expectedMasterVersion", siteVersion.(string))
+}
+
+func TestSlaveRPCSignRequired(t *testing.T) {
+	a := assert.New(t)
+	np := &cluster.NodePool{}
+	np.Init()
+	slaveRPCSignRequiredFunc := SlaveRPCSignRequired(np)
+	rec := httptest.NewRecorder()
+
+	// id parse failed
+	{
+		c, _ := gin.CreateTestContext(rec)
+		c.Request = httptest.NewRequest("GET", "/", nil)
+		c.Request.Header.Set("X-Cr-Node-Id", "unknown")
+		slaveRPCSignRequiredFunc(c)
+		a.True(c.IsAborted())
+	}
+
+	// node id not exist
+	{
+		c, _ := gin.CreateTestContext(rec)
+		c.Request = httptest.NewRequest("GET", "/", nil)
+		c.Request.Header.Set("X-Cr-Node-Id", "38")
+		slaveRPCSignRequiredFunc(c)
+		a.True(c.IsAborted())
+	}
+
+	// success
+	{
+		authInstance := auth.HMACAuth{SecretKey: []byte("")}
+		np.Add(&model.Node{Model: gorm.Model{
+			ID: 38,
+		}})
+
+		c, _ := gin.CreateTestContext(rec)
+		c.Request = httptest.NewRequest("POST", "/", nil)
+		c.Request.Header.Set("X-Cr-Node-Id", "38")
+		c.Request = auth.SignRequest(authInstance, c.Request, 0)
+		slaveRPCSignRequiredFunc(c)
+		a.False(c.IsAborted())
+	}
+}
+
+func TestUseSlaveAria2Instance(t *testing.T) {
+	a := assert.New(t)
+
+	// MasterSiteID not set
+	{
+		testController := &controllermock.SlaveControllerMock{}
+		useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
+		c, _ := gin.CreateTestContext(httptest.NewRecorder())
+		c.Request = httptest.NewRequest("GET", "/", nil)
+		useSlaveAria2InstanceFunc(c)
+		a.True(c.IsAborted())
+	}
+
+	// Cannot get aria2 instances
+	{
+		testController := &controllermock.SlaveControllerMock{}
+		useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
+		c, _ := gin.CreateTestContext(httptest.NewRecorder())
+		c.Request = httptest.NewRequest("GET", "/", nil)
+		c.Set("MasterSiteID", "expectedSiteID")
+		testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, errors.New("error"))
+		useSlaveAria2InstanceFunc(c)
+		a.True(c.IsAborted())
+		testController.AssertExpectations(t)
+	}
+
+	// Success
+	{
+		testController := &controllermock.SlaveControllerMock{}
+		useSlaveAria2InstanceFunc := UseSlaveAria2Instance(testController)
+		c, _ := gin.CreateTestContext(httptest.NewRecorder())
+		c.Request = httptest.NewRequest("GET", "/", nil)
+		c.Set("MasterSiteID", "expectedSiteID")
+		testController.On("GetAria2Instance", "expectedSiteID").Return(&common.DummyAria2{}, nil)
+		useSlaveAria2InstanceFunc(c)
+		a.False(c.IsAborted())
+		res, _ := c.Get("MasterAria2Instance")
+		a.NotNil(res)
+		testController.AssertExpectations(t)
+	}
+}

models/download_test.go

@@ -177,3 +177,14 @@ func TestDownload_Delete(t *testing.T) {
 	}
 
 }
+
+func TestDownload_GetNodeID(t *testing.T) {
+	a := assert.New(t)
+	record := Download{}
+
+	// compatible with 3.4
+	a.EqualValues(1, record.GetNodeID())
+
+	record.NodeID = 5
+	a.EqualValues(5, record.GetNodeID())
+}

models/migration.go

@@ -1,12 +1,17 @@
 package model
 
 import (
+	"context"
+	"github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
 	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 	"github.com/fatih/color"
 	"github.com/gofrs/uuid"
+	"github.com/hashicorp/go-version"
 	"github.com/jinzhu/gorm"
+	"sort"
+	"strings"
 )
 
 // 是否需要迁移
@@ -54,6 +59,9 @@ func migration() {
 	// 向设置数据表添加初始设置
 	addDefaultSettings()
 
+	// 执行数据库升级脚本
+	execUpgradeScripts()
+
 	util.Log().Info("数据库初始化结束")
 
 }
@@ -181,6 +189,7 @@ Neue',Helvetica,Arial,sans-serif; box-sizing: border-box; font-size: 14px; verti
 		{Name: "pwa_display", Value: "standalone", Type: "pwa"},
 		{Name: "pwa_theme_color", Value: "#000000", Type: "pwa"},
 		{Name: "pwa_background_color", Value: "#ffffff", Type: "pwa"},
+		{Name: "office_preview_service", Value: "https://view.officeapps.live.com/op/view.aspx?src={$src}", Type: "preview"},
 	}
 
 	for _, value := range defaultSettings {
@@ -290,3 +299,17 @@ func addDefaultNode() {
 		}
 	}
 }
+
+func execUpgradeScripts() {
+	s := invoker.ListPrefix("UpgradeTo")
+	versions := make([]*version.Version, len(s))
+	for i, raw := range s {
+		v, _ := version.NewVersion(strings.TrimPrefix(raw, "UpgradeTo"))
+		versions[i] = v
+	}
+	sort.Sort(version.Collection(versions))
+
+	for i := 0; i < len(versions); i++ {
+		invoker.RunDBScript("UpgradeTo"+versions[i].String(), context.Background())
+	}
+}

models/node_test.go

@@ -0,0 +1,64 @@
+package model
+
+import (
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestGetNodeByID(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+	res, err := GetNodeByID(1)
+	a.NoError(err)
+	a.EqualValues(1, res.ID)
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestGetNodesByStatus(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"status"}).AddRow(NodeActive))
+	res, err := GetNodesByStatus(NodeActive)
+	a.NoError(err)
+	a.Len(res, 1)
+	a.EqualValues(NodeActive, res[0].Status)
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestNode_AfterFind(t *testing.T) {
+	a := assert.New(t)
+	node := &Node{}
+
+	// No aria2 options
+	{
+		a.NoError(node.AfterFind())
+	}
+
+	// with aria2 options
+	{
+		node.Aria2Options = `{"timeout":1}`
+		a.NoError(node.AfterFind())
+		a.Equal(1, node.Aria2OptionsSerialized.Timeout)
+	}
+}
+
+func TestNode_BeforeSave(t *testing.T) {
+	a := assert.New(t)
+	node := &Node{}
+
+	node.Aria2OptionsSerialized.Timeout = 1
+	a.NoError(node.BeforeSave())
+	a.Contains(node.Aria2Options, "1")
+}
+
+func TestNode_SetStatus(t *testing.T) {
+	a := assert.New(t)
+	node := &Node{}
+
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)nodes").WithArgs(NodeActive, sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+	a.NoError(node.SetStatus(NodeActive))
+	a.Equal(NodeActive, node.Status)
+	a.NoError(mock.ExpectationsWereMet())
+}

models/policy_test.go

@@ -327,3 +327,18 @@ func TestPolicy_IsThumbExist(t *testing.T) {
 		asserts.Equal(testCase.expect, policy.IsThumbExist(testCase.name))
 	}
 }
+
+func TestPolicy_UpdateAccessKeyAndClearCache(t *testing.T) {
+	a := assert.New(t)
+	cache.Set("policy_1331", Policy{}, 3600)
+	p := &Policy{}
+	p.ID = 1331
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WithArgs("ak", sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.NoError(p.UpdateAccessKeyAndClearCache("ak"))
+	a.NoError(mock.ExpectationsWereMet())
+	_, ok := cache.Get("policy_1331")
+	a.False(ok)
+}

models/scripts/init.go

@@ -0,0 +1,9 @@
+package scripts
+
+import "github.com/cloudreve/Cloudreve/v3/models/scripts/invoker"
+
+func Init() {
+	invoker.Register("ResetAdminPassword", ResetAdminPassword(0))
+	invoker.Register("CalibrateUserStorage", UserStorageCalibration(0))
+	invoker.Register("UpgradeTo3.4.0", UpgradeTo340(0))
+}

models/scripts/invoker.go

@@ -1,25 +0,0 @@
-package scripts
-
-import (
-	"context"
-	"fmt"
-)
-
-type DBScript interface {
-	Run(ctx context.Context)
-}
-
-var availableScripts = make(map[string]DBScript)
-
-func RunDBScript(name string, ctx context.Context) error {
-	if script, ok := availableScripts[name]; ok {
-		script.Run(ctx)
-		return nil
-	}
-
-	return fmt.Errorf("数据库脚本 [%s] 不存在", name)
-}
-
-func register(name string, script DBScript) {
-	availableScripts[name] = script
-}

models/scripts/invoker/invoker.go

@@ -0,0 +1,38 @@
+package invoker
+
+import (
+	"context"
+	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"strings"
+)
+
+type DBScript interface {
+	Run(ctx context.Context)
+}
+
+var availableScripts = make(map[string]DBScript)
+
+func RunDBScript(name string, ctx context.Context) error {
+	if script, ok := availableScripts[name]; ok {
+		util.Log().Info("开始执行数据库脚本 [%s]", name)
+		script.Run(ctx)
+		return nil
+	}
+
+	return fmt.Errorf("数据库脚本 [%s] 不存在", name)
+}
+
+func Register(name string, script DBScript) {
+	availableScripts[name] = script
+}
+
+func ListPrefix(prefix string) []string {
+	var scripts []string
+	for name := range availableScripts {
+		if strings.HasPrefix(name, prefix) {
+			scripts = append(scripts, name)
+		}
+	}
+	return scripts
+}

models/scripts/invoker/invoker_test.go

@@ -0,0 +1,39 @@
+package invoker
+
+import (
+	"context"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+type TestScript int
+
+func (script TestScript) Run(ctx context.Context) {
+
+}
+
+func TestRunDBScript(t *testing.T) {
+	asserts := assert.New(t)
+	Register("test", TestScript(0))
+
+	// 不存在
+	{
+		asserts.Error(RunDBScript("else", context.Background()))
+	}
+
+	// 存在
+	{
+		asserts.NoError(RunDBScript("test", context.Background()))
+	}
+}
+
+func TestListPrefix(t *testing.T) {
+	asserts := assert.New(t)
+	Register("U1", TestScript(0))
+	Register("U2", TestScript(0))
+	Register("U3", TestScript(0))
+	Register("P1", TestScript(0))
+
+	res := ListPrefix("U")
+	asserts.Len(res, 3)
+}

models/scripts/invoker_test.go

@@ -1,49 +0,0 @@
-package scripts
-
-import (
-	"context"
-	"database/sql"
-	"github.com/DATA-DOG/go-sqlmock"
-	model "github.com/cloudreve/Cloudreve/v3/models"
-	"github.com/jinzhu/gorm"
-	"github.com/stretchr/testify/assert"
-	"testing"
-)
-
-var mock sqlmock.Sqlmock
-var mockDB *gorm.DB
-
-type TestScript int
-
-func (script TestScript) Run(ctx context.Context) {
-
-}
-
-// TestMain 初始化数据库Mock
-func TestMain(m *testing.M) {
-	var db *sql.DB
-	var err error
-	db, mock, err = sqlmock.New()
-	if err != nil {
-		panic("An error was not expected when opening a stub database connection")
-	}
-	model.DB, _ = gorm.Open("mysql", db)
-	mockDB = model.DB
-	defer db.Close()
-	m.Run()
-}
-
-func TestRunDBScript(t *testing.T) {
-	asserts := assert.New(t)
-	register("test", TestScript(0))
-
-	// 不存在
-	{
-		asserts.Error(RunDBScript("else", context.Background()))
-	}
-
-	// 存在
-	{
-		asserts.NoError(RunDBScript("test", context.Background()))
-	}
-}

models/scripts/reset.go

@@ -9,10 +9,6 @@ import (
 
 type ResetAdminPassword int
 
-func init() {
-	register("ResetAdminPassword", ResetAdminPassword(0))
-}
-
 // Run 运行脚本从社区版升级至 Pro 版
 func (script ResetAdminPassword) Run(ctx context.Context) {
 	// 查找用户

models/scripts/storage.go

@@ -8,10 +8,6 @@ import (
 
 type UserStorageCalibration int
 
-func init() {
-	register("CalibrateUserStorage", UserStorageCalibration(0))
-}
-
 type storageResult struct {
 	Total uint64
 }

models/scripts/storage_test.go

@@ -2,11 +2,31 @@ package scripts
 
 import (
 	"context"
+	"database/sql"
 	"github.com/DATA-DOG/go-sqlmock"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 	"testing"
 )
 
+var mock sqlmock.Sqlmock
+var mockDB *gorm.DB
+
+// TestMain 初始化数据库Mock
+func TestMain(m *testing.M) {
+	var db *sql.DB
+	var err error
+	db, mock, err = sqlmock.New()
+	if err != nil {
+		panic("An error was not expected when opening a stub database connection")
+	}
+	model.DB, _ = gorm.Open("mysql", db)
+	mockDB = model.DB
+	defer db.Close()
+	m.Run()
+}
+
 func TestUserStorageCalibration_Run(t *testing.T) {
 	asserts := assert.New(t)
 	script := UserStorageCalibration(0)

models/scripts/upgrade.go

@@ -0,0 +1,43 @@
+package scripts
+
+import (
+	"context"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"strconv"
+)
+
+type UpgradeTo340 int
+
+// Run upgrade from older version to 3.4.0
+func (script UpgradeTo340) Run(ctx context.Context) {
+	// 取回老版本 aria2 设定
+	old := model.GetSettingByType([]string{"aria2"})
+	if len(old) == 0 {
+		return
+	}
+
+	// 写入到新版本的节点设定
+	n, err := model.GetNodeByID(1)
+	if err != nil {
+		util.Log().Error("找不到主机节点, %s", err)
+	}
+
+	n.Aria2Enabled = old["aria2_rpcurl"] != ""
+	n.Aria2OptionsSerialized.Options = old["aria2_options"]
+	n.Aria2OptionsSerialized.Server = old["aria2_rpcurl"]
+
+	interval, err := strconv.Atoi(old["aria2_interval"])
+	if err != nil {
+		interval = 10
+	}
+	n.Aria2OptionsSerialized.Interval = interval
+	n.Aria2OptionsSerialized.TempPath = old["aria2_temp_path"]
+	n.Aria2OptionsSerialized.Token = old["aria2_token"]
+	if err := model.DB.Save(&n).Error; err != nil {
+		util.Log().Error("无法保存主机节点 Aria2 配置信息, %s", err)
+	} else {
+		model.DB.Where("type = ?", "aria2").Delete(model.Setting{})
+		util.Log().Info("Aria2 配置信息已成功迁移至 3.4.0+ 版本的模式")
+	}
+}

models/scripts/upgrade_test.go

@@ -0,0 +1,66 @@
+package scripts
+
+import (
+	"context"
+	"errors"
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestUpgradeTo340_Run(t *testing.T) {
+	a := assert.New(t)
+	script := UpgradeTo340(0)
+
+	// skip
+	{
+		mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name"}))
+		script.Run(context.Background())
+		a.NoError(mock.ExpectationsWereMet())
+	}
+
+	// node not found
+	{
+		mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name"}).AddRow("1"))
+		mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}))
+		script.Run(context.Background())
+		a.NoError(mock.ExpectationsWereMet())
+	}
+
+	// success
+	{
+		mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name", "value"}).
+			AddRow("aria2_rpcurl", "expected_aria2_rpcurl").
+			AddRow("aria2_interval", "expected_aria2_interval").
+			AddRow("aria2_temp_path", "expected_aria2_temp_path").
+			AddRow("aria2_token", "expected_aria2_token").
+			AddRow("aria2_options", "{}"))
+
+		mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+		script.Run(context.Background())
+		a.NoError(mock.ExpectationsWereMet())
+	}
+
+	// failed
+	{
+		mock.ExpectQuery("SELECT(.+)settings").WillReturnRows(sqlmock.NewRows([]string{"name", "value"}).
+			AddRow("aria2_rpcurl", "expected_aria2_rpcurl").
+			AddRow("aria2_interval", "expected_aria2_interval").
+			AddRow("aria2_temp_path", "expected_aria2_temp_path").
+			AddRow("aria2_token", "expected_aria2_token").
+			AddRow("aria2_options", "{}"))
+
+		mock.ExpectQuery("SELECT(.+)nodes").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
+		mock.ExpectRollback()
+		script.Run(context.Background())
+		a.NoError(mock.ExpectationsWereMet())
+	}
+}

models/share_test.go

@@ -188,18 +188,6 @@ func TestShare_CanBeDownloadBy(t *testing.T) {
 		asserts.Error(share.CanBeDownloadBy(user))
 	}
 
-	// 未登录，需要积分
-	{
-		user := &User{
-			Group: Group{
-				OptionsSerialized: GroupOption{
-					ShareDownload: true,
-				},
-			},
-		}
-		asserts.Error(share.CanBeDownloadBy(user))
-	}
-
 	// 成功
 	{
 		user := &User{

models/tag_test.go

@@ -55,9 +55,9 @@ func TestGetTagsByUID(t *testing.T) {
 
 func TestGetTagsByID(t *testing.T) {
 	asserts := assert.New(t)
-	mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
-	res, err := GetTasksByID(1)
+	mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"name"}).AddRow("tag"))
+	res, err := GetTagsByUID(1)
 	asserts.NoError(mock.ExpectationsWereMet())
 	asserts.NoError(err)
-	asserts.EqualValues(1, res.ID)
+	asserts.EqualValues("tag", res[0].Name)
 }

models/task_test.go

@@ -91,3 +91,14 @@ func TestListTasks(t *testing.T) {
 	asserts.EqualValues(5, total)
 	asserts.Len(res, 1)
 }
+
+func TestGetTasksByStatus(t *testing.T) {
+	a := assert.New(t)
+
+	mock.ExpectQuery("SELECT(.+)").
+		WithArgs(1, 2).
+		WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+	res := GetTasksByStatus(1, 2)
+	a.NoError(mock.ExpectationsWereMet())
+	a.Len(res, 1)
+}

models/user.go

@@ -35,7 +35,7 @@ type User struct {
 	Storage   uint64
 	TwoFactor string
 	Avatar    string
-	Options   string `json:"-",gorm:"type:text"`
+	Options   string `json:"-" gorm:"type:text"`
 	Authn     string `gorm:"type:text"`
 
 	// 关联模型

models/user_test.go

@@ -177,10 +177,10 @@ func TestNewUser(t *testing.T) {
 
 func TestUser_AfterFind(t *testing.T) {
 	asserts := assert.New(t)
-	cache.Deletes([]string{"1"}, "policy_")
+	cache.Deletes([]string{"0"}, "policy_")
 
 	policyRows := sqlmock.NewRows([]string{"id", "name"}).
-		AddRow(1, "默认存储策略")
+		AddRow(144, "默认存储策略")
 	mock.ExpectQuery("^SELECT (.+)").WillReturnRows(policyRows)
 
 	newUser := NewUser()
@@ -240,11 +240,6 @@ func TestUser_GetRemainingCapacity(t *testing.T) {
 	newUser.Group.MaxStorage = 100
 	newUser.Storage = 200
 	asserts.Equal(uint64(0), newUser.GetRemainingCapacity())
-
-	cache.Set("pack_size_0", uint64(10), 0)
-	newUser.Group.MaxStorage = 100
-	newUser.Storage = 101
-	asserts.Equal(uint64(9), newUser.GetRemainingCapacity())
 }
 
 func TestUser_DeductionCapacity(t *testing.T) {
@@ -280,10 +275,6 @@ func TestUser_DeductionCapacity(t *testing.T) {
 	asserts.Equal(false, newUser.IncreaseStorage(1))
 	asserts.Equal(uint64(100), newUser.Storage)
 
-	cache.Set("pack_size_1", uint64(1), 0)
-	asserts.Equal(true, newUser.IncreaseStorage(1))
-	asserts.Equal(uint64(101), newUser.Storage)
-
 	asserts.True(newUser.IncreaseStorage(0))
 }
 

models/webdav_test.go

@@ -55,6 +55,6 @@ func TestDeleteWebDAVAccountByID(t *testing.T) {
 	mock.ExpectBegin()
 	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
 	mock.ExpectCommit()
-	asserts.NoError(DeleteTagByID(1, 1))
+	DeleteWebDAVAccountByID(1, 1)
 	asserts.NoError(mock.ExpectationsWereMet())
 }

pkg/aria2/aria2.go

@@ -3,6 +3,8 @@ package aria2
 import (
 	"context"
 	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"net/url"
 	"sync"
 	"time"
@@ -31,7 +33,7 @@ func GetLoadBalancer() balancer.Balancer {
 }
 
 // Init 初始化
-func Init(isReload bool) {
+func Init(isReload bool, pool cluster.Pool, mqClient mq.MQ) {
 	Lock.Lock()
 	LB = balancer.NewBalancer("RoundRobin")
 	Lock.Unlock()
@@ -42,7 +44,7 @@ func Init(isReload bool) {
 
 		for i := 0; i < len(unfinished); i++ {
 			// 创建任务监控
-			monitor.NewMonitor(&unfinished[i])
+			monitor.NewMonitor(&unfinished[i], pool, mqClient)
 		}
 	}
 }

pkg/aria2/aria2_test.go

@@ -2,14 +2,15 @@ package aria2
 
 import (
 	"database/sql"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mocks"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
+	"github.com/stretchr/testify/assert"
+	testMock "github.com/stretchr/testify/mock"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
 	model "github.com/cloudreve/Cloudreve/v3/models"
-	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/monitor"
-	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"github.com/jinzhu/gorm"
-	"github.com/stretchr/testify/assert"
 )
 
 var mock sqlmock.Sqlmock
@@ -27,66 +28,39 @@ func TestMain(m *testing.M) {
 	m.Run()
 }
 
-func TestDummyAria2(t *testing.T) {
-	asserts := assert.New(t)
-	instance := DummyAria2{}
-	asserts.Error(instance.CreateTask(nil, nil))
-	_, err := instance.Status(nil)
-	asserts.Error(err)
-	asserts.Error(instance.Cancel(nil))
-	asserts.Error(instance.Select(nil, nil))
-}
-
 func TestInit(t *testing.T) {
-	monitor.MAX_RETRY = 0
-	asserts := assert.New(t)
-	cache.Set("setting_aria2_token", "1", 0)
-	cache.Set("setting_aria2_call_timeout", "5", 0)
-	cache.Set("setting_aria2_options", `[]`, 0)
+	a := assert.New(t)
+	mockPool := &mocks.NodePoolMock{}
+	mockPool.On("GetNodeByID", testMock.Anything).Return(nil)
+	mockQueue := mq.NewMQ()
 
-	// 未指定RPC地址，跳过
+	mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+	Init(false, mockPool, mockQueue)
+	a.NoError(mock.ExpectationsWereMet())
+	mockPool.AssertExpectations(t)
+}
+
+func TestTestRPCConnection(t *testing.T) {
+	a := assert.New(t)
+
+	// url not legal
 	{
-		cache.Set("setting_aria2_rpcurl", "", 0)
-		Init(false)
-		asserts.IsType(&DummyAria2{}, Instance)
+		res, err := TestRPCConnection(string([]byte{0x7f}), "", 10)
+		a.Error(err)
+		a.Empty(res.Version)
 	}
 
-	// 无法解析服务器地址
+	// rpc failed
 	{
-		cache.Set("setting_aria2_rpcurl", string(byte(0x7f)), 0)
-		Init(false)
-		asserts.IsType(&DummyAria2{}, Instance)
-	}
-
-	// 无法解析全局配置
-	{
-		Instance = &RPCService{}
-		cache.Set("setting_aria2_options", "?", 0)
-		cache.Set("setting_aria2_rpcurl", "ws://127.0.0.1:1234", 0)
-		Init(false)
-		asserts.IsType(&DummyAria2{}, Instance)
-	}
-
-	// 连接失败
-	{
-		cache.Set("setting_aria2_options", "{}", 0)
-		cache.Set("setting_aria2_rpcurl", "http://127.0.0.1:1234", 0)
-		cache.Set("setting_aria2_call_timeout", "1", 0)
-		cache.Set("setting_aria2_interval", "100", 0)
-		mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"g_id"}).AddRow("1"))
-		Init(false)
-		asserts.NoError(mock.ExpectationsWereMet())
-		asserts.IsType(&RPCService{}, Instance)
+		res, err := TestRPCConnection("ws://0.0.0.0", "", 0)
+		a.Error(err)
+		a.Empty(res.Version)
 	}
 }
 
-func TestGetStatus(t *testing.T) {
-	asserts := assert.New(t)
-	asserts.Equal(4, GetStatus("complete"))
-	asserts.Equal(1, GetStatus("active"))
-	asserts.Equal(0, GetStatus("waiting"))
-	asserts.Equal(2, GetStatus("paused"))
-	asserts.Equal(3, GetStatus("error"))
-	asserts.Equal(5, GetStatus("removed"))
-	asserts.Equal(6, GetStatus("?"))
+func TestGetLoadBalancer(t *testing.T) {
+	a := assert.New(t)
+	a.NotPanics(func() {
+		GetLoadBalancer()
+	})
 }

pkg/aria2/caller.go

@@ -1,114 +0,0 @@
-package aria2
-
-import (
-	"context"
-	"path/filepath"
-	"strconv"
-	"strings"
-	"time"
-
-	model "github.com/cloudreve/Cloudreve/v3/models"
-	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
-	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
-	"github.com/cloudreve/Cloudreve/v3/pkg/util"
-)
-
-// RPCService 通过RPC服务的Aria2任务管理器
-type RPCService struct {
-	options *clientOptions
-	Caller  rpc.Client
-}
-
-type clientOptions struct {
-	Options map[string]interface{} // 创建下载时额外添加的设置
-}
-
-// Init 初始化
-func (client *RPCService) Init(server, secret string, timeout int, options map[string]interface{}) error {
-	// 客户端已存在，则关闭先前连接
-	if client.Caller != nil {
-		client.Caller.Close()
-	}
-
-	client.options = &clientOptions{
-		Options: options,
-	}
-	caller, err := rpc.New(context.Background(), server, secret, time.Duration(timeout)*time.Second,
-		mq.GlobalMQ)
-	client.Caller = caller
-	return err
-}
-
-// Status 查询下载状态
-func (client *RPCService) Status(task *model.Download) (rpc.StatusInfo, error) {
-	res, err := client.Caller.TellStatus(task.GID)
-	if err != nil {
-		// 失败后重试
-		util.Log().Debug("无法获取离线下载状态，%s，10秒钟后重试", err)
-		time.Sleep(time.Duration(10) * time.Second)
-		res, err = client.Caller.TellStatus(task.GID)
-	}
-
-	return res, err
-}
-
-// Cancel 取消下载
-func (client *RPCService) Cancel(task *model.Download) error {
-	// 取消下载任务
-	_, err := client.Caller.Remove(task.GID)
-	if err != nil {
-		util.Log().Warning("无法取消离线下载任务[%s], %s", task.GID, err)
-	}
-
-	//// 删除临时文件
-	//util.Log().Debug("离线下载任务[%s]已取消，1 分钟后删除临时文件", task.GID)
-	//go func(task *model.Download) {
-	//	select {
-	//	case <-time.After(time.Duration(60) * time.Second):
-	//		err := os.RemoveAll(task.Parent)
-	//		if err != nil {
-	//			util.Log().Warning("无法删除离线下载临时目录[%s], %s", task.Parent, err)
-	//		}
-	//	}
-	//}(task)
-
-	return err
-}
-
-// Select 选取要下载的文件
-func (client *RPCService) Select(task *model.Download, files []int) error {
-	var selected = make([]string, len(files))
-	for i := 0; i < len(files); i++ {
-		selected[i] = strconv.Itoa(files[i])
-	}
-	_, err := client.Caller.ChangeOption(task.GID, map[string]interface{}{"select-file": strings.Join(selected, ",")})
-	return err
-}
-
-// CreateTask 创建新任务
-func (client *RPCService) CreateTask(task *model.Download, groupOptions map[string]interface{}) (string, error) {
-	// 生成存储路径
-	path := filepath.Join(
-		model.GetSettingByName("aria2_temp_path"),
-		"aria2",
-		strconv.FormatInt(time.Now().UnixNano(), 10),
-	)
-
-	// 创建下载任务
-	options := map[string]interface{}{
-		"dir": path,
-	}
-	for k, v := range client.options.Options {
-		options[k] = v
-	}
-	for k, v := range groupOptions {
-		options[k] = v
-	}
-
-	gid, err := client.Caller.AddURI(task.Source, options)
-	if err != nil || gid == "" {
-		return "", err
-	}
-
-	return gid, nil
-}

pkg/aria2/caller_test.go

@@ -1,52 +0,0 @@
-package aria2
-
-import (
-	"testing"
-
-	model "github.com/cloudreve/Cloudreve/v3/models"
-	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestRPCService_Init(t *testing.T) {
-	asserts := assert.New(t)
-	caller := &RPCService{}
-	asserts.Error(caller.Init("ws://", "", 1, nil))
-	asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
-}
-
-func TestRPCService_Status(t *testing.T) {
-	asserts := assert.New(t)
-	caller := &RPCService{}
-	asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
-
-	_, err := caller.Status(&model.Download{})
-	asserts.Error(err)
-}
-
-func TestRPCService_Cancel(t *testing.T) {
-	asserts := assert.New(t)
-	caller := &RPCService{}
-	asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
-
-	err := caller.Cancel(&model.Download{Parent: "test"})
-	asserts.Error(err)
-}
-
-func TestRPCService_Select(t *testing.T) {
-	asserts := assert.New(t)
-	caller := &RPCService{}
-	asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
-
-	err := caller.Select(&model.Download{Parent: "test"}, []int{1, 2, 3})
-	asserts.Error(err)
-}
-
-func TestRPCService_CreateTask(t *testing.T) {
-	asserts := assert.New(t)
-	caller := &RPCService{}
-	asserts.NoError(caller.Init("http://127.0.0.1", "", 1, nil))
-	cache.Set("setting_aria2_temp_path", "test", 0)
-	err := caller.CreateTask(&model.Download{Parent: "test"}, map[string]interface{}{"1": "1"})
-	asserts.Error(err)
-}

pkg/aria2/common/common_test.go

@@ -0,0 +1,45 @@
+package common
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestDummyAria2(t *testing.T) {
+	a := assert.New(t)
+	d := &DummyAria2{}
+
+	a.NoError(d.Init())
+
+	res, err := d.CreateTask(&model.Download{}, map[string]interface{}{})
+	a.Empty(res)
+	a.Error(err)
+
+	_, err = d.Status(&model.Download{})
+	a.Error(err)
+
+	err = d.Cancel(&model.Download{})
+	a.Error(err)
+
+	err = d.Select(&model.Download{}, []int{})
+	a.Error(err)
+
+	configRes := d.GetConfig()
+	a.NotNil(configRes)
+
+	err = d.DeleteTempFile(&model.Download{})
+	a.Error(err)
+}
+
+func TestGetStatus(t *testing.T) {
+	a := assert.New(t)
+
+	a.Equal(GetStatus("complete"), Complete)
+	a.Equal(GetStatus("active"), Downloading)
+	a.Equal(GetStatus("waiting"), Ready)
+	a.Equal(GetStatus("paused"), Paused)
+	a.Equal(GetStatus("error"), Error)
+	a.Equal(GetStatus("removed"), Canceled)
+	a.Equal(GetStatus("unknown"), Unknown)
+}

pkg/aria2/monitor/monitor.go

@@ -33,29 +33,29 @@ type Monitor struct {
 var MAX_RETRY = 10
 
 // NewMonitor 新建离线下载状态监控
-func NewMonitor(task *model.Download) {
+func NewMonitor(task *model.Download, pool cluster.Pool, mqClient mq.MQ) {
 	monitor := &Monitor{
 		Task:     task,
 		notifier: make(chan mq.Message),
-		node:     cluster.Default.GetNodeByID(task.GetNodeID()),
+		node:     pool.GetNodeByID(task.GetNodeID()),
 	}
 
 	if monitor.node != nil {
 		monitor.Interval = time.Duration(monitor.node.GetAria2Instance().GetConfig().Interval) * time.Second
-		go monitor.Loop()
+		go monitor.Loop(mqClient)
 
-		monitor.notifier = mq.GlobalMQ.Subscribe(monitor.Task.GID, 0)
+		monitor.notifier = mqClient.Subscribe(monitor.Task.GID, 0)
 	} else {
 		monitor.setErrorStatus(errors.New("节点不可用"))
 	}
 }
 
 // Loop 开启监控循环
-func (monitor *Monitor) Loop() {
-	defer mq.GlobalMQ.Unsubscribe(monitor.Task.GID, monitor.notifier)
+func (monitor *Monitor) Loop(mqClient mq.MQ) {
+	defer mqClient.Unsubscribe(monitor.Task.GID, monitor.notifier)
 
 	// 首次循环立即更新
-	interval := time.Duration(0)
+	interval := 50 * time.Millisecond
 
 	for {
 		select {
@@ -112,7 +112,7 @@ func (monitor *Monitor) Update() bool {
 
 	switch status.Status {
 	case "complete":
-		return monitor.Complete(status)
+		return monitor.Complete(task.TaskPoll)
 	case "error":
 		return monitor.Error(status)
 	case "active", "waiting", "paused":
@@ -235,7 +235,7 @@ func (monitor *Monitor) RemoveTempFolder() {
 }
 
 // Complete 完成下载，返回是否中断监控
-func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
+func (monitor *Monitor) Complete(pool task.Pool) bool {
 	// 创建中转任务
 	file := make([]string, 0, len(monitor.Task.StatusInfo.Files))
 	sizes := make(map[string]uint64, len(monitor.Task.StatusInfo.Files))
@@ -259,11 +259,12 @@ func (monitor *Monitor) Complete(status rpc.StatusInfo) bool {
 	)
 	if err != nil {
 		monitor.setErrorStatus(err)
+		monitor.RemoveTempFolder()
 		return true
 	}
 
 	// 提交中转任务
-	task.TaskPoll.Submit(job)
+	pool.Submit(job)
 
 	// 更新任务ID
 	monitor.Task.TaskID = job.Model().ID

pkg/aria2/monitor/monitor_test.go

@@ -1,326 +1,438 @@
 package monitor
 
 import (
+	"database/sql"
 	"errors"
-	"testing"
-	"time"
-
 	"github.com/DATA-DOG/go-sqlmock"
 	model "github.com/cloudreve/Cloudreve/v3/models"
-	"github.com/cloudreve/Cloudreve/v3/pkg/aria2"
 	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
 	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
-	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
-	"github.com/cloudreve/Cloudreve/v3/pkg/task"
-	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mocks"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/jinzhu/gorm"
 	"github.com/stretchr/testify/assert"
 	testMock "github.com/stretchr/testify/mock"
+	"testing"
 )
 
-type InstanceMock struct {
-	testMock.Mock
-}
+var mock sqlmock.Sqlmock
 
-func (m InstanceMock) CreateTask(task *model.Download, options map[string]interface{}) error {
-	args := m.Called(task, options)
-	return args.Error(0)
-}
-
-func (m InstanceMock) Status(task *model.Download) (rpc.StatusInfo, error) {
-	args := m.Called(task)
-	return args.Get(0).(rpc.StatusInfo), args.Error(1)
-}
-
-func (m InstanceMock) Cancel(task *model.Download) error {
-	args := m.Called(task)
-	return args.Error(0)
-}
-
-func (m InstanceMock) Select(task *model.Download, files []int) error {
-	args := m.Called(task, files)
-	return args.Error(0)
+// TestMain 初始化数据库Mock
+func TestMain(m *testing.M) {
+	var db *sql.DB
+	var err error
+	db, mock, err = sqlmock.New()
+	if err != nil {
+		panic("An error was not expected when opening a stub database connection")
+	}
+	model.DB, _ = gorm.Open("mysql", db)
+	defer db.Close()
+	m.Run()
 }
 
 func TestNewMonitor(t *testing.T) {
-	asserts := assert.New(t)
-	NewMonitor(&model.Download{GID: "gid"})
-	_, ok := common.EventNotifier.Subscribes.Load("gid")
-	asserts.True(ok)
+	a := assert.New(t)
+	mockMQ := mq.NewMQ()
+
+	// node not available
+	{
+		mockPool := &mocks.NodePoolMock{}
+		mockPool.On("GetNodeByID", uint(1)).Return(nil)
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+
+		task := &model.Download{
+			Model: gorm.Model{ID: 1},
+		}
+		NewMonitor(task, mockPool, mockMQ)
+		mockPool.AssertExpectations(t)
+		a.NoError(mock.ExpectationsWereMet())
+		a.NotEmpty(task.Error)
+	}
+
+	// success
+	{
+		mockNode := &mocks.NodeMock{}
+		mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
+		mockPool := &mocks.NodePoolMock{}
+		mockPool.On("GetNodeByID", uint(1)).Return(mockNode)
+
+		task := &model.Download{
+			Model: gorm.Model{ID: 1},
+		}
+		NewMonitor(task, mockPool, mockMQ)
+		mockNode.AssertExpectations(t)
+		mockPool.AssertExpectations(t)
+	}
+
 }
 
 func TestMonitor_Loop(t *testing.T) {
-	asserts := assert.New(t)
-	notifier := make(chan common.StatusEvent)
-	MAX_RETRY = 0
-	monitor := &Monitor{
-		Task:     &model.Download{GID: "gid"},
-		Interval: time.Duration(1) * time.Second,
-		notifier: notifier,
-	}
-	asserts.NotPanics(func() {
-		monitor.Loop()
-	})
-}
-
-func TestMonitor_Update(t *testing.T) {
-	asserts := assert.New(t)
-	monitor := &Monitor{
-		Task: &model.Download{
-			GID:    "gid",
-			Parent: "TestMonitor_Update",
-		},
-		Interval: time.Duration(1) * time.Second,
+	a := assert.New(t)
+	mockMQ := mq.NewMQ()
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
+	m := &Monitor{
+		retried:  MAX_RETRY,
+		node:     mockNode,
+		Task:     &model.Download{Model: gorm.Model{ID: 1}},
+		notifier: mockMQ.Subscribe("test", 1),
 	}
 
-	// 无法获取状态
+	// into interval loop
 	{
-		MAX_RETRY = 1
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, errors.New("error"))
-		file, _ := util.CreatNestedFile("TestMonitor_Update/1")
-		file.Close()
-		aria2.Instance = testInstance
-		asserts.False(monitor.Update())
-		asserts.True(monitor.Update())
-		testInstance.AssertExpectations(t)
-		asserts.False(util.Exists("TestMonitor_Update"))
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+		m.Loop(mockMQ)
+		a.NoError(mock.ExpectationsWereMet())
+		a.NotEmpty(m.Task.Error)
 	}
 
-	// 磁力链下载重定向
+	// into notifier loop
 	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{
-			FollowedBy: []string{"1"},
-		}, nil)
-		monitor.Task.ID = 1
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.False(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-		asserts.EqualValues("1", monitor.Task.GID)
-	}
-
-	// 无法更新任务信息
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{}, nil)
-		monitor.Task.ID = 1
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
-		aria2.mock.ExpectRollback()
-		aria2.Instance = testInstance
-		asserts.True(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-	}
-
-	// 返回未知状态
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "?"}, nil)
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.True(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-	}
-
-	// 返回被取消状态
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "removed"}, nil)
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.True(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-	}
-
-	// 返回活跃状态
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "active"}, nil)
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.False(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-	}
-
-	// 返回错误状态
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "error"}, nil)
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.True(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
-	}
-
-	// 返回完成
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveStatus", testMock.Anything).Return(rpc.StatusInfo{Status: "complete"}, nil)
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		aria2.Instance = testInstance
-		asserts.True(monitor.Update())
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		testInstance.AssertExpectations(t)
+		m.Task.Error = ""
+		mockMQ.Publish("test", mq.Message{})
+		mock.ExpectBegin()
+		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+		mock.ExpectCommit()
+		m.Loop(mockMQ)
+		a.NoError(mock.ExpectationsWereMet())
+		a.NotEmpty(m.Task.Error)
 	}
 }
 
-func TestMonitor_UpdateTaskInfo(t *testing.T) {
-	asserts := assert.New(t)
-	monitor := &Monitor{
-		Task: &model.Download{
-			Model:  gorm.Model{ID: 1},
-			GID:    "gid",
-			Parent: "TestMonitor_UpdateTaskInfo",
-		},
+func TestMonitor_UpdateFailedAfterRetry(t *testing.T) {
+	a := assert.New(t)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	for i := 0; i < MAX_RETRY; i++ {
+		a.False(m.Update())
 	}
 
-	// 失败
-	{
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
-		aria2.mock.ExpectRollback()
-		err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		asserts.Error(err)
+	mockNode.AssertExpectations(t)
+	a.True(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	a.NotEmpty(m.Task.Error)
+}
+
+func TestMonitor_UpdateMagentoFollow(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		FollowedBy: []string{"next"},
+	}, nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.False(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	a.Equal("next", m.Task.GID)
+	mockAria2.AssertExpectations(t)
+}
+
+func TestMonitor_UpdateFailedToUpdateInfo(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{}, nil)
+	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnError(errors.New("error"))
+	mock.ExpectRollback()
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+	a.NotEmpty(m.Task.Error)
+}
+
+func TestMonitor_UpdateCompleted(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		Status: "complete",
+	}, nil)
+	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	mockNode.On("ID").Return(uint(1))
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+	mock.ExpectQuery("SELECT(.+)users(.+)").WillReturnError(errors.New("error"))
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+	a.NotEmpty(m.Task.Error)
+}
+
+func TestMonitor_UpdateError(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		Status:       "error",
+		ErrorMessage: "error",
+	}, nil)
+	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+	a.NotEmpty(m.Task.Error)
+}
+
+func TestMonitor_UpdateActive(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		Status: "active",
+	}, nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.False(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+}
+
+func TestMonitor_UpdateRemoved(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		Status: "removed",
+	}, nil)
+	mockAria2.On("DeleteTempFile", testMock.Anything).Return(nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Update())
+	a.Equal(common.Canceled, m.Task.Status)
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+}
+
+func TestMonitor_UpdateUnknown(t *testing.T) {
+	a := assert.New(t)
+	mockAria2 := &mocks.Aria2Mock{}
+	mockAria2.On("Status", testMock.Anything).Return(rpc.StatusInfo{
+		Status: "unknown",
+	}, nil)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(mockAria2)
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
+	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Update())
+	a.NoError(mock.ExpectationsWereMet())
+	mockAria2.AssertExpectations(t)
+	mockNode.AssertExpectations(t)
+}
+
+func TestMonitor_UpdateTaskInfoValidateFailed(t *testing.T) {
+	a := assert.New(t)
+	status := rpc.StatusInfo{
+		Status:          "completed",
+		TotalLength:     "100",
+		CompletedLength: "50",
+		DownloadSpeed:   "20",
+	}
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
+	m := &Monitor{
+		node: mockNode,
+		Task: &model.Download{Model: gorm.Model{ID: 1}},
 	}
 
-	// 更新成功，无需校验
-	{
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		err := monitor.UpdateTaskInfo(rpc.StatusInfo{})
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		asserts.NoError(err)
-	}
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
 
-	// 更新成功，大小改变，需要校验，校验失败
-	{
-		testInstance := new(InstanceMock)
-		testInstance.On("SlaveCancel", testMock.Anything).Return(nil)
-		aria2.Instance = testInstance
-		aria2.mock.ExpectBegin()
-		aria2.mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
-		aria2.mock.ExpectCommit()
-		err := monitor.UpdateTaskInfo(rpc.StatusInfo{TotalLength: "1"})
-		asserts.NoError(aria2.mock.ExpectationsWereMet())
-		asserts.Error(err)
-		testInstance.AssertExpectations(t)
-	}
+	err := m.UpdateTaskInfo(status)
+	a.Error(err)
+	a.NoError(mock.ExpectationsWereMet())
+	mockNode.AssertExpectations(t)
 }
 
 func TestMonitor_ValidateFile(t *testing.T) {
-	asserts := assert.New(t)
-	monitor := &Monitor{
+	a := assert.New(t)
+	m := &Monitor{
 		Task: &model.Download{
-			Model:  gorm.Model{ID: 1},
-			GID:    "gid",
-			Parent: "TestMonitor_ValidateFile",
+			Model:     gorm.Model{ID: 1},
+			TotalSize: 100,
 		},
 	}
 
-	// 无法创建文件系统
+	// failed to create filesystem
 	{
-		monitor.Task.User = &model.User{
+		m.Task.User = &model.User{
 			Policy: model.Policy{
-				Type: "unknown",
+				Type: "random",
 			},
 		}
-		asserts.Error(monitor.ValidateFile())
+		a.Equal(filesystem.ErrUnknownPolicyType, m.ValidateFile())
 	}
 
-	// 文件大小超出容量配额
+	// User capacity not enough
 	{
-		cache.Set("pack_size_0", uint64(0), 0)
-		monitor.Task.TotalSize = 11
-		monitor.Task.User = &model.User{
-			Policy: model.Policy{
-				Type: "mock",
-			},
+		m.Task.User = &model.User{
 			Group: model.Group{
-				MaxStorage: 10,
+				MaxStorage: 99,
+			},
+			Policy: model.Policy{
+				Type: "local",
 			},
 		}
-		asserts.Equal(filesystem.ErrInsufficientCapacity, monitor.ValidateFile())
+		a.Equal(filesystem.ErrInsufficientCapacity, m.ValidateFile())
 	}
 
-	// 单文件大小超出容量配额
+	// single file too big
 	{
-		cache.Set("pack_size_0", uint64(0), 0)
-		monitor.Task.TotalSize = 10
-		monitor.Task.StatusInfo.Files = []rpc.FileInfo{
+		m.Task.StatusInfo.Files = []rpc.FileInfo{
 			{
+				Length:   "100",
 				Selected: "true",
-				Length:   "6",
 			},
 		}
-		monitor.Task.User = &model.User{
-			Policy: model.Policy{
-				Type:    "mock",
-				MaxSize: 5,
-			},
+		m.Task.User = &model.User{
 			Group: model.Group{
-				MaxStorage: 10,
+				MaxStorage: 100,
+			},
+			Policy: model.Policy{
+				Type:    "local",
+				MaxSize: 99,
 			},
 		}
-		asserts.Equal(filesystem.ErrFileSizeTooBig, monitor.ValidateFile())
+		a.Equal(filesystem.ErrFileSizeTooBig, m.ValidateFile())
+	}
+
+	// all pass
+	{
+		m.Task.StatusInfo.Files = []rpc.FileInfo{
+			{
+				Length:   "100",
+				Selected: "true",
+			},
+		}
+		m.Task.User = &model.User{
+			Group: model.Group{
+				MaxStorage: 100,
+			},
+			Policy: model.Policy{
+				Type:    "local",
+				MaxSize: 100,
+			},
+		}
+		a.NoError(m.ValidateFile())
 	}
 }
 
 func TestMonitor_Complete(t *testing.T) {
-	asserts := assert.New(t)
-	monitor := &Monitor{
+	a := assert.New(t)
+	mockNode := &mocks.NodeMock{}
+	mockNode.On("ID").Return(uint(1))
+	mockPool := &mocks.TaskPoolMock{}
+	mockPool.On("Submit", testMock.Anything)
+	m := &Monitor{
+		node: mockNode,
 		Task: &model.Download{
-			Model:  gorm.Model{ID: 1},
-			GID:    "gid",
-			Parent: "TestMonitor_Complete",
-			StatusInfo: rpc.StatusInfo{
-				Files: []rpc.FileInfo{
-					{
-						Selected: "true",
-						Path:     "TestMonitor_Complete",
-					},
-				},
-			},
+			Model:     gorm.Model{ID: 1},
+			TotalSize: 100,
+			UserID:    9414,
+		},
+	}
+	m.Task.StatusInfo.Files = []rpc.FileInfo{
+		{
+			Length:   "100",
+			Selected: "true",
 		},
 	}
 
-	cache.Set("setting_max_worker_num", "1", 0)
-	aria2.mock.ExpectQuery("SELECT(.+)tasks").WillReturnRows(sqlmock.NewRows([]string{"id"}))
-	task.Init()
-	aria2.mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
-	aria2.mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
-	aria2.mock.ExpectBegin()
-	aria2.mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
-	aria2.mock.ExpectCommit()
+	mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(9414))
 
-	aria2.mock.ExpectBegin()
-	aria2.mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
-	aria2.mock.ExpectCommit()
-	asserts.True(monitor.Complete(rpc.StatusInfo{}))
-	asserts.NoError(aria2.mock.ExpectationsWereMet())
+	mock.ExpectBegin()
+	mock.ExpectExec("INSERT(.+)tasks").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	mock.ExpectBegin()
+	mock.ExpectExec("UPDATE(.+)downloads").WillReturnResult(sqlmock.NewResult(1, 1))
+	mock.ExpectCommit()
+
+	a.True(m.Complete(mockPool))
+	a.NoError(mock.ExpectationsWereMet())
+	mockNode.AssertExpectations(t)
+	mockPool.AssertExpectations(t)
 }

pkg/aria2/notification_test.go

@@ -1,52 +0,0 @@
-package aria2
-
-import (
-	"testing"
-
-	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestNotifier_Notify(t *testing.T) {
-	asserts := assert.New(t)
-	notifier2 := &Notifier{}
-	notifyChan := make(chan StatusEvent, 10)
-	notifier2.Subscribe(notifyChan, "1")
-
-	// 未订阅
-	{
-		notifier2.Notify([]rpc.Event{rpc.Event{Gid: ""}}, 1)
-		asserts.Len(notifyChan, 0)
-	}
-
-	// 订阅
-	{
-		notifier2.Notify([]rpc.Event{{Gid: "1"}}, 1)
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnBtDownloadComplete([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnDownloadStart([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnDownloadPause([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnDownloadStop([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnDownloadComplete([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-
-		notifier2.OnDownloadError([]rpc.Event{{Gid: "1"}})
-		asserts.Len(notifyChan, 1)
-		<-notifyChan
-	}
-}

pkg/auth/auth.go

@@ -17,8 +17,10 @@ import (
 )
 
 var (
-	ErrAuthFailed = serializer.NewError(serializer.CodeNoPermissionErr, "鉴权失败", nil)
-	ErrExpired    = serializer.NewError(serializer.CodeSignExpired, "签名已过期", nil)
+	ErrAuthFailed        = serializer.NewError(serializer.CodeNoPermissionErr, "鉴权失败", nil)
+	ErrAuthHeaderMissing = serializer.NewError(serializer.CodeNoPermissionErr, "authorization header is missing", nil)
+	ErrExpiresMissing    = serializer.NewError(serializer.CodeNoPermissionErr, "expire timestamp is missing", nil)
+	ErrExpired           = serializer.NewError(serializer.CodeSignExpired, "签名已过期", nil)
 )
 
 // General 通用的认证接口
@@ -33,7 +35,7 @@ type Auth interface {
 }
 
 // SignRequest 对PUT\POST等复杂HTTP请求签名，只会对URI部分、
-// 请求正文、`X-`开头的header进行签名
+// 请求正文、`X-Cr-`开头的header进行签名
 func SignRequest(instance Auth, r *http.Request, expires int64) *http.Request {
 	// 处理有效期
 	if expires > 0 {
@@ -55,7 +57,7 @@ func CheckRequest(instance Auth, r *http.Request) error {
 		ok   bool
 	)
 	if sign, ok = r.Header["Authorization"]; !ok || len(sign) == 0 {
-		return ErrAuthFailed
+		return ErrAuthHeaderMissing
 	}
 	sign[0] = strings.TrimPrefix(sign[0], "Bearer ")
 
@@ -67,7 +69,7 @@ func CheckRequest(instance Auth, r *http.Request) error {
 func getSignContent(r *http.Request) (rawSignString string) {
 	// 读取所有body正文
 	var body = []byte{}
-	if _, ok := r.Header["X-Policy"]; !ok {
+	if _, ok := r.Header["X-Cr-Policy"]; !ok {
 		if r.Body != nil {
 			body, _ = ioutil.ReadAll(r.Body)
 			_ = r.Body.Close()
@@ -78,7 +80,7 @@ func getSignContent(r *http.Request) (rawSignString string) {
 	// 决定要签名的header
 	var signedHeader []string
 	for k, _ := range r.Header {
-		if strings.HasPrefix(k, "X-") && k != "X-Filename" {
+		if strings.HasPrefix(k, "X-Cr-") && k != "X-Cr-Filename" {
 			signedHeader = append(signedHeader, fmt.Sprintf("%s=%s", k, r.Header.Get(k)))
 		}
 	}

pkg/auth/auth_test.go

@@ -70,7 +70,7 @@ func TestSignRequest(t *testing.T) {
 			strings.NewReader("I am body."),
 		)
 		asserts.NoError(err)
-		req.Header["X-Policy"] = []string{"I am Policy"}
+		req.Header["X-Cr-Policy"] = []string{"I am Policy"}
 		req = SignRequest(General, req, 10)
 		asserts.NotEmpty(req.Header["Authorization"])
 	}
@@ -80,6 +80,19 @@ func TestCheckRequest(t *testing.T) {
 	asserts := assert.New(t)
 	General = HMACAuth{SecretKey: []byte(util.RandStringRunes(256))}
 
+	// 缺少请求头
+	{
+		req, err := http.NewRequest(
+			"POST",
+			"http://127.0.0.1/api/v3/upload",
+			strings.NewReader("I am body."),
+		)
+		asserts.NoError(err)
+		err = CheckRequest(General, req)
+		asserts.Error(err)
+		asserts.Equal(ErrAuthHeaderMissing, err)
+	}
+
 	// 非上传请求 验证成功
 	{
 		req, err := http.NewRequest(
@@ -101,7 +114,7 @@ func TestCheckRequest(t *testing.T) {
 			strings.NewReader("I am body."),
 		)
 		asserts.NoError(err)
-		req.Header["X-Policy"] = []string{"I am Policy"}
+		req.Header["X-Cr-Policy"] = []string{"I am Policy"}
 		req = SignRequest(General, req, 0)
 		err = CheckRequest(General, req)
 		asserts.NoError(err)

pkg/auth/hmac.go

@@ -33,7 +33,7 @@ func (auth HMACAuth) Check(body string, sign string) error {
 	signSlice := strings.Split(sign, ":")
 	// 如果未携带expires字段
 	if signSlice[len(signSlice)-1] == "" {
-		return ErrAuthFailed
+		return ErrExpiresMissing
 	}
 
 	// 验证是否过期

pkg/balancer/balancer_test.go

@@ -0,0 +1,12 @@
+package balancer
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewBalancer(t *testing.T) {
+	a := assert.New(t)
+	a.NotNil(NewBalancer(""))
+	a.IsType(&RoundRobin{}, NewBalancer("RoundRobin"))
+}

pkg/balancer/roundrobin_test.go

@@ -0,0 +1,42 @@
+package balancer
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestRoundRobin_NextIndex(t *testing.T) {
+	a := assert.New(t)
+	r := &RoundRobin{}
+	total := 5
+	for i := 1; i < total; i++ {
+		a.Equal(i, r.NextIndex(total))
+	}
+	for i := 0; i < total; i++ {
+		a.Equal(i, r.NextIndex(total))
+	}
+}
+
+func TestRoundRobin_NextPeer(t *testing.T) {
+	a := assert.New(t)
+	r := &RoundRobin{}
+
+	// not slice
+	{
+		err, _ := r.NextPeer("s")
+		a.Equal(ErrInputNotSlice, err)
+	}
+
+	// no nodes
+	{
+		err, _ := r.NextPeer([]string{})
+		a.Equal(ErrNoAvaliableNode, err)
+	}
+
+	// pass
+	{
+		err, res := r.NextPeer([]string{"a"})
+		a.NoError(err)
+		a.Equal("a", res.(string))
+	}
+}

pkg/cluster/controller.go

@@ -1,4 +1,4 @@
-package slave
+package cluster
 
 import (
 	"bytes"
@@ -8,7 +8,6 @@ import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
 	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
 	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
-	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/cloudreve/Cloudreve/v3/pkg/request"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
@@ -51,13 +50,13 @@ type MasterInfo struct {
 	TTL int
 	URL *url.URL
 	// used to invoke aria2 rpc calls
-	Instance cluster.Node
+	Instance Node
 	Client   request.Client
 
 	jobTracker map[string]bool
 }
 
-func Init() {
+func InitController() {
 	DefaultController = &slaveController{
 		masters: make(map[string]MasterInfo),
 	}
@@ -95,7 +94,7 @@ func (c *slaveController) HandleHeartBeat(req *serializer.NodePingReq) (serializ
 				}, int64(req.CredentialTTL)),
 			),
 			jobTracker: make(map[string]bool),
-			Instance: cluster.NewNodeFromDBModel(&model.Node{
+			Instance: NewNodeFromDBModel(&model.Node{
 				Model:                  gorm.Model{ID: req.Node.ID},
 				MasterKey:              req.Node.MasterKey,
 				Type:                   model.MasterNodeType,
@@ -162,6 +161,7 @@ func (c *slaveController) SubmitTask(id string, job interface{}, hash string, su
 			return nil
 		}
 
+		node.jobTracker[hash] = true
 		submitter(job)
 		return nil
 	}

pkg/cluster/controller_test.go

@@ -0,0 +1,385 @@
+package cluster
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
+	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
+	"github.com/cloudreve/Cloudreve/v3/pkg/request"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+	"github.com/stretchr/testify/assert"
+	testMock "github.com/stretchr/testify/mock"
+	"io"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func TestInitController(t *testing.T) {
+	assert.NotPanics(t, func() {
+		InitController()
+	})
+}
+
+func TestSlaveController_HandleHeartBeat(t *testing.T) {
+	a := assert.New(t)
+	c := &slaveController{
+		masters: make(map[string]MasterInfo),
+	}
+
+	// first heart beat
+	{
+		_, err := c.HandleHeartBeat(&serializer.NodePingReq{
+			SiteID: "1",
+			Node:   &model.Node{},
+		})
+		a.NoError(err)
+
+		_, err = c.HandleHeartBeat(&serializer.NodePingReq{
+			SiteID: "2",
+			Node:   &model.Node{},
+		})
+		a.NoError(err)
+
+		a.Len(c.masters, 2)
+	}
+
+	// second heart beat, no fresh
+	{
+		_, err := c.HandleHeartBeat(&serializer.NodePingReq{
+			SiteID:  "1",
+			SiteURL: "http://127.0.0.1",
+			Node:    &model.Node{},
+		})
+		a.NoError(err)
+		a.Len(c.masters, 2)
+		a.Empty(c.masters["1"].URL)
+	}
+
+	// second heart beat, fresh
+	{
+		_, err := c.HandleHeartBeat(&serializer.NodePingReq{
+			SiteID:   "1",
+			IsUpdate: true,
+			SiteURL:  "http://127.0.0.1",
+			Node:     &model.Node{},
+		})
+		a.NoError(err)
+		a.Len(c.masters, 2)
+		a.Equal("http://127.0.0.1", c.masters["1"].URL.String())
+	}
+
+	// second heart beat, fresh, url illegal
+	{
+		_, err := c.HandleHeartBeat(&serializer.NodePingReq{
+			SiteID:   "1",
+			IsUpdate: true,
+			SiteURL:  string([]byte{0x7f}),
+			Node:     &model.Node{},
+		})
+		a.Error(err)
+		a.Len(c.masters, 2)
+		a.Equal("http://127.0.0.1", c.masters["1"].URL.String())
+	}
+}
+
+type nodeMock struct {
+	testMock.Mock
+}
+
+func (n nodeMock) Init(node *model.Node) {
+	n.Called(node)
+}
+
+func (n nodeMock) IsFeatureEnabled(feature string) bool {
+	args := n.Called(feature)
+	return args.Bool(0)
+}
+
+func (n nodeMock) SubscribeStatusChange(callback func(isActive bool, id uint)) {
+	n.Called(callback)
+}
+
+func (n nodeMock) Ping(req *serializer.NodePingReq) (*serializer.NodePingResp, error) {
+	args := n.Called(req)
+	return args.Get(0).(*serializer.NodePingResp), args.Error(1)
+}
+
+func (n nodeMock) IsActive() bool {
+	args := n.Called()
+	return args.Bool(0)
+}
+
+func (n nodeMock) GetAria2Instance() common.Aria2 {
+	args := n.Called()
+	return args.Get(0).(common.Aria2)
+}
+
+func (n nodeMock) ID() uint {
+	args := n.Called()
+	return args.Get(0).(uint)
+}
+
+func (n nodeMock) Kill() {
+	n.Called()
+}
+
+func (n nodeMock) IsMater() bool {
+	args := n.Called()
+	return args.Bool(0)
+}
+
+func (n nodeMock) MasterAuthInstance() auth.Auth {
+	args := n.Called()
+	return args.Get(0).(auth.Auth)
+}
+
+func (n nodeMock) SlaveAuthInstance() auth.Auth {
+	args := n.Called()
+	return args.Get(0).(auth.Auth)
+}
+
+func (n nodeMock) DBModel() *model.Node {
+	args := n.Called()
+	return args.Get(0).(*model.Node)
+}
+
+func TestSlaveController_GetAria2Instance(t *testing.T) {
+	a := assert.New(t)
+	mockNode := &nodeMock{}
+	mockNode.On("GetAria2Instance").Return(&common.DummyAria2{})
+	c := &slaveController{
+		masters: map[string]MasterInfo{
+			"1": {Instance: mockNode},
+		},
+	}
+
+	// node node found
+	{
+		res, err := c.GetAria2Instance("2")
+		a.Nil(res)
+		a.Equal(ErrMasterNotFound, err)
+	}
+
+	// node found
+	{
+		res, err := c.GetAria2Instance("1")
+		a.NotNil(res)
+		a.NoError(err)
+		mockNode.AssertExpectations(t)
+	}
+
+}
+
+type requestMock struct {
+	testMock.Mock
+}
+
+func (r requestMock) Request(method, target string, body io.Reader, opts ...request.Option) *request.Response {
+	return r.Called(method, target, body, opts).Get(0).(*request.Response)
+}
+
+func TestSlaveController_SendNotification(t *testing.T) {
+	a := assert.New(t)
+	c := &slaveController{
+		masters: map[string]MasterInfo{
+			"1": {},
+		},
+	}
+
+	// node not exit
+	{
+		a.Equal(ErrMasterNotFound, c.SendNotification("2", "", mq.Message{}))
+	}
+
+	// gob encode error
+	{
+		type randomType struct{}
+		a.Error(c.SendNotification("1", "", mq.Message{
+			Content: randomType{},
+		}))
+	}
+
+	// return none 200
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s1", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{StatusCode: http.StatusConflict},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		a.Error(c.SendNotification("1", "s1", mq.Message{}))
+		mockRequest.AssertExpectations(t)
+	}
+
+	// master return error
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s2", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		a.Equal(1, c.SendNotification("1", "s2", mq.Message{}).(serializer.AppError).Code)
+		mockRequest.AssertExpectations(t)
+	}
+
+	// success
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "PUT", "/api/v3/slave/notification/s3", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":0}")),
+			},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		a.NoError(c.SendNotification("1", "s3", mq.Message{}))
+		mockRequest.AssertExpectations(t)
+	}
+}
+
+func TestSlaveController_SubmitTask(t *testing.T) {
+	a := assert.New(t)
+	c := &slaveController{
+		masters: map[string]MasterInfo{
+			"1": {
+				jobTracker: map[string]bool{},
+			},
+		},
+	}
+
+	// node not exit
+	{
+		a.Equal(ErrMasterNotFound, c.SubmitTask("2", "", "", nil))
+	}
+
+	// success
+	{
+		submitted := false
+		a.NoError(c.SubmitTask("1", "", "hash", func(i interface{}) {
+			submitted = true
+		}))
+		a.True(submitted)
+	}
+
+	// job already submitted
+	{
+		submitted := false
+		a.NoError(c.SubmitTask("1", "", "hash", func(i interface{}) {
+			submitted = true
+		}))
+		a.False(submitted)
+	}
+}
+
+func TestSlaveController_GetMasterInfo(t *testing.T) {
+	a := assert.New(t)
+	c := &slaveController{
+		masters: map[string]MasterInfo{
+			"1": {},
+		},
+	}
+
+	// node not exit
+	{
+		res, err := c.GetMasterInfo("2")
+		a.Equal(ErrMasterNotFound, err)
+		a.Nil(res)
+	}
+
+	// success
+	{
+		res, err := c.GetMasterInfo("1")
+		a.NoError(err)
+		a.NotNil(res)
+	}
+}
+
+func TestSlaveController_GetOneDriveToken(t *testing.T) {
+	a := assert.New(t)
+	c := &slaveController{
+		masters: map[string]MasterInfo{
+			"1": {},
+		},
+	}
+
+	// node not exit
+	{
+		res, err := c.GetOneDriveToken("2", 1)
+		a.Equal(ErrMasterNotFound, err)
+		a.Empty(res)
+	}
+
+	// return none 200
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{StatusCode: http.StatusConflict},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		res, err := c.GetOneDriveToken("1", 1)
+		a.Error(err)
+		a.Empty(res)
+		mockRequest.AssertExpectations(t)
+	}
+
+	// master return error
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		res, err := c.GetOneDriveToken("1", 1)
+		a.Equal(1, err.(serializer.AppError).Code)
+		a.Empty(res)
+		mockRequest.AssertExpectations(t)
+	}
+
+	// success
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "GET", "/api/v3/slave/credential/onedrive/1", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"expected\"}")),
+			},
+		})
+		c := &slaveController{
+			masters: map[string]MasterInfo{
+				"1": {Client: mockRequest},
+			},
+		}
+		res, err := c.GetOneDriveToken("1", 1)
+		a.NoError(err)
+		a.Equal("expected", res)
+		mockRequest.AssertExpectations(t)
+	}
+
+}

pkg/cluster/errors.go

@@ -1,8 +1,12 @@
 package cluster
 
-import "errors"
+import (
+	"errors"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+)
 
 var (
 	ErrFeatureNotExist = errors.New("No nodes in nodepool match the feature specificed")
 	ErrIlegalPath      = errors.New("path out of boundary of setting temp folder")
+	ErrMasterNotFound  = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
 )

pkg/cluster/master.go

@@ -20,7 +20,10 @@ import (
 	"time"
 )
 
-const deleteTempFileDuration = 60 * time.Second
+const (
+	deleteTempFileDuration = 60 * time.Second
+	statusRetryDuration    = 10 * time.Second
+)
 
 type MasterNode struct {
 	Model    *model.Node
@@ -33,8 +36,10 @@ type rpcService struct {
 	Caller      rpc.Client
 	Initialized bool
 
-	parent  *MasterNode
-	options *clientOptions
+	retryDuration         time.Duration
+	deletePaddingDuration time.Duration
+	parent                *MasterNode
+	options               *clientOptions
 }
 
 type clientOptions struct {
@@ -46,6 +51,8 @@ func (node *MasterNode) Init(nodeModel *model.Node) {
 	node.lock.Lock()
 	node.Model = nodeModel
 	node.aria2RPC.parent = node
+	node.aria2RPC.retryDuration = statusRetryDuration
+	node.aria2RPC.deletePaddingDuration = deleteTempFileDuration
 	node.lock.Unlock()
 
 	node.lock.RLock()
@@ -214,8 +221,8 @@ func (r *rpcService) Status(task *model.Download) (rpc.StatusInfo, error) {
 	res, err := r.Caller.TellStatus(task.GID)
 	if err != nil {
 		// 失败后重试
-		util.Log().Debug("无法获取离线下载状态，%s，10秒钟后重试", err)
-		time.Sleep(time.Duration(10) * time.Second)
+		util.Log().Debug("无法获取离线下载状态，%s，稍后重试", err)
+		time.Sleep(r.retryDuration)
 		res, err = r.Caller.TellStatus(task.GID)
 	}
 
@@ -253,13 +260,13 @@ func (s *rpcService) DeleteTempFile(task *model.Download) error {
 	defer s.parent.lock.RUnlock()
 
 	// 避免被aria2占用，异步执行删除
-	go func(src string) {
-		time.Sleep(deleteTempFileDuration)
+	go func(d time.Duration, src string) {
+		time.Sleep(d)
 		err := os.RemoveAll(src)
 		if err != nil {
 			util.Log().Warning("无法删除离线下载临时目录[%s], %s", src, err)
 		}
-	}(task.Parent)
+	}(s.deletePaddingDuration, task.Parent)
 
 	return nil
 }

pkg/cluster/master_test.go

@@ -0,0 +1,186 @@
+package cluster
+
+import (
+	"context"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
+	"github.com/stretchr/testify/assert"
+	"os"
+	"testing"
+	"time"
+)
+
+func TestMasterNode_Init(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{}
+	m.Init(&model.Node{Status: model.NodeSuspend})
+	a.Equal(model.NodeSuspend, m.DBModel().Status)
+	m.Init(&model.Node{Aria2Enabled: true})
+}
+
+func TestMasterNode_DummyMethods(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{
+		Model: &model.Node{},
+	}
+
+	m.Model.ID = 5
+	a.Equal(m.Model.ID, m.ID())
+
+	res, err := m.Ping(&serializer.NodePingReq{})
+	a.NoError(err)
+	a.NotNil(res)
+
+	a.True(m.IsActive())
+	a.True(m.IsMater())
+
+	m.SubscribeStatusChange(func(isActive bool, id uint) {})
+}
+
+func TestMasterNode_IsFeatureEnabled(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{
+		Model: &model.Node{},
+	}
+
+	a.False(m.IsFeatureEnabled("aria2"))
+	a.False(m.IsFeatureEnabled("random"))
+	m.Model.Aria2Enabled = true
+	a.True(m.IsFeatureEnabled("aria2"))
+}
+
+func TestMasterNode_AuthInstance(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{
+		Model: &model.Node{},
+	}
+
+	a.NotNil(m.MasterAuthInstance())
+	a.NotNil(m.SlaveAuthInstance())
+}
+
+func TestMasterNode_Kill(t *testing.T) {
+	m := &MasterNode{
+		Model: &model.Node{},
+	}
+
+	m.Kill()
+
+	caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
+	m.aria2RPC.Caller = caller
+	m.Kill()
+}
+
+func TestMasterNode_GetAria2Instance(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{
+		Model:    &model.Node{},
+		aria2RPC: rpcService{},
+	}
+
+	m.aria2RPC.parent = m
+
+	a.NotNil(m.GetAria2Instance())
+	m.Model.Aria2Enabled = true
+	a.NotNil(m.GetAria2Instance())
+	m.aria2RPC.Initialized = true
+	a.NotNil(m.GetAria2Instance())
+}
+
+func TestRpcService_Init(t *testing.T) {
+	a := assert.New(t)
+	m := &MasterNode{
+		Model: &model.Node{
+			Aria2OptionsSerialized: model.Aria2Option{
+				Options: "{",
+			},
+		},
+		aria2RPC: rpcService{},
+	}
+	m.aria2RPC.parent = m
+
+	// failed to decode address
+	{
+		m.Model.Aria2OptionsSerialized.Server = string([]byte{0x7f})
+		a.Error(m.aria2RPC.Init())
+	}
+
+	// failed to decode options
+	{
+		m.Model.Aria2OptionsSerialized.Server = ""
+		a.Error(m.aria2RPC.Init())
+	}
+
+	// failed to initialized
+	{
+		m.Model.Aria2OptionsSerialized.Server = ""
+		m.Model.Aria2OptionsSerialized.Options = "{}"
+		caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
+		m.aria2RPC.Caller = caller
+		a.Error(m.aria2RPC.Init())
+		a.False(m.aria2RPC.Initialized)
+	}
+}
+
+func getTestRPCNode() *MasterNode {
+	m := &MasterNode{
+		Model: &model.Node{
+			Aria2OptionsSerialized: model.Aria2Option{},
+		},
+		aria2RPC: rpcService{
+			options: &clientOptions{
+				Options: map[string]interface{}{"1": "1"},
+			},
+		},
+	}
+	m.aria2RPC.parent = m
+	caller, _ := rpc.New(context.Background(), "http://", "", 0, nil)
+	m.aria2RPC.Caller = caller
+	return m
+}
+
+func TestRpcService_CreateTask(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNode()
+
+	res, err := m.aria2RPC.CreateTask(&model.Download{}, map[string]interface{}{"1": "1"})
+	a.Error(err)
+	a.Empty(res)
+}
+
+func TestRpcService_Status(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNode()
+
+	res, err := m.aria2RPC.Status(&model.Download{})
+	a.Error(err)
+	a.Empty(res)
+}
+
+func TestRpcService_Cancel(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNode()
+
+	a.Error(m.aria2RPC.Cancel(&model.Download{}))
+}
+
+func TestRpcService_Select(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNode()
+
+	a.NotNil(m.aria2RPC.GetConfig())
+	a.Error(m.aria2RPC.Select(&model.Download{}, []int{1, 2, 3}))
+}
+
+func TestRpcService_DeleteTempFile(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNode()
+	fdName := "TestRpcService_DeleteTempFile"
+	a.NoError(os.Mkdir(fdName, 0644))
+
+	a.NoError(m.aria2RPC.DeleteTempFile(&model.Download{Parent: fdName}))
+	time.Sleep(500 * time.Millisecond)
+	a.False(util.Exists(fdName))
+}

pkg/cluster/node_test.go

@@ -0,0 +1,17 @@
+package cluster
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewNodeFromDBModel(t *testing.T) {
+	a := assert.New(t)
+	a.IsType(&SlaveNode{}, NewNodeFromDBModel(&model.Node{
+		Type: model.SlaveNodeType,
+	}))
+	a.IsType(&MasterNode{}, NewNodeFromDBModel(&model.Node{
+		Type: model.MasterNodeType,
+	}))
+}

pkg/cluster/pool.go

@@ -39,14 +39,22 @@ type NodePool struct {
 
 // Init 初始化从机节点池
 func Init() {
-	Default = &NodePool{
-		featureMap: make(map[string][]Node),
-	}
+	Default = &NodePool{}
+	Default.Init()
 	if err := Default.initFromDB(); err != nil {
 		util.Log().Warning("节点池初始化失败, %s", err)
 	}
 }
 
+func (pool *NodePool) Init() {
+	pool.lock.Lock()
+	defer pool.lock.Unlock()
+
+	pool.featureMap = make(map[string][]Node)
+	pool.active = make(map[uint]Node)
+	pool.inactive = make(map[uint]Node)
+}
+
 func (pool *NodePool) buildIndexMap() {
 	pool.lock.Lock()
 	for _, feature := range featureGroup {
@@ -76,14 +84,19 @@ func (pool *NodePool) GetNodeByID(id uint) Node {
 
 func (pool *NodePool) nodeStatusChange(isActive bool, id uint) {
 	util.Log().Debug("从机节点 [ID=%d] 状态变更 [Active=%t]", id, isActive)
+	var node Node
 	pool.lock.Lock()
-	if isActive {
-		node := pool.inactive[id]
+	if n, ok := pool.inactive[id]; ok {
+		node = n
 		delete(pool.inactive, id)
+	} else {
+		node = pool.active[id]
+		delete(pool.active, id)
+	}
+
+	if isActive {
 		pool.active[id] = node
 	} else {
-		node := pool.active[id]
-		delete(pool.active, id)
 		pool.inactive[id] = node
 	}
 	pool.lock.Unlock()
@@ -98,8 +111,6 @@ func (pool *NodePool) initFromDB() error {
 	}
 
 	pool.lock.Lock()
-	pool.active = make(map[uint]Node)
-	pool.inactive = make(map[uint]Node)
 	for i := 0; i < len(nodes); i++ {
 		pool.add(&nodes[i])
 	}

pkg/cluster/pool_test.go

@@ -0,0 +1,161 @@
+package cluster
+
+import (
+	"database/sql"
+	"errors"
+	"github.com/DATA-DOG/go-sqlmock"
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
+	"github.com/jinzhu/gorm"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var mock sqlmock.Sqlmock
+
+// TestMain 初始化数据库Mock
+func TestMain(m *testing.M) {
+	var db *sql.DB
+	var err error
+	db, mock, err = sqlmock.New()
+	if err != nil {
+		panic("An error was not expected when opening a stub database connection")
+	}
+	model.DB, _ = gorm.Open("mysql", db)
+	defer db.Close()
+	m.Run()
+}
+
+func TestInitFailed(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)").WillReturnError(errors.New("error"))
+	Init()
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestInitSuccess(t *testing.T) {
+	a := assert.New(t)
+	mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "aria2_enabled", "type"}).AddRow(1, true, model.MasterNodeType))
+	Init()
+	a.NoError(mock.ExpectationsWereMet())
+}
+
+func TestNodePool_GetNodeByID(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+	mockNode := &nodeMock{}
+
+	// inactive
+	{
+		p.inactive[1] = mockNode
+		a.Equal(mockNode, p.GetNodeByID(1))
+	}
+
+	// active
+	{
+		delete(p.inactive, 1)
+		p.active[1] = mockNode
+		a.Equal(mockNode, p.GetNodeByID(1))
+	}
+}
+
+func TestNodePool_NodeStatusChange(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	n := &MasterNode{Model: &model.Node{}}
+	p.Init()
+	p.inactive[1] = n
+
+	p.nodeStatusChange(true, 1)
+	a.Len(p.inactive, 0)
+	a.Equal(n, p.active[1])
+
+	p.nodeStatusChange(false, 1)
+	a.Len(p.active, 0)
+	a.Equal(n, p.inactive[1])
+
+	p.nodeStatusChange(false, 1)
+	a.Len(p.active, 0)
+	a.Equal(n, p.inactive[1])
+}
+
+func TestNodePool_Add(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// new node
+	{
+		p.Add(&model.Node{})
+		a.Len(p.active, 1)
+	}
+
+	// old node
+	{
+		p.inactive[0] = p.active[0]
+		delete(p.active, 0)
+		p.Add(&model.Node{})
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 1)
+	}
+}
+
+func TestNodePool_Delete(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// active
+	{
+		mockNode := &nodeMock{}
+		mockNode.On("Kill")
+		p.active[0] = mockNode
+		p.Delete(0)
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 0)
+		mockNode.AssertExpectations(t)
+	}
+
+	p.Init()
+
+	// inactive
+	{
+		mockNode := &nodeMock{}
+		mockNode.On("Kill")
+		p.inactive[0] = mockNode
+		p.Delete(0)
+		a.Len(p.active, 0)
+		a.Len(p.inactive, 0)
+		mockNode.AssertExpectations(t)
+	}
+}
+
+func TestNodePool_BalanceNodeByFeature(t *testing.T) {
+	a := assert.New(t)
+	p := &NodePool{}
+	p.Init()
+
+	// success
+	{
+		p.featureMap["test"] = []Node{&MasterNode{}}
+		err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
+		a.NoError(err)
+		a.Equal(p.featureMap["test"][0], res)
+	}
+
+	// NoNodes
+	{
+		p.featureMap["test"] = []Node{}
+		err, res := p.BalanceNodeByFeature("test", balancer.NewBalancer("round-robin"))
+		a.Error(err)
+		a.Nil(res)
+	}
+
+	// No match feature
+	{
+		err, res := p.BalanceNodeByFeature("test2", balancer.NewBalancer("round-robin"))
+		a.Error(err)
+		a.Nil(res)
+	}
+}

pkg/cluster/slave_test.go

@@ -0,0 +1,443 @@
+package cluster
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
+	"github.com/cloudreve/Cloudreve/v3/pkg/request"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+	"github.com/stretchr/testify/assert"
+	testMock "github.com/stretchr/testify/mock"
+	"io/ioutil"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestSlaveNode_InitAndKill(t *testing.T) {
+	a := assert.New(t)
+	n := &SlaveNode{
+		callback: func(b bool, u uint) {
+
+		},
+	}
+
+	a.NotPanics(func() {
+		n.Init(&model.Node{})
+		time.Sleep(time.Millisecond * 500)
+		n.Init(&model.Node{})
+		n.Kill()
+	})
+}
+
+func TestSlaveNode_DummyMethods(t *testing.T) {
+	a := assert.New(t)
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+
+	m.Model.ID = 5
+	a.Equal(m.Model.ID, m.ID())
+	a.Equal(m.Model.ID, m.DBModel().ID)
+
+	a.False(m.IsActive())
+	a.False(m.IsMater())
+
+	m.SubscribeStatusChange(func(isActive bool, id uint) {})
+}
+
+func TestSlaveNode_IsFeatureEnabled(t *testing.T) {
+	a := assert.New(t)
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+
+	a.False(m.IsFeatureEnabled("aria2"))
+	a.False(m.IsFeatureEnabled("random"))
+	m.Model.Aria2Enabled = true
+	a.True(m.IsFeatureEnabled("aria2"))
+}
+
+func TestSlaveNode_Ping(t *testing.T) {
+	a := assert.New(t)
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+
+	// master return error code
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.Ping(&serializer.NodePingReq{})
+		a.Error(err)
+		a.Nil(res)
+		a.Equal(1, err.(serializer.AppError).Code)
+	}
+
+	// return unexpected json
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"233\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.Ping(&serializer.NodePingReq{})
+		a.Error(err)
+		a.Nil(res)
+	}
+
+	// return success
+	{
+		mockRequest := &requestMock{}
+		mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"{}\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.Ping(&serializer.NodePingReq{})
+		a.NoError(err)
+		a.NotNil(res)
+	}
+}
+
+func TestSlaveNode_GetAria2Instance(t *testing.T) {
+	a := assert.New(t)
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+
+	a.NotNil(m.GetAria2Instance())
+	m.Model.Aria2Enabled = true
+	a.NotNil(m.GetAria2Instance())
+	a.NotNil(m.GetAria2Instance())
+}
+
+func TestSlaveNode_StartPingLoop(t *testing.T) {
+	callbackCount := 0
+	finishedChan := make(chan struct{})
+	mockRequest := requestMock{}
+	mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
+		Response: &http.Response{
+			StatusCode: 404,
+		},
+	})
+	m := &SlaveNode{
+		Active: true,
+		Model:  &model.Node{},
+		callback: func(b bool, u uint) {
+			callbackCount++
+			if callbackCount == 2 {
+				close(finishedChan)
+			}
+			if callbackCount == 1 {
+				mockRequest.AssertExpectations(t)
+				mockRequest = requestMock{}
+				mockRequest.On("Request", "POST", "heartbeat", testMock.Anything, testMock.Anything).Return(&request.Response{
+					Response: &http.Response{
+						StatusCode: 200,
+						Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"{}\"}")),
+					},
+				})
+			}
+		},
+	}
+	cache.Set("setting_slave_ping_interval", "0", 0)
+	cache.Set("setting_slave_recover_interval", "0", 0)
+	cache.Set("setting_slave_node_retry", "1", 0)
+
+	m.caller.Client = &mockRequest
+	go func() {
+		select {
+		case <-finishedChan:
+			m.Kill()
+		}
+	}()
+	m.StartPingLoop()
+	mockRequest.AssertExpectations(t)
+}
+
+func TestSlaveNode_AuthInstance(t *testing.T) {
+	a := assert.New(t)
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+
+	a.NotNil(m.MasterAuthInstance())
+	a.NotNil(m.SlaveAuthInstance())
+}
+
+func TestSlaveNode_ChangeStatus(t *testing.T) {
+	a := assert.New(t)
+	isActive := false
+	m := &SlaveNode{
+		Model: &model.Node{},
+		callback: func(b bool, u uint) {
+			isActive = b
+		},
+	}
+
+	a.NotPanics(func() {
+		m.changeStatus(false)
+	})
+	m.changeStatus(true)
+	a.True(isActive)
+}
+
+func getTestRPCNodeSlave() *SlaveNode {
+	m := &SlaveNode{
+		Model: &model.Node{},
+	}
+	m.caller.parent = m
+	return m
+}
+
+func TestSlaveCaller_CreateTask(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNodeSlave()
+
+	// master return 404
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 404,
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.CreateTask(&model.Download{}, nil)
+		a.Empty(res)
+		a.Error(err)
+	}
+
+	// master return error
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.CreateTask(&model.Download{}, nil)
+		a.Empty(res)
+		a.Error(err)
+	}
+
+	// master return success
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/task", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.CreateTask(&model.Download{}, nil)
+		a.Equal("res", res)
+		a.NoError(err)
+	}
+}
+
+func TestSlaveCaller_Status(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNodeSlave()
+
+	// master return 404
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 404,
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.Status(&model.Download{})
+		a.Empty(res.Status)
+		a.Error(err)
+	}
+
+	// master return error
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.Status(&model.Download{})
+		a.Empty(res.Status)
+		a.Error(err)
+	}
+
+	// master return success
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/status", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"re456456s\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		res, err := m.caller.Status(&model.Download{})
+		a.Empty(res.Status)
+		a.NoError(err)
+	}
+}
+
+func TestSlaveCaller_Cancel(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNodeSlave()
+
+	// master return 404
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 404,
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Cancel(&model.Download{})
+		a.Error(err)
+	}
+
+	// master return error
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Cancel(&model.Download{})
+		a.Error(err)
+	}
+
+	// master return success
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/cancel", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Cancel(&model.Download{})
+		a.NoError(err)
+	}
+}
+
+func TestSlaveCaller_Select(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNodeSlave()
+	m.caller.Init()
+	m.caller.GetConfig()
+
+	// master return 404
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 404,
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Select(&model.Download{}, nil)
+		a.Error(err)
+	}
+
+	// master return error
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Select(&model.Download{}, nil)
+		a.Error(err)
+	}
+
+	// master return success
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/select", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.Select(&model.Download{}, nil)
+		a.NoError(err)
+	}
+}
+
+func TestSlaveCaller_DeleteTempFile(t *testing.T) {
+	a := assert.New(t)
+	m := getTestRPCNodeSlave()
+	m.caller.Init()
+	m.caller.GetConfig()
+
+	// master return 404
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 404,
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.DeleteTempFile(&model.Download{})
+		a.Error(err)
+	}
+
+	// master return error
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"code\":1}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.DeleteTempFile(&model.Download{})
+		a.Error(err)
+	}
+
+	// master return success
+	{
+		mockRequest := requestMock{}
+		mockRequest.On("Request", "POST", "aria2/delete", testMock.Anything, testMock.Anything).Return(&request.Response{
+			Response: &http.Response{
+				StatusCode: 200,
+				Body:       ioutil.NopCloser(strings.NewReader("{\"data\":\"res\"}")),
+			},
+		})
+		m.caller.Client = mockRequest
+		err := m.caller.DeleteTempFile(&model.Download{})
+		a.NoError(err)
+	}
+}

pkg/conf/conf.go

@@ -70,9 +70,13 @@ type redis struct {
 
 // 缩略图 配置
 type thumb struct {
-	MaxWidth   uint
-	MaxHeight  uint
-	FileSuffix string `validate:"min=1"`
+	MaxWidth      uint
+	MaxHeight     uint
+	FileSuffix    string `validate:"min=1"`
+	MaxTaskCount  int
+	EncodeMethod  string `validate:"eq=jpg|eq=png"`
+	EncodeQuality int    `validate:"gte=1,lte=100"`
+	GCAfterGen    bool
 }
 
 // 跨域配置

pkg/conf/defaults.go

@@ -44,16 +44,20 @@ var CaptchaConfig = &captcha{
 var CORSConfig = &cors{
 	AllowOrigins:     []string{"UNSET"},
 	AllowMethods:     []string{"PUT", "POST", "GET", "OPTIONS"},
-	AllowHeaders:     []string{"Cookie", "X-Policy", "Authorization", "Content-Length", "Content-Type", "X-Path", "X-FileName"},
+	AllowHeaders:     []string{"Cookie", "X-Cr-Policy", "Authorization", "Content-Length", "Content-Type", "X-Path", "X-FileName"},
 	AllowCredentials: false,
 	ExposeHeaders:    nil,
 }
 
 // ThumbConfig 缩略图配置
 var ThumbConfig = &thumb{
-	MaxWidth:   400,
-	MaxHeight:  300,
-	FileSuffix: "._thumb",
+	MaxWidth:      400,
+	MaxHeight:     300,
+	FileSuffix:    "._thumb",
+	MaxTaskCount:  -1,
+	EncodeMethod:  "jpg",
+	GCAfterGen:    false,
+	EncodeQuality: 85,
 }
 
 // SlaveConfig 从机配置

pkg/conf/version.go

@@ -1,13 +1,13 @@
 package conf
 
 // BackendVersion 当前后端版本号
-var BackendVersion = "3.4.0-beta1"
+var BackendVersion = "3.4.0"
 
 // RequiredDBVersion 与当前版本匹配的数据库版本
-var RequiredDBVersion = "3.4.0-beta1"
+var RequiredDBVersion = "3.4.0"
 
 // RequiredStaticVersion 与当前版本匹配的静态资源版本
-var RequiredStaticVersion = "3.4.0-beta1"
+var RequiredStaticVersion = "3.4.0"
 
 // IsPro 是否为Pro版本
 var IsPro = "false"

pkg/filesystem/driver/local/handler_test.go

@@ -190,7 +190,7 @@ func TestHandler_Source(t *testing.T) {
 
 	// 设定了CDN，解析失败
 	{
-		handler.Policy.BaseURL = string(0x7f)
+		handler.Policy.BaseURL = string([]byte{0x7f})
 		file := model.File{
 			Model: gorm.Model{
 				ID: 1,

pkg/filesystem/driver/onedrive/api.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -573,7 +574,7 @@ func sysError(err error) *RespError {
 
 func (client *Client) request(ctx context.Context, method string, url string, body io.Reader, option ...request.Option) (string, *RespError) {
 	// 获取凭证
-	err := client.UpdateCredential(ctx)
+	err := client.UpdateCredential(ctx, conf.SystemConfig.Mode == "slave")
 	if err != nil {
 		return "", sysError(err)
 	}

pkg/filesystem/driver/onedrive/client.go

@@ -2,6 +2,7 @@ package onedrive
 
 import (
 	"errors"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
 	"github.com/cloudreve/Cloudreve/v3/pkg/request"
@@ -28,7 +29,8 @@ type Client struct {
 	ClientSecret string
 	Redirect     string
 
-	Request request.Client
+	Request           request.Client
+	ClusterController cluster.Controller
 }
 
 // Endpoints OneDrive客户端相关设置
@@ -51,11 +53,12 @@ func NewClient(policy *model.Policy) (*Client, error) {
 		Credential: &Credential{
 			RefreshToken: policy.AccessKey,
 		},
-		Policy:       policy,
-		ClientID:     policy.BucketName,
-		ClientSecret: policy.SecretKey,
-		Redirect:     policy.OptionsSerialized.OdRedirect,
-		Request:      request.NewClient(),
+		Policy:            policy,
+		ClientID:          policy.BucketName,
+		ClientSecret:      policy.SecretKey,
+		Redirect:          policy.OptionsSerialized.OdRedirect,
+		Request:           request.NewClient(),
+		ClusterController: cluster.DefaultController,
 	}
 
 	if client.Endpoints.DriverResource == "" {

pkg/filesystem/driver/onedrive/oauth.go

@@ -10,9 +10,7 @@ import (
 	"time"
 
 	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
-	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/request"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 )
 
@@ -125,8 +123,8 @@ func (client *Client) ObtainToken(ctx context.Context, opts ...Option) (*Credent
 }
 
 // UpdateCredential 更新凭证，并检查有效期
-func (client *Client) UpdateCredential(ctx context.Context) error {
-	if conf.SystemConfig.Mode == "slave" {
+func (client *Client) UpdateCredential(ctx context.Context, isSlave bool) error {
+	if isSlave {
 		return client.fetchCredentialFromMaster(ctx)
 	}
 
@@ -179,7 +177,7 @@ func (client *Client) UpdateCredential(ctx context.Context) error {
 
 // UpdateCredential 更新凭证，并检查有效期
 func (client *Client) fetchCredentialFromMaster(ctx context.Context) error {
-	res, err := slave.DefaultController.GetOneDriveToken(client.Policy.MasterID, client.Policy.ID)
+	res, err := client.ClusterController.GetOneDriveToken(client.Policy.MasterID, client.Policy.ID)
 	if err != nil {
 		return err
 	}

pkg/filesystem/driver/onedrive/oauth_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mocks/controllermock"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -269,10 +270,10 @@ func TestClient_UpdateCredential(t *testing.T) {
 
 	// 无有效的RefreshToken
 	{
-		err := client.UpdateCredential(context.Background())
+		err := client.UpdateCredential(context.Background(), false)
 		asserts.Equal(ErrInvalidRefreshToken, err)
 		client.Credential = nil
-		err = client.UpdateCredential(context.Background())
+		err = client.UpdateCredential(context.Background(), false)
 		asserts.Equal(ErrInvalidRefreshToken, err)
 	}
 
@@ -299,7 +300,7 @@ func TestClient_UpdateCredential(t *testing.T) {
 		mock.ExpectBegin()
 		mock.ExpectExec("UPDATE(.+)").WillReturnResult(sqlmock.NewResult(1, 1))
 		mock.ExpectCommit()
-		err := client.UpdateCredential(context.Background())
+		err := client.UpdateCredential(context.Background(), false)
 		clientMock.AssertExpectations(t)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.NoError(err)
@@ -331,7 +332,7 @@ func TestClient_UpdateCredential(t *testing.T) {
 		client.Credential = &Credential{
 			RefreshToken: "old_refresh_token",
 		}
-		err := client.UpdateCredential(context.Background())
+		err := client.UpdateCredential(context.Background(), false)
 		clientMock.AssertExpectations(t)
 		asserts.Error(err)
 	}
@@ -346,7 +347,7 @@ func TestClient_UpdateCredential(t *testing.T) {
 		client.Credential = &Credential{
 			RefreshToken: "old_refresh_token",
 		}
-		err := client.UpdateCredential(context.Background())
+		err := client.UpdateCredential(context.Background(), false)
 		asserts.NoError(err)
 		asserts.Equal("AccessToken", client.Credential.AccessToken)
 		asserts.Equal("RefreshToken", client.Credential.RefreshToken)
@@ -359,8 +360,27 @@ func TestClient_UpdateCredential(t *testing.T) {
 			AccessToken:  "AccessToken2",
 			ExpiresIn:    time.Now().Add(time.Duration(10) * time.Second).Unix(),
 		}
-		err := client.UpdateCredential(context.Background())
+		err := client.UpdateCredential(context.Background(), false)
 		asserts.NoError(err)
 		asserts.Equal("AccessToken2", client.Credential.AccessToken)
 	}
+
+	// slave failed
+	{
+		mockController := &controllermock.SlaveControllerMock{}
+		mockController.On("GetOneDriveToken", testMock.Anything, testMock.Anything).Return("", errors.New("error"))
+		client.ClusterController = mockController
+		err := client.UpdateCredential(context.Background(), true)
+		asserts.Error(err)
+	}
+
+	// slave success
+	{
+		mockController := &controllermock.SlaveControllerMock{}
+		mockController.On("GetOneDriveToken", testMock.Anything, testMock.Anything).Return("AccessToken3", nil)
+		client.ClusterController = mockController
+		err := client.UpdateCredential(context.Background(), true)
+		asserts.NoError(err)
+		asserts.Equal("AccessToken3", client.Credential.AccessToken)
+	}
 }

pkg/filesystem/driver/remote/handler.go

@@ -153,10 +153,7 @@ func (handler Driver) Put(ctx context.Context, file io.ReadCloser, dst string, s
 	}
 
 	// 对文件名进行URLEncode
-	fileName, err := url.QueryUnescape(path.Base(dst))
-	if err != nil {
-		return err
-	}
+	fileName := url.QueryEscape(path.Base(dst))
 
 	// 决定是否要禁用文件覆盖
 	overwrite := "true"
@@ -170,9 +167,9 @@ func (handler Driver) Put(ctx context.Context, file io.ReadCloser, dst string, s
 		handler.Policy.GetUploadURL(),
 		file,
 		request.WithHeader(map[string][]string{
-			"X-Policy":    {credential.Policy},
-			"X-FileName":  {fileName},
-			"X-Overwrite": {overwrite},
+			"X-Cr-Policy":    {credential.Policy},
+			"X-Cr-FileName":  {fileName},
+			"X-Cr-Overwrite": {overwrite},
 		}),
 		request.WithContentLength(int64(size)),
 		request.WithTimeout(time.Duration(0)),
@@ -294,7 +291,7 @@ func (handler Driver) Source(
 	sourcePath := base64.RawURLEncoding.EncodeToString([]byte(path))
 	signedURI, err = auth.SignURI(
 		handler.AuthInstance,
-		fmt.Sprintf("%s/%d/%s/%s", controller, speed, sourcePath, fileName),
+		fmt.Sprintf("%s/%d/%s/%s", controller, speed, sourcePath, url.PathEscape(fileName)),
 		ttl,
 	)
 
@@ -335,8 +332,8 @@ func (handler Driver) getUploadCredential(ctx context.Context, policy serializer
 	// 签名上传策略
 	uploadRequest, _ := http.NewRequest("POST", "/api/v3/slave/upload", nil)
 	uploadRequest.Header = map[string][]string{
-		"X-Policy":    {policyEncoded},
-		"X-Overwrite": {"false"},
+		"X-Cr-Policy":    {policyEncoded},
+		"X-Cr-Overwrite": {"false"},
 	}
 	auth.SignRequest(handler.AuthInstance, uploadRequest, TTL)
 

pkg/filesystem/driver/remote/handler_test.go

@@ -105,7 +105,7 @@ func TestHandler_Source(t *testing.T) {
 	// 解析失败 自定义CDN
 	{
 		handler := Driver{
-			Policy:       &model.Policy{Server: "/", BaseURL: string(0x7f)},
+			Policy:       &model.Policy{Server: "/", BaseURL: string([]byte{0x7f})},
 			AuthInstance: auth.HMACAuth{},
 		}
 		file := model.File{

pkg/filesystem/file_test.go

@@ -228,25 +228,25 @@ func TestFileSystem_deleteGroupedFile(t *testing.T) {
 		},
 	}
 
-	// 全部失败
+	// 全部不存在
 	{
 		failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
 		asserts.Equal(map[uint][]string{
-			1: {"1_1.txt", "1_2.txt"},
-			2: {"2_1.txt", "2_2.txt"},
-			3: {"3_1.txt"},
+			1: {},
+			2: {},
+			3: {},
 		}, failed)
 	}
-	// 部分失败
+	// 部分不存在
 	{
 		file, err := os.Create(util.RelativePath("1_1.txt"))
 		asserts.NoError(err)
 		_ = file.Close()
 		failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
 		asserts.Equal(map[uint][]string{
-			1: {"1_2.txt"},
-			2: {"2_1.txt", "2_2.txt"},
-			3: {"3_1.txt"},
+			1: {},
+			2: {},
+			3: {},
 		}, failed)
 	}
 	// 部分失败,包含整组未知存储策略导致的失败
@@ -259,9 +259,9 @@ func TestFileSystem_deleteGroupedFile(t *testing.T) {
 		files[3].Policy.Type = "unknown"
 		failed := fs.deleteGroupedFile(ctx, fs.GroupFileByPolicy(ctx, files))
 		asserts.Equal(map[uint][]string{
-			1: {"1_2.txt"},
+			1: {},
 			2: {"2_1.txt", "2_2.txt"},
-			3: {"3_1.txt"},
+			3: {},
 		}, failed)
 	}
 }

pkg/filesystem/filesystem_test.go

@@ -1,6 +1,9 @@
 package filesystem
 
 import (
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/shadow/masterinslave"
+	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/shadow/slaveinmaster"
 	"net/http/httptest"
 
 	"github.com/DATA-DOG/go-sqlmock"
@@ -104,6 +107,10 @@ func TestDispatchHandler(t *testing.T) {
 	err = fs.DispatchHandler()
 	asserts.NoError(err)
 
+	fs.Policy = &model.Policy{Type: "cos"}
+	err = fs.DispatchHandler()
+	asserts.NoError(err)
+
 	fs.Policy = &model.Policy{Type: "s3"}
 	err = fs.DispatchHandler()
 	asserts.NoError(err)
@@ -262,3 +269,40 @@ func TestFileSystem_SetTargetByInterface(t *testing.T) {
 		asserts.Len(fs.FileTarget, 1)
 	}
 }
+
+func TestFileSystem_SwitchToSlaveHandler(t *testing.T) {
+	a := assert.New(t)
+	fs := FileSystem{
+		User: &model.User{},
+	}
+	mockNode := &cluster.MasterNode{
+		Model: &model.Node{},
+	}
+	fs.SwitchToSlaveHandler(mockNode)
+	a.IsType(&slaveinmaster.Driver{}, fs.Handler)
+}
+
+func TestFileSystem_SwitchToShadowHandler(t *testing.T) {
+	a := assert.New(t)
+	fs := FileSystem{
+		User:   &model.User{},
+		Policy: &model.Policy{},
+	}
+	mockNode := &cluster.MasterNode{
+		Model: &model.Node{},
+	}
+
+	// remote to local
+	{
+		fs.Policy.Type = "remote"
+		fs.SwitchToShadowHandler(mockNode, "", "")
+		a.IsType(&masterinslave.Driver{}, fs.Handler)
+	}
+
+	// local to remote
+	{
+		fs.Policy.Type = "local"
+		fs.SwitchToShadowHandler(mockNode, "", "")
+		a.IsType(&masterinslave.Driver{}, fs.Handler)
+	}
+}

pkg/filesystem/hooks_test.go

@@ -708,3 +708,29 @@ func TestHookGiveBackCapacity(t *testing.T) {
 		asserts.EqualValues(7, fs.User.Storage)
 	}
 }
+
+func TestHookValidateCapacityWithoutIncrease(t *testing.T) {
+	a := assert.New(t)
+	fs := &FileSystem{
+		User: &model.User{
+			Model:   gorm.Model{ID: 1},
+			Storage: 10,
+			Group:   model.Group{},
+		},
+	}
+	ctx := context.WithValue(context.Background(), fsctx.FileHeaderCtx, local.FileStream{Size: 1})
+
+	// not enough
+	{
+		fs.User.Group.MaxStorage = 10
+		a.Error(HookValidateCapacityWithoutIncrease(ctx, fs))
+		a.EqualValues(10, fs.User.Storage)
+	}
+
+	// enough
+	{
+		fs.User.Group.MaxStorage = 11
+		a.NoError(HookValidateCapacityWithoutIncrease(ctx, fs))
+		a.EqualValues(10, fs.User.Storage)
+	}
+}

pkg/filesystem/image.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"fmt"
 	"strconv"
+	"sync"
+
+	"runtime"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
 	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
@@ -35,18 +38,53 @@ func (fs *FileSystem) GetThumb(ctx context.Context, id uint) (*response.ContentR
 	ctx = context.WithValue(ctx, fsctx.ThumbSizeCtx, [2]uint{w, h})
 	ctx = context.WithValue(ctx, fsctx.FileModelCtx, fs.FileTarget[0])
 	res, err := fs.Handler.Thumb(ctx, fs.FileTarget[0].SourceName)
-	if err == nil && conf.SystemConfig.Mode == "master" {
-		res.MaxAge = model.GetIntSetting("preview_timeout", 60)
-	}
 
 	// 本地存储策略出错时重新生成缩略图
 	if err != nil && fs.Policy.Type == "local" {
 		fs.GenerateThumbnail(ctx, &fs.FileTarget[0])
+		res, err = fs.Handler.Thumb(ctx, fs.FileTarget[0].SourceName)
+	}
+
+	if err == nil && conf.SystemConfig.Mode == "master" {
+		res.MaxAge = model.GetIntSetting("preview_timeout", 60)
 	}
 
 	return res, err
 }
 
+// thumbPool 要使用的任务池
+var thumbPool *Pool
+var once sync.Once
+
+// Pool 带有最大配额的任务池
+type Pool struct {
+	// 容量
+	worker chan int
+}
+
+// Init 初始化任务池
+func getThumbWorker() *Pool {
+	once.Do(func() {
+		maxWorker := conf.ThumbConfig.MaxTaskCount
+		if maxWorker <= 0 {
+			maxWorker = runtime.GOMAXPROCS(0)
+		}
+		thumbPool = &Pool{
+			worker: make(chan int, maxWorker),
+		}
+		util.Log().Debug("初始化Thumb任务队列，WorkerNum = %d", maxWorker)
+	})
+	return thumbPool
+}
+func (pool *Pool) addWorker() {
+	pool.worker <- 1
+	util.Log().Debug("Thumb任务队列，addWorker")
+}
+func (pool *Pool) releaseWorker() {
+	util.Log().Debug("Thumb任务队列，releaseWorker")
+	<-pool.worker
+}
+
 // GenerateThumbnail 尝试为本地策略文件生成缩略图并获取图像原始大小
 // TODO 失败时，如果之前还有图像信息，则清除
 func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
@@ -65,6 +103,8 @@ func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
 		return
 	}
 	defer source.Close()
+	getThumbWorker().addWorker()
+	defer getThumbWorker().releaseWorker()
 
 	image, err := thumb.NewThumbFromFile(source, file.Name)
 	if err != nil {
@@ -79,6 +119,12 @@ func (fs *FileSystem) GenerateThumbnail(ctx context.Context, file *model.File) {
 	image.GetThumb(fs.GenerateThumbnailSize(w, h))
 	// 保存到文件
 	err = image.Save(util.RelativePath(file.SourceName + conf.ThumbConfig.FileSuffix))
+	image = nil
+	if conf.ThumbConfig.GCAfterGen {
+		util.Log().Debug("GenerateThumbnail runtime.GC")
+		runtime.GC()
+	}
+
 	if err != nil {
 		util.Log().Warning("无法保存缩略图：%s", err)
 		return

pkg/filesystem/image_test.go

@@ -38,3 +38,12 @@ func TestFileSystem_GetThumb(t *testing.T) {
 		asserts.EqualValues(50, res.MaxAge)
 	}
 }
+
+func TestFileSystem_ThumbWorker(t *testing.T) {
+	asserts := assert.New(t)
+
+	asserts.NotPanics(func() {
+		getThumbWorker().addWorker()
+		getThumbWorker().releaseWorker()
+	})
+}

pkg/filesystem/manage_test.go

@@ -356,47 +356,6 @@ func TestFileSystem_Delete(t *testing.T) {
 	}}
 	ctx := context.Background()
 
-	//全部未成功
-	{
-		// 列出要删除的目录
-		mock.ExpectQuery("SELECT(.+)").
-			WillReturnRows(
-				sqlmock.NewRows([]string{"id"}).
-					AddRow(1).
-					AddRow(2).
-					AddRow(3),
-			)
-		mock.ExpectQuery("SELECT(.+)").
-			WithArgs(1, 2, 3).
-			WillReturnRows(
-				sqlmock.NewRows([]string{"id", "name", "source_name", "policy_id", "size"}).
-					AddRow(4, "1.txt", "1.txt", 2, 1),
-			)
-		// 查询顶级的文件
-		mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "name", "source_name", "policy_id", "size"}).AddRow(1, "1.txt", "1.txt", 603, 2))
-		mock.ExpectQuery("SELECT(.+)files(.+)").
-			WillReturnRows(sqlmock.NewRows([]string{"id", "policy_id", "source_name"}))
-		// 查找软连接
-		mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id"}))
-		// 查询上传策略
-		mock.ExpectQuery("SELECT(.+)").WillReturnRows(sqlmock.NewRows([]string{"id", "type"}).AddRow(603, "local"))
-		// 删除文件记录
-		mock.ExpectBegin()
-		mock.ExpectExec("DELETE(.+)files").
-			WillReturnResult(sqlmock.NewResult(0, 3))
-		mock.ExpectCommit()
-		// 删除对应分享
-		mock.ExpectBegin()
-		mock.ExpectExec("UPDATE(.+)shares").
-			WillReturnResult(sqlmock.NewResult(0, 3))
-		mock.ExpectCommit()
-
-		err := fs.Delete(ctx, []uint{1}, []uint{1}, false)
-		asserts.Error(err)
-		asserts.Equal(203, err.(serializer.AppError).Code)
-		asserts.Equal(uint64(3), fs.User.Storage)
-		asserts.NoError(mock.ExpectationsWereMet())
-	}
 	//全部未成功，强制
 	{
 		fs.CleanTargets()

pkg/mocks/controllermock/c.go

@@ -0,0 +1,43 @@
+package controllermock
+
+import (
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+	"github.com/stretchr/testify/mock"
+)
+
+type SlaveControllerMock struct {
+	mock.Mock
+}
+
+func (s SlaveControllerMock) HandleHeartBeat(pingReq *serializer.NodePingReq) (serializer.NodePingResp, error) {
+	args := s.Called(pingReq)
+	return args.Get(0).(serializer.NodePingResp), args.Error(1)
+}
+
+func (s SlaveControllerMock) GetAria2Instance(s2 string) (common.Aria2, error) {
+	args := s.Called(s2)
+	return args.Get(0).(common.Aria2), args.Error(1)
+}
+
+func (s SlaveControllerMock) SendNotification(s3 string, s2 string, message mq.Message) error {
+	args := s.Called(s3, s2, message)
+	return args.Error(0)
+}
+
+func (s SlaveControllerMock) SubmitTask(s3 string, i interface{}, s2 string, f func(interface{})) error {
+	args := s.Called(s3, i, s2, f)
+	return args.Error(0)
+}
+
+func (s SlaveControllerMock) GetMasterInfo(s2 string) (*cluster.MasterInfo, error) {
+	args := s.Called(s2)
+	return args.Get(0).(*cluster.MasterInfo), args.Error(1)
+}
+
+func (s SlaveControllerMock) GetOneDriveToken(s2 string, u uint) (string, error) {
+	args := s.Called(s2, u)
+	return args.String(0), args.Error(1)
+}

pkg/mocks/mocks.go

@@ -0,0 +1,161 @@
+package mocks
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/common"
+	"github.com/cloudreve/Cloudreve/v3/pkg/aria2/rpc"
+	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
+	"github.com/cloudreve/Cloudreve/v3/pkg/balancer"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/request"
+	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
+	"github.com/cloudreve/Cloudreve/v3/pkg/task"
+	testMock "github.com/stretchr/testify/mock"
+	"io"
+)
+
+type NodePoolMock struct {
+	testMock.Mock
+}
+
+func (n NodePoolMock) BalanceNodeByFeature(feature string, lb balancer.Balancer) (error, cluster.Node) {
+	args := n.Called(feature, lb)
+	return args.Error(0), args.Get(1).(cluster.Node)
+}
+
+func (n NodePoolMock) GetNodeByID(id uint) cluster.Node {
+	args := n.Called(id)
+	if res, ok := args.Get(0).(cluster.Node); ok {
+		return res
+	}
+
+	return nil
+}
+
+func (n NodePoolMock) Add(node *model.Node) {
+	n.Called(node)
+}
+
+func (n NodePoolMock) Delete(id uint) {
+	n.Called(id)
+}
+
+type NodeMock struct {
+	testMock.Mock
+}
+
+func (n NodeMock) Init(node *model.Node) {
+	n.Called(node)
+}
+
+func (n NodeMock) IsFeatureEnabled(feature string) bool {
+	args := n.Called(feature)
+	return args.Bool(0)
+}
+
+func (n NodeMock) SubscribeStatusChange(callback func(isActive bool, id uint)) {
+	n.Called(callback)
+}
+
+func (n NodeMock) Ping(req *serializer.NodePingReq) (*serializer.NodePingResp, error) {
+	args := n.Called(req)
+	return args.Get(0).(*serializer.NodePingResp), args.Error(1)
+}
+
+func (n NodeMock) IsActive() bool {
+	args := n.Called()
+	return args.Bool(0)
+}
+
+func (n NodeMock) GetAria2Instance() common.Aria2 {
+	args := n.Called()
+	return args.Get(0).(common.Aria2)
+}
+
+func (n NodeMock) ID() uint {
+	args := n.Called()
+	return args.Get(0).(uint)
+}
+
+func (n NodeMock) Kill() {
+	n.Called()
+}
+
+func (n NodeMock) IsMater() bool {
+	args := n.Called()
+	return args.Bool(0)
+}
+
+func (n NodeMock) MasterAuthInstance() auth.Auth {
+	args := n.Called()
+	return args.Get(0).(auth.Auth)
+}
+
+func (n NodeMock) SlaveAuthInstance() auth.Auth {
+	args := n.Called()
+	return args.Get(0).(auth.Auth)
+}
+
+func (n NodeMock) DBModel() *model.Node {
+	args := n.Called()
+	return args.Get(0).(*model.Node)
+}
+
+type Aria2Mock struct {
+	testMock.Mock
+}
+
+func (a Aria2Mock) Init() error {
+	args := a.Called()
+	return args.Error(0)
+}
+
+func (a Aria2Mock) CreateTask(task *model.Download, options map[string]interface{}) (string, error) {
+	args := a.Called(task, options)
+	return args.String(0), args.Error(1)
+}
+
+func (a Aria2Mock) Status(task *model.Download) (rpc.StatusInfo, error) {
+	args := a.Called(task)
+	return args.Get(0).(rpc.StatusInfo), args.Error(1)
+}
+
+func (a Aria2Mock) Cancel(task *model.Download) error {
+	args := a.Called(task)
+	return args.Error(0)
+}
+
+func (a Aria2Mock) Select(task *model.Download, files []int) error {
+	args := a.Called(task, files)
+	return args.Error(0)
+}
+
+func (a Aria2Mock) GetConfig() model.Aria2Option {
+	args := a.Called()
+	return args.Get(0).(model.Aria2Option)
+}
+
+func (a Aria2Mock) DeleteTempFile(download *model.Download) error {
+	args := a.Called(download)
+	return args.Error(0)
+}
+
+type TaskPoolMock struct {
+	testMock.Mock
+}
+
+func (t TaskPoolMock) Add(num int) {
+	t.Called(num)
+}
+
+func (t TaskPoolMock) Submit(job task.Job) {
+	t.Called(job)
+}
+
+type RequestMock struct {
+	testMock.Mock
+}
+
+func (r RequestMock) Request(method, target string, body io.Reader, opts ...request.Option) *request.Response {
+	return r.Called(method, target, body, opts).Get(0).(*request.Response)
+}

pkg/request/request.go

@@ -97,13 +97,13 @@ func (c HTTPClient) Request(method, target string, body io.Reader, opts ...Optio
 	}
 
 	if options.masterMeta && conf.SystemConfig.Mode == "master" {
-		req.Header.Add("X-Site-Url", model.GetSiteURL().String())
-		req.Header.Add("X-Site-Id", model.GetSettingByName("siteID"))
-		req.Header.Add("X-Cloudreve-Version", conf.BackendVersion)
+		req.Header.Add("X-Cr-Site-Url", model.GetSiteURL().String())
+		req.Header.Add("X-Cr-Site-Id", model.GetSettingByName("siteID"))
+		req.Header.Add("X-Cr-Cloudreve-Version", conf.BackendVersion)
 	}
 
 	if options.slaveNodeID != "" && conf.SystemConfig.Mode == "slave" {
-		req.Header.Add("X-Node-Id", options.slaveNodeID)
+		req.Header.Add("X-Cr-Node-Id", options.slaveNodeID)
 	}
 
 	if options.contentLength != -1 {

pkg/request/request_test.go

@@ -3,6 +3,7 @@ package request
 import (
 	"context"
 	"errors"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cache"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -54,16 +55,19 @@ func TestWithContext(t *testing.T) {
 
 func TestHTTPClient_Request(t *testing.T) {
 	asserts := assert.New(t)
-	client := HTTPClient{}
+	client := NewClient(WithSlaveMeta("test"))
 
 	// 正常
 	{
 		resp := client.Request(
-			"GET",
-			"http://cloudreveisnotexist.com",
+			"POST",
+			"/test",
 			strings.NewReader(""),
+			WithContentLength(0),
+			WithEndpoint("http://cloudreveisnotexist.com"),
 			WithTimeout(time.Duration(1)*time.Microsecond),
 			WithCredential(auth.HMACAuth{SecretKey: []byte("123")}, 10),
+			WithoutHeader([]string{"origin", "origin"}),
 		)
 		asserts.Error(resp.Err)
 		asserts.Nil(resp.Response)
@@ -224,3 +228,11 @@ func TestNopRSCloser_SetFirstFakeChunk(t *testing.T) {
 	rsc.SetContentLength(20)
 	asserts.EqualValues(20, rsc.status.Size)
 }
+
+func TestBlackHole(t *testing.T) {
+	a := assert.New(t)
+	cache.Set("setting_reset_after_upload_failed", "true", 0)
+	a.NotPanics(func() {
+		BlackHole(strings.NewReader("TestBlackHole"))
+	})
+}

pkg/serializer/error_test.go

@@ -0,0 +1,42 @@
+package serializer
+
+import (
+	"errors"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewError(t *testing.T) {
+	a := assert.New(t)
+	err := NewError(400, "Bad Request", errors.New("error"))
+	a.Error(err)
+	a.EqualValues(400, err.Code)
+
+	err.WithError(errors.New("error2"))
+	a.Equal("error2", err.RawError.Error())
+	a.Equal("Bad Request", err.Error())
+
+	resp := &Response{
+		Code:  400,
+		Msg:   "Bad Request",
+		Error: "error",
+	}
+	err = NewErrorFromResponse(resp)
+	a.Error(err)
+}
+
+func TestDBErr(t *testing.T) {
+	a := assert.New(t)
+	resp := DBErr("", nil)
+	a.NotEmpty(resp.Msg)
+
+	resp = ParamErr("", nil)
+	a.NotEmpty(resp.Msg)
+}
+
+func TestErr(t *testing.T) {
+	a := assert.New(t)
+	err := NewError(400, "Bad Request", errors.New("error"))
+	resp := Err(400, "", err)
+	a.Equal("Bad Request", resp.Msg)
+}

pkg/serializer/response_test.go

@@ -0,0 +1,33 @@
+package serializer
+
+import (
+	"encoding/json"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestNewResponseWithGobData(t *testing.T) {
+	a := assert.New(t)
+	type args struct {
+		data interface{}
+	}
+
+	res := NewResponseWithGobData(args{})
+	a.Equal(CodeInternalSetting, res.Code)
+
+	res = NewResponseWithGobData("TestNewResponseWithGobData")
+	a.Equal(0, res.Code)
+	a.NotEmpty(res.Data)
+}
+
+func TestResponse_GobDecode(t *testing.T) {
+	a := assert.New(t)
+	res := NewResponseWithGobData("TestResponse_GobDecode")
+	jsonContent, err := json.Marshal(res)
+	a.NoError(err)
+	resDecoded := &Response{}
+	a.NoError(json.Unmarshal(jsonContent, resDecoded))
+	var target string
+	resDecoded.GobDecode(&target)
+	a.Equal("TestResponse_GobDecode", target)
+}

pkg/serializer/slave_test.go

@@ -0,0 +1,20 @@
+package serializer
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestSlaveTransferReq_Hash(t *testing.T) {
+	a := assert.New(t)
+	s1 := &SlaveTransferReq{
+		Src:    "1",
+		Policy: &model.Policy{},
+	}
+	s2 := &SlaveTransferReq{
+		Src:    "2",
+		Policy: &model.Policy{},
+	}
+	a.NotEqual(s1.Hash("1"), s2.Hash("1"))
+}

pkg/slave/errors.go

@@ -1,7 +0,0 @@
-package slave
-
-import "github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-
-var (
-	ErrMasterNotFound = serializer.NewError(serializer.CodeMasterNotFound, "未知的主机节点", nil)
-)

pkg/task/compress_test.go

@@ -144,7 +144,7 @@ func TestCompressTask_Do(t *testing.T) {
 		task.Do()
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.NotEmpty(task.GetError().Msg)
-		asserts.True(util.IsEmpty("test/compress"))
+		asserts.True(util.IsEmpty(util.RelativePath("test/compress")))
 	}
 }
 

pkg/task/job.go

@@ -82,7 +82,7 @@ func Record(job Job) (*model.Task, error) {
 }
 
 // Resume 从数据库中恢复未完成任务
-func Resume() {
+func Resume(p Pool) {
 	tasks := model.GetTasksByStatus(Queued, Processing)
 	if len(tasks) == 0 {
 		return
@@ -97,7 +97,7 @@ func Resume() {
 		}
 
 		if job != nil {
-			TaskPoll.Submit(job)
+			p.Submit(job)
 		}
 	}
 }

pkg/task/job_test.go

@@ -2,6 +2,7 @@ package task
 
 import (
 	"errors"
+	testMock "github.com/stretchr/testify/mock"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
@@ -21,15 +22,46 @@ func TestRecord(t *testing.T) {
 	asserts.NoError(err)
 }
 
+type taskPoolMock struct {
+	testMock.Mock
+}
+
+func (t taskPoolMock) Add(num int) {
+	t.Called(num)
+}
+
+func (t taskPoolMock) Submit(job Job) {
+	t.Called(job)
+}
+
 func TestResume(t *testing.T) {
 	asserts := assert.New(t)
+	mockPool := taskPoolMock{}
 
 	// 没有任务
 	{
-		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued).WillReturnRows(sqlmock.NewRows([]string{"type"}))
-		Resume()
+		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}))
+		Resume(mockPool)
 		asserts.NoError(mock.ExpectationsWereMet())
 	}
+
+	// 有任务, 类型未知
+	{
+		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(233))
+		Resume(mockPool)
+		asserts.NoError(mock.ExpectationsWereMet())
+	}
+
+	// 有任务
+	{
+		mockPool.On("Submit", testMock.Anything)
+		mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type", "props"}).AddRow(CompressTaskType, "{}"))
+		mock.ExpectQuery("SELECT(.+)users").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+		mock.ExpectQuery("SELECT(.+)policies").WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
+		Resume(mockPool)
+		asserts.NoError(mock.ExpectationsWereMet())
+		mockPool.AssertExpectations(t)
+	}
 }
 
 func TestGetJobFromModel(t *testing.T) {

pkg/task/pool.go

@@ -7,23 +7,28 @@ import (
 )
 
 // TaskPoll 要使用的任务池
-var TaskPoll *Pool
+var TaskPoll Pool
 
-// Pool 带有最大配额的任务池
-type Pool struct {
+type Pool interface {
+	Add(num int)
+	Submit(job Job)
+}
+
+// AsyncPool 带有最大配额的任务池
+type AsyncPool struct {
 	// 容量
 	idleWorker chan int
 }
 
 // Add 增加可用Worker数量
-func (pool *Pool) Add(num int) {
+func (pool *AsyncPool) Add(num int) {
 	for i := 0; i < num; i++ {
 		pool.idleWorker <- 1
 	}
 }
 
 // ObtainWorker 阻塞直到获取新的Worker
-func (pool *Pool) ObtainWorker() Worker {
+func (pool *AsyncPool) obtainWorker() Worker {
 	select {
 	case <-pool.idleWorker:
 		// 有空闲Worker名额时，返回新Worker
@@ -32,32 +37,32 @@ func (pool *Pool) ObtainWorker() Worker {
 }
 
 // FreeWorker 添加空闲Worker
-func (pool *Pool) FreeWorker() {
+func (pool *AsyncPool) freeWorker() {
 	pool.Add(1)
 }
 
 // Submit 开始提交任务
-func (pool *Pool) Submit(job Job) {
+func (pool *AsyncPool) Submit(job Job) {
 	go func() {
 		util.Log().Debug("等待获取Worker")
-		worker := pool.ObtainWorker()
+		worker := pool.obtainWorker()
 		util.Log().Debug("获取到Worker")
 		worker.Do(job)
 		util.Log().Debug("释放Worker")
-		pool.FreeWorker()
+		pool.freeWorker()
 	}()
 }
 
 // Init 初始化任务池
 func Init() {
 	maxWorker := model.GetIntSetting("max_worker_num", 10)
-	TaskPoll = &Pool{
+	TaskPoll = &AsyncPool{
 		idleWorker: make(chan int, maxWorker),
 	}
 	TaskPoll.Add(maxWorker)
 	util.Log().Info("初始化任务队列，WorkerNum = %d", maxWorker)
 
 	if conf.SystemConfig.Mode == "master" {
-		Resume()
+		Resume(TaskPoll)
 	}
 }

pkg/task/pool_test.go

@@ -29,15 +29,15 @@ func TestMain(m *testing.M) {
 func TestInit(t *testing.T) {
 	asserts := assert.New(t)
 	cache.Set("setting_max_worker_num", "10", 0)
-	mock.ExpectQuery("SELECT(.+)").WithArgs(Queued).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(-1))
+	mock.ExpectQuery("SELECT(.+)").WithArgs(Queued, Processing).WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow(-1))
 	Init()
 	asserts.NoError(mock.ExpectationsWereMet())
-	asserts.Len(TaskPoll.idleWorker, 10)
+	asserts.Len(TaskPoll.(*AsyncPool).idleWorker, 10)
 }
 
 func TestPool_Submit(t *testing.T) {
 	asserts := assert.New(t)
-	pool := &Pool{
+	pool := &AsyncPool{
 		idleWorker: make(chan int, 1),
 	}
 	pool.Add(1)

pkg/task/slavetask/transfer.go

@@ -3,11 +3,11 @@ package slavetask
 import (
 	"context"
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
 	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/cloudreve/Cloudreve/v3/pkg/task"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 	"os"
@@ -68,8 +68,8 @@ func (job *TransferTask) SetErrorMsg(msg string, err error) {
 		},
 	}
 
-	if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
-		util.Log().Warning("无法发送转存失败通知到从机, ", err)
+	if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
+		util.Log().Warning("无法发送转存失败通知到从机, %s", err)
 	}
 }
 
@@ -94,7 +94,7 @@ func (job *TransferTask) Do() {
 		return
 	}
 
-	master, err := slave.DefaultController.GetMasterInfo(job.MasterID)
+	master, err := cluster.DefaultController.GetMasterInfo(job.MasterID)
 	if err != nil {
 		job.SetErrorMsg("找不到主机节点", err)
 		return
@@ -131,8 +131,8 @@ func (job *TransferTask) Do() {
 		Content:     serializer.SlaveTransferResult{},
 	}
 
-	if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
-		util.Log().Warning("无法发送转存成功通知到从机, ", err)
+	if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
+		util.Log().Warning("无法发送转存成功通知到从机, %s", err)
 	}
 }
 

pkg/task/transfer_test.go

@@ -138,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{}, "/", "/", false)
+		job, err := NewTransferTask(1, []string{}, "/", "/", false, 0, nil)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.NotNil(job)
 		asserts.NoError(err)
@@ -150,7 +150,7 @@ func TestNewTransferTask(t *testing.T) {
 		mock.ExpectBegin()
 		mock.ExpectExec("INSERT(.+)").WillReturnError(errors.New("error"))
 		mock.ExpectRollback()
-		job, err := NewTransferTask(1, []string{}, "/", "/", false)
+		job, err := NewTransferTask(1, []string{}, "/", "/", false, 0, nil)
 		asserts.NoError(mock.ExpectationsWereMet())
 		asserts.Nil(job)
 		asserts.Error(err)

pkg/thumb/image.go

@@ -12,9 +12,11 @@ import (
 	"strings"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 
-	"github.com/nfnt/resize"
+	//"github.com/nfnt/resize"
+	"golang.org/x/image/draw"
 )
 
 // Thumb 缩略图
@@ -58,7 +60,8 @@ func NewThumbFromFile(file io.Reader, name string) (*Thumb, error) {
 
 // GetThumb 生成给定最大尺寸的缩略图
 func (image *Thumb) GetThumb(width, height uint) {
-	image.src = resize.Thumbnail(width, height, image.src, resize.Lanczos3)
+	//image.src = resize.Thumbnail(width, height, image.src, resize.Lanczos3)
+	image.src = Thumbnail(width, height, image.src)
 }
 
 // GetSize 获取图像尺寸
@@ -75,12 +78,59 @@ func (image *Thumb) Save(path string) (err error) {
 		return err
 	}
 	defer out.Close()
+	switch conf.ThumbConfig.EncodeMethod {
+	case "png":
+		err = png.Encode(out, image.src)
+	default:
+		err = jpeg.Encode(out, image.src, &jpeg.Options{Quality: conf.ThumbConfig.EncodeQuality})
+	}
 
-	err = png.Encode(out, image.src)
 	return err
 
 }
 
+// Thumbnail will downscale provided image to max width and height preserving
+// original aspect ratio and using the interpolation function interp.
+// It will return original image, without processing it, if original sizes
+// are already smaller than provided constraints.
+func Thumbnail(maxWidth, maxHeight uint, img image.Image) image.Image {
+	origBounds := img.Bounds()
+	origWidth := uint(origBounds.Dx())
+	origHeight := uint(origBounds.Dy())
+	newWidth, newHeight := origWidth, origHeight
+
+	// Return original image if it have same or smaller size as constraints
+	if maxWidth >= origWidth && maxHeight >= origHeight {
+		return img
+	}
+
+	// Preserve aspect ratio
+	if origWidth > maxWidth {
+		newHeight = uint(origHeight * maxWidth / origWidth)
+		if newHeight < 1 {
+			newHeight = 1
+		}
+		newWidth = maxWidth
+	}
+
+	if newHeight > maxHeight {
+		newWidth = uint(newWidth * maxHeight / newHeight)
+		if newWidth < 1 {
+			newWidth = 1
+		}
+		newHeight = maxHeight
+	}
+	return Resize(newWidth, newHeight, img)
+}
+
+func Resize(newWidth, newHeight uint, img image.Image) image.Image {
+	// Set the expected size that you want:
+	dst := image.NewRGBA(image.Rect(0, 0, int(newWidth), int(newHeight)))
+	// Resize:
+	draw.BiLinear.Scale(dst, dst.Rect, img, img.Bounds(), draw.Src, nil)
+	return dst
+}
+
 // CreateAvatar 创建头像
 func (image *Thumb) CreateAvatar(uid uint) error {
 	// 读取头像相关设定
@@ -92,7 +142,8 @@ func (image *Thumb) CreateAvatar(uid uint) error {
 	// 生成头像缩略图
 	src := image.src
 	for k, size := range []int{s, m, l} {
-		image.src = resize.Resize(uint(size), uint(size), src, resize.Lanczos3)
+		//image.src = resize.Resize(uint(size), uint(size), src, resize.Lanczos3)
+		image.src = Resize(uint(size), uint(size), src)
 		err := image.Save(filepath.Join(savePath, fmt.Sprintf("avatar_%d_%d.png", uid, k)))
 		if err != nil {
 			return err

pkg/thumb/image_test.go

@@ -86,6 +86,30 @@ func TestThumb_GetThumb(t *testing.T) {
 	})
 }
 
+func TestThumb_Thumbnail(t *testing.T) {
+	asserts := assert.New(t)
+	{
+		img := image.NewRGBA(image.Rect(0, 0, 500, 200))
+		thumb := Thumbnail(100, 100, img)
+		asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 40))
+	}
+	{
+		img := image.NewRGBA(image.Rect(0, 0, 200, 200))
+		thumb := Thumbnail(100, 100, img)
+		asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 100))
+	}
+	{
+		img := image.NewRGBA(image.Rect(0, 0, 500, 500))
+		thumb := Thumbnail(100, 100, img)
+		asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 100, 100))
+	}
+	{
+		img := image.NewRGBA(image.Rect(0, 0, 200, 500))
+		thumb := Thumbnail(100, 100, img)
+		asserts.Equal(thumb.Bounds(), image.Rect(0, 0, 40, 100))
+	}
+}
+
 func TestThumb_Save(t *testing.T) {
 	asserts := assert.New(t)
 	file := CreateTestImage()

pkg/webdav/file.go

@@ -38,23 +38,26 @@ func moveFiles(ctx context.Context, fs *filesystem.FileSystem, src FileInfo, dst
 		fileIDs = []uint{src.(*model.File).ID}
 	}
 
-	// 判断是否为重命名
-	if src.GetPosition() == path.Dir(dst) {
-		err = fs.Rename(
-			ctx,
-			folderIDs,
-			fileIDs,
-			path.Base(dst),
-		)
-	} else {
-		err = fs.Move(
-			ctx,
-			folderIDs,
-			fileIDs,
-			src.GetPosition(),
-			path.Dir(dst),
-		)
-	}
+        // 判断是否需要移动
+        if src.GetPosition() != path.Dir(dst) {
+                err = fs.Move(
+                        ctx,
+                        folderIDs,
+                        fileIDs,
+                        src.GetPosition(),
+                        path.Dir(dst),
+                )
+        }
+
+        // 判断是否需要重命名
+        if err == nil && src.GetName() != path.Base(dst) {
+                err = fs.Rename(
+                        ctx,
+                        folderIDs,
+                        fileIDs,
+                        path.Base(dst),
+                )
+        }
 
 	if err != nil {
 		return http.StatusInternalServerError, err

routers/controllers/admin.go

@@ -1,6 +1,8 @@
 package controllers
 
 import (
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"io"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
@@ -72,7 +74,7 @@ func AdminReloadService(c *gin.Context) {
 	case "email":
 		email.Init()
 	case "aria2":
-		aria2.Init(true)
+		aria2.Init(true, cluster.Default, mq.GlobalMQ)
 	}
 
 	c.JSON(200, serializer.Response{})

routers/controllers/file.go

@@ -1,6 +1,5 @@
 package controllers
 
-import "C"
 import (
 	"context"
 	"fmt"
@@ -10,6 +9,7 @@ import (
 	"sync"
 
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/local"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
@@ -180,7 +180,7 @@ func Thumb(c *gin.Context) {
 	}
 
 	defer resp.Content.Close()
-	http.ServeContent(c.Writer, c.Request, "thumb.png", fs.FileTarget[0].UpdatedAt, resp.Content)
+	http.ServeContent(c.Writer, c.Request, "thumb."+conf.ThumbConfig.EncodeMethod, fs.FileTarget[0].UpdatedAt, resp.Content)
 
 }
 
@@ -308,8 +308,8 @@ func FileUploadStream(c *gin.Context) {
 	}
 
 	// 解码文件名和路径
-	fileName, err := url.QueryUnescape(c.Request.Header.Get("X-FileName"))
-	filePath, err := url.QueryUnescape(c.Request.Header.Get("X-Path"))
+	fileName, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-FileName"))
+	filePath, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-Path"))
 	if err != nil {
 		c.JSON(200, ErrorResponse(err))
 		return

routers/controllers/slave.go

@@ -32,7 +32,7 @@ func SlaveUpload(c *gin.Context) {
 	fs.Handler = local.Driver{}
 
 	// 从请求中取得上传策略
-	uploadPolicyRaw := c.GetHeader("X-Policy")
+	uploadPolicyRaw := c.GetHeader("X-Cr-Policy")
 	if uploadPolicyRaw == "" {
 		c.JSON(200, serializer.ParamErr("未指定上传策略", nil))
 		return
@@ -54,7 +54,7 @@ func SlaveUpload(c *gin.Context) {
 	}
 
 	// 解码文件名和路径
-	fileName, err := url.QueryUnescape(c.Request.Header.Get("X-FileName"))
+	fileName, err := url.QueryUnescape(c.Request.Header.Get("X-Cr-FileName"))
 	if err != nil {
 		c.JSON(200, ErrorResponse(err))
 		return
@@ -74,7 +74,7 @@ func SlaveUpload(c *gin.Context) {
 	fs.Use("AfterValidateFailed", filesystem.HookDeleteTempFile)
 
 	// 是否允许覆盖
-	if c.Request.Header.Get("X-Overwrite") == "false" {
+	if c.Request.Header.Get("X-Cr-Overwrite") == "false" {
 		ctx = context.WithValue(ctx, fsctx.DisableOverwrite, true)
 	}
 

routers/file_router_test.go

@@ -73,7 +73,7 @@ func TestLocalFileUpload(t *testing.T) {
 					strings.NewReader("2333"),
 				)
 				req.Header.Add("Content-Length", "4")
-				req.Header.Add("X-FileName", "大地的%sfsf")
+				req.Header.Add("X-Cr-FileName", "大地的%sfsf")
 				return req
 			},
 			ExpectCode: 40002,
@@ -87,8 +87,8 @@ func TestLocalFileUpload(t *testing.T) {
 					strings.NewReader("2333"),
 				)
 				req.Header.Add("Content-Length", "4")
-				req.Header.Add("X-FileName", "TestFileUploadRoute.txt")
-				req.Header.Add("X-Path", "/")
+				req.Header.Add("X-Cr-FileName", "TestFileUploadRoute.txt")
+				req.Header.Add("X-Cr-Path", "/")
 				return req
 			},
 			ExpectCode: 0,

routers/router.go

@@ -3,6 +3,7 @@ package routers
 import (
 	"github.com/cloudreve/Cloudreve/v3/middleware"
 	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/hashid"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
@@ -59,7 +60,7 @@ func InitSlaveRouter() *gin.Engine {
 
 		// 离线下载
 		aria2 := v3.Group("aria2")
-		aria2.Use(middleware.UseSlaveAria2Instance())
+		aria2.Use(middleware.UseSlaveAria2Instance(cluster.DefaultController))
 		{
 			// 创建离线下载任务
 			aria2.POST("task", controllers.SlaveAria2Create)
@@ -205,7 +206,7 @@ func InitMasterRouter() *gin.Engine {
 
 		// 从机的 RPC 通信
 		slave := v3.Group("slave")
-		slave.Use(middleware.SlaveRPCSignRequired())
+		slave.Use(middleware.SlaveRPCSignRequired(cluster.Default))
 		{
 			// 事件通知
 			slave.PUT("notification/:subject", controllers.SlaveNotificationPush)

service/admin/node.go

@@ -24,7 +24,10 @@ func (service *AddNodeService) Add() serializer.Response {
 		}
 	}
 
-	cluster.Default.Add(&service.Node)
+	if service.Node.Status == model.NodeActive {
+		cluster.Default.Add(&service.Node)
+	}
+
 	return serializer.Response{Data: service.Node.ID}
 }
 

service/aria2/add.go

@@ -9,7 +9,6 @@ import (
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
 	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 	"github.com/gin-gonic/gin"
 )
@@ -49,9 +48,7 @@ func (service *AddURLService) Add(c *gin.Context, taskType int) serializer.Respo
 	}
 
 	// 获取 Aria2 负载均衡器
-	aria2.Lock.RLock()
-	lb := aria2.LB
-	aria2.Lock.RUnlock()
+	lb := aria2.GetLoadBalancer()
 
 	// 获取 Aria2 实例
 	err, node := cluster.Default.BalanceNodeByFeature("aria2", lb)
@@ -73,7 +70,7 @@ func (service *AddURLService) Add(c *gin.Context, taskType int) serializer.Respo
 	}
 
 	// 创建任务监控
-	monitor.NewMonitor(task)
+	monitor.NewMonitor(task, cluster.Default, mq.GlobalMQ)
 
 	return serializer.Response{}
 }
@@ -91,7 +88,7 @@ func Add(c *gin.Context, service *serializer.SlaveAria2Call) serializer.Response
 	// 创建事件通知回调
 	siteID, _ := c.Get("MasterSiteID")
 	mq.GlobalMQ.SubscribeCallback(gid, func(message mq.Message) {
-		if err := slave.DefaultController.SendNotification(siteID.(string), message.TriggeredBy, message); err != nil {
+		if err := cluster.DefaultController.SendNotification(siteID.(string), message.TriggeredBy, message); err != nil {
 			util.Log().Warning("无法发送离线下载任务状态变更通知, %s", err)
 		}
 	})

service/explorer/file.go

@@ -2,8 +2,10 @@ package explorer
 
 import (
 	"context"
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
+	"github.com/cloudreve/Cloudreve/v3/pkg/util"
 	"io/ioutil"
 	"net/http"
 	"net/url"
@@ -222,15 +224,15 @@ func (service *FileIDService) CreateDocPreviewSession(ctx context.Context, c *gi
 	}
 
 	// 生成最终的预览器地址
-	// TODO 从配置文件中读取
-	viewerBase, _ := url.Parse("https://view.officeapps.live.com/op/view.aspx")
-	params := viewerBase.Query()
-	params.Set("src", downloadURL)
-	viewerBase.RawQuery = params.Encode()
-
+	srcB64 := base64.StdEncoding.EncodeToString([]byte(downloadURL))
+	srcEncoded := url.QueryEscape(downloadURL)
+	srcB64Encoded := url.QueryEscape(srcB64)
 	return serializer.Response{
 		Code: 0,
-		Data: viewerBase.String(),
+		Data: util.Replace(map[string]string{
+			"{$src}":    srcEncoded,
+			"{$srcB64}": srcB64Encoded,
+		}, model.GetSettingByName("office_preview_service")),
 	}
 }
 

service/explorer/slave.go

@@ -6,10 +6,10 @@ import (
 	"encoding/json"
 	"fmt"
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/fsctx"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/cloudreve/Cloudreve/v3/pkg/task"
 	"github.com/cloudreve/Cloudreve/v3/pkg/task/slavetask"
 	"github.com/gin-gonic/gin"
@@ -153,7 +153,7 @@ func CreateTransferTask(c *gin.Context, req *serializer.SlaveTransferReq) serial
 			MasterID: id.(string),
 		}
 
-		if err := slave.DefaultController.SubmitTask(job.MasterID, job, req.Hash(job.MasterID), func(job interface{}) {
+		if err := cluster.DefaultController.SubmitTask(job.MasterID, job, req.Hash(job.MasterID), func(job interface{}) {
 			task.TaskPoll.Submit(job.(task.Job))
 		}); err != nil {
 			return serializer.Err(serializer.CodeInternalSetting, "任务创建失败", err)

service/node/fabric.go

@@ -3,10 +3,11 @@ package node
 import (
 	"encoding/gob"
 	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"github.com/cloudreve/Cloudreve/v3/pkg/conf"
 	"github.com/cloudreve/Cloudreve/v3/pkg/filesystem/driver/onedrive"
 	"github.com/cloudreve/Cloudreve/v3/pkg/mq"
 	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
-	"github.com/cloudreve/Cloudreve/v3/pkg/slave"
 	"github.com/gin-gonic/gin"
 )
 
@@ -19,7 +20,7 @@ type OneDriveCredentialService struct {
 }
 
 func HandleMasterHeartbeat(req *serializer.NodePingReq) serializer.Response {
-	res, err := slave.DefaultController.HandleHeartBeat(req)
+	res, err := cluster.DefaultController.HandleHeartBeat(req)
 	if err != nil {
 		return serializer.Err(serializer.CodeInternalSetting, "Cannot initialize slave controller", err)
 	}
@@ -54,7 +55,7 @@ func (s *OneDriveCredentialService) Get(c *gin.Context) serializer.Response {
 		return serializer.Err(serializer.CodeInternalSetting, "Cannot initialize OneDrive client", err)
 	}
 
-	if err := client.UpdateCredential(c); err != nil {
+	if err := client.UpdateCredential(c, conf.SystemConfig.Mode == "slave"); err != nil {
 		return serializer.Err(serializer.CodeInternalSetting, "Cannot refresh OneDrive credential", err)
 	}