github.com/cloudreve/Cloudreve/v3@v3.0.0-20240224133659-3edb00a6484c/middleware/cluster.go (about)

     1  package middleware
     2  
     3  import (
     4  	"github.com/cloudreve/Cloudreve/v3/pkg/auth"
     5  	"github.com/cloudreve/Cloudreve/v3/pkg/cluster"
     6  	"github.com/cloudreve/Cloudreve/v3/pkg/serializer"
     7  	"github.com/gin-gonic/gin"
     8  	"strconv"
     9  )
    10  
    11  // MasterMetadata 解析主机节点发来请求的包含主机节点信息的元数据
    12  func MasterMetadata() gin.HandlerFunc {
    13  	return func(c *gin.Context) {
    14  		c.Set("MasterSiteID", c.GetHeader(auth.CrHeaderPrefix+"Site-Id"))
    15  		c.Set("MasterSiteURL", c.GetHeader(auth.CrHeaderPrefix+"Site-Url"))
    16  		c.Set("MasterVersion", c.GetHeader(auth.CrHeaderPrefix+"Cloudreve-Version"))
    17  		c.Next()
    18  	}
    19  }
    20  
    21  // UseSlaveAria2Instance 从机用于获取对应主机节点的Aria2实例
    22  func UseSlaveAria2Instance(clusterController cluster.Controller) gin.HandlerFunc {
    23  	return func(c *gin.Context) {
    24  		if siteID, exist := c.Get("MasterSiteID"); exist {
    25  			// 获取对应主机节点的从机Aria2实例
    26  			caller, err := clusterController.GetAria2Instance(siteID.(string))
    27  			if err != nil {
    28  				c.JSON(200, serializer.Err(serializer.CodeNotSet, "Failed to get Aria2 instance", err))
    29  				c.Abort()
    30  				return
    31  			}
    32  
    33  			c.Set("MasterAria2Instance", caller)
    34  			c.Next()
    35  			return
    36  		}
    37  
    38  		c.JSON(200, serializer.ParamErr("Unknown master node ID", nil))
    39  		c.Abort()
    40  	}
    41  }
    42  
    43  func SlaveRPCSignRequired(nodePool cluster.Pool) gin.HandlerFunc {
    44  	return func(c *gin.Context) {
    45  		nodeID, err := strconv.ParseUint(c.GetHeader(auth.CrHeaderPrefix+"Node-Id"), 10, 64)
    46  		if err != nil {
    47  			c.JSON(200, serializer.ParamErr("Unknown master node ID", err))
    48  			c.Abort()
    49  			return
    50  		}
    51  
    52  		slaveNode := nodePool.GetNodeByID(uint(nodeID))
    53  		if slaveNode == nil {
    54  			c.JSON(200, serializer.ParamErr("Unknown master node ID", err))
    55  			c.Abort()
    56  			return
    57  		}
    58  
    59  		SignRequired(slaveNode.MasterAuthInstance())(c)
    60  
    61  	}
    62  }