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 }