Refactor: move slave pkg inside of cluster

Test: middleware for node communication

bootstrap/init.go

@@ -9,7 +9,6 @@ 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/task"
 	"github.com/gin-gonic/gin"
 )
@@ -78,7 +77,7 @@ func Init(path string) {
 		{
 			"slave",
 			func() {
-				slave.Init()
+				cluster.InitController()
 			},
 		},
 		{

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,14 +792,13 @@ 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"},
 		}
 		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"
 )
@@ -19,11 +18,11 @@ func MasterMetadata() gin.HandlerFunc {
 }
 
 // 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,7 +39,7 @@ 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)
 		if err != nil {
@@ -49,7 +48,7 @@ func SlaveRPCSignRequired() gin.HandlerFunc {
 			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,80 @@
+package middleware
+
+import (
+	model "github.com/cloudreve/Cloudreve/v3/models"
+	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
+	"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-Site-Id":           {"expectedSiteID"},
+		"X-Site-Url":          {"expectedSiteURL"},
+		"X-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-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-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-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)
+
+}

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"`
 
 	// 关联模型

pkg/slave/slave.go → 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,

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/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 {
@@ -98,8 +106,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/filesystem/driver/onedrive/oauth.go

@@ -3,6 +3,7 @@ package onedrive
 import (
 	"context"
 	"encoding/json"
+	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
 	"io/ioutil"
 	"net/http"
 	"net/url"
@@ -12,7 +13,6 @@ import (
 	"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"
 )
 
@@ -179,7 +179,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 := cluster.DefaultController.GetOneDriveToken(client.Policy.MasterID, client.Policy.ID)
 	if err != nil {
 		return err
 	}

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,7 +68,7 @@ func (job *TransferTask) SetErrorMsg(msg string, err error) {
 		},
 	}
 
-	if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
+	if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), notifyMsg); err != nil {
 		util.Log().Warning("无法发送转存失败通知到从机, ", 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,7 +131,7 @@ func (job *TransferTask) Do() {
 		Content:     serializer.SlaveTransferResult{},
 	}
 
-	if err := slave.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
+	if err := cluster.DefaultController.SendNotification(job.MasterID, job.Req.Hash(job.MasterID), msg); err != nil {
 		util.Log().Warning("无法发送转存成功通知到从机, ", err)
 	}
 }

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)