3x-ui/sub/subController.go

package sub

import (
	"encoding/base64"
	"net"
	"strconv"
	"strings"

	"x-ui/util/common"

	"github.com/gin-gonic/gin"
)

type SUBController struct {
	subTitle       string
	subPath        string
	subJsonPath    string
	subEncrypt     bool
	updateInterval string

	subService     *SubService
	subJsonService *SubJsonService
}

func NewSUBController(
	g *gin.RouterGroup,
	subPath string,
	jsonPath string,
	encrypt bool,
	showInfo bool,
	rModel string,
	update string,
	jsonFragment string,
	jsonNoise string,
	jsonMux string,
	jsonRules string,
	subTitle string,
) *SUBController {
	sub := NewSubService(showInfo, rModel)
	a := &SUBController{
		subTitle:       subTitle,
		subPath:        subPath,
		subJsonPath:    jsonPath,
		subEncrypt:     encrypt,
		updateInterval: update,

		subService:     sub,
		subJsonService: NewSubJsonService(jsonFragment, jsonNoise, jsonMux, jsonRules, sub),
	}
	a.initRouter(g)
	return a
}

func (a *SUBController) initRouter(g *gin.RouterGroup) {
	gLink := g.Group(a.subPath)
	gJson := g.Group(a.subJsonPath)

	gLink.GET(":subid", a.subs)
	gJson.GET(":subid", a.subJsons)
}

func (a *SUBController) subs(c *gin.Context) {
	subId := c.Param("subid")
	var host string
	if h, err := getHostFromXFH(c.GetHeader("X-Forwarded-Host")); err == nil {
		host = h
	}
	if host == "" {
		host = c.GetHeader("X-Real-IP")
	}
	if host == "" {
		var err error
		host, _, err = net.SplitHostPort(c.Request.Host)
		if err != nil {
			host = c.Request.Host
		}
	}
	subs, header, lastOnline, err := a.subService.GetSubs(subId, host)
	if err != nil || len(subs) == 0 {
		c.String(400, "Error!")
	} else {
		result := ""
		for _, sub := range subs {
			result += sub + "\n"
		}

		// Add headers
		c.Writer.Header().Set("Subscription-Userinfo", header)
		c.Writer.Header().Set("Profile-Update-Interval", a.updateInterval)
		c.Writer.Header().Set("Profile-Title", "base64:"+base64.StdEncoding.EncodeToString([]byte(a.subTitle)))
		// Also include whole subscription content in base64 as requested
		c.Writer.Header().Set("Subscription-Content-Base64", base64.StdEncoding.EncodeToString([]byte(result)))

		// If the request expects HTML (e.g., browser) or explicitly asked (?html=1 or ?view=html), render the info page here
		accept := c.GetHeader("Accept")
		if strings.Contains(strings.ToLower(accept), "text/html") || c.Query("html") == "1" || strings.EqualFold(c.Query("view"), "html") {
			// Determine scheme
			scheme := "http"
			if c.Request.TLS != nil || strings.EqualFold(c.GetHeader("X-Forwarded-Proto"), "https") {
				scheme = "https"
			}

			// Parse header values
			var uploadByte, downloadByte, totalByte, expire int64
			parts := strings.Split(header, ";")
			for _, p := range parts {
				kv := strings.Split(strings.TrimSpace(p), "=")
				if len(kv) != 2 {
					continue
				}
				key := strings.ToLower(strings.TrimSpace(kv[0]))
				val := strings.TrimSpace(kv[1])
				switch key {
				case "upload":
					if v, err := parseInt64(val); err == nil {
						uploadByte = v
					}
				case "download":
					if v, err := parseInt64(val); err == nil {
						downloadByte = v
					}
				case "total":
					if v, err := parseInt64(val); err == nil {
						totalByte = v
					}
				case "expire":
					if v, err := parseInt64(val); err == nil {
						expire = v
					}
				}
			}

			download := common.FormatTraffic(downloadByte)
			upload := common.FormatTraffic(uploadByte)
			total := "∞"
			used := common.FormatTraffic(uploadByte + downloadByte)
			remained := ""
			if totalByte > 0 {
				total = common.FormatTraffic(totalByte)
				left := max(totalByte-(uploadByte+downloadByte), 0)
				remained = common.FormatTraffic(left)
			}

			// Build sub URL
			subURL := scheme + "://" + host + strings.TrimRight(a.subPath, "/") + "/" + subId
			if strings.HasSuffix(a.subPath, "/") {
				subURL = scheme + "://" + host + a.subPath + subId
			}

			basePath := "/"
			hostHeader := c.GetHeader("X-Forwarded-Host")
			if hostHeader == "" {
				hostHeader = c.GetHeader("X-Real-IP")
			}
			if hostHeader == "" {
				hostHeader = host
			}
			c.HTML(200, "subscription.html", gin.H{
				"title":        "subscription.title",
				"host":         hostHeader,
				"base_path":    basePath,
				"sId":          subId,
				"download":     download,
				"upload":       upload,
				"total":        total,
				"used":         used,
				"remained":     remained,
				"expire":       expire,
				"lastOnline":   lastOnline,
				"datepicker":   a.subService.datepicker,
				"downloadByte": downloadByte,
				"uploadByte":   uploadByte,
				"totalByte":    totalByte,
				"subUrl":       subURL,
				"result":       subs,
			})
			return
		}

		if a.subEncrypt {
			c.String(200, base64.StdEncoding.EncodeToString([]byte(result)))
		} else {
			c.String(200, result)
		}
	}
}

func (a *SUBController) subJsons(c *gin.Context) {
	subId := c.Param("subid")
	var host string
	if h, err := getHostFromXFH(c.GetHeader("X-Forwarded-Host")); err == nil {
		host = h
	}
	if host == "" {
		host = c.GetHeader("X-Real-IP")
	}
	if host == "" {
		var err error
		host, _, err = net.SplitHostPort(c.Request.Host)
		if err != nil {
			host = c.Request.Host
		}
	}
	jsonSub, header, err := a.subJsonService.GetJson(subId, host)
	if err != nil || len(jsonSub) == 0 {
		c.String(400, "Error!")
	} else {

		// Add headers
		c.Writer.Header().Set("Subscription-Userinfo", header)
		c.Writer.Header().Set("Profile-Update-Interval", a.updateInterval)
		c.Writer.Header().Set("Profile-Title", "base64:"+base64.StdEncoding.EncodeToString([]byte(a.subTitle)))

		c.String(200, jsonSub)
	}
}

func getHostFromXFH(s string) (string, error) {
	if strings.Contains(s, ":") {
		realHost, _, err := net.SplitHostPort(s)
		if err != nil {
			return "", err
		}
		return realHost, nil
	}
	return s, nil
}

func parseInt64(s string) (int64, error) {
	var n int64
	var err error
	// handle potential quotes
	s = strings.Trim(s, "\"'")
	n, err = strconv.ParseInt(s, 10, 64)
	return n, err
}