sub header enhancement

sub/subHandler.go

@@ -21,7 +21,7 @@ func NewSubHandler(g *gin.RouterGroup) {
 
 func (s *SubHandler) initRouter(g *gin.RouterGroup) {
 	g.GET("/:subid", s.subs)
-	g.HEAD("/:subid", s.subs)
+	g.HEAD("/:subid", s.subHeaders)
 }
 
 func (s *SubHandler) subs(c *gin.Context) {
@@ -50,16 +50,29 @@ func (s *SubHandler) subs(c *gin.Context) {
 			return
 		}
 	}
-	// Add headers
-	c.Writer.Header().Set("Subscription-Userinfo", headers[0])
-	c.Writer.Header().Set("Profile-Update-Interval", headers[1])
-	c.Writer.Header().Set("Profile-Title", headers[2])
 
-	// For HEAD requests
-	if c.Request.Method == "HEAD" {
-		c.Status(200)
-		return
-	}
+	s.addHeaders(c, headers)
 
 	c.String(200, *result)
 }
+
+func (s *SubHandler) subHeaders(c *gin.Context) {
+	subId := c.Param("subid")
+	client, err := s.SubService.getClientBySubId(subId)
+	if err != nil {
+		logger.Error(err)
+		c.String(400, "Error!")
+		return
+	}
+
+	headers := s.SubService.getClientHeaders(client)
+	s.addHeaders(c, headers)
+
+	c.Status(200)
+}
+
+func (s *SubHandler) addHeaders(c *gin.Context, headers []string) {
+	c.Writer.Header().Set("Subscription-Userinfo", headers[0])
+	c.Writer.Header().Set("Profile-Update-Interval", headers[1])
+	c.Writer.Header().Set("Profile-Title", headers[2])
+}

sub/subService.go

@@ -20,9 +20,7 @@ type SubService struct {
 func (s *SubService) GetSubs(subId string) (*string, []string, error) {
 	var err error
 
-	db := database.GetDB()
-	client := &model.Client{}
-	err = db.Model(model.Client{}).Where("enable = true and name = ?", subId).First(client).Error
+	client, err := s.getClientBySubId(subId)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -36,8 +34,7 @@ func (s *SubService) GetSubs(subId string) (*string, []string, error) {
 	linksArray := s.LinkService.GetLinks(&client.Links, "all", clientInfo)
 	result := strings.Join(linksArray, "\n")
 
-	updateInterval, _ := s.SettingService.GetSubUpdates()
-	headers := util.GetHeaders(client, updateInterval)
+	headers := s.getClientHeaders(client)
 
 	subEncode, _ := s.SettingService.GetSubEncode()
 	if subEncode {
@@ -47,6 +44,21 @@ func (s *SubService) GetSubs(subId string) (*string, []string, error) {
 	return &result, headers, nil
 }
 
+func (j *SubService) getClientBySubId(subId string) (*model.Client, error) {
+	db := database.GetDB()
+	client := &model.Client{}
+	err := db.Model(model.Client{}).Where("enable = true and name = ?", subId).First(client).Error
+	if err != nil {
+		return nil, err
+	}
+	return client, nil
+}
+
+func (s *SubService) getClientHeaders(client *model.Client) []string {
+	updateInterval, _ := s.SettingService.GetSubUpdates()
+	return util.GetHeaders(client, updateInterval)
+}
+
 func (s *SubService) getClientInfo(c *model.Client) string {
 	now := time.Now().Unix()