new-pai-linghe/setting/secondary_auto_group.go

package setting

import (
	"net"
	"strings"
	"sync"

	"github.com/QuantumNous/new-api/common"
)

// SecondaryAutoGroup 二级 auto 定义（高价池 / 低价池等）。
//
// 设计要点：
//   - 对用户隐藏：名字不进 GroupRatio / UserUsableGroups，因此不会出现在用户分组下拉，
//     用户只能通过顶层 "auto" 分组 + 来源域名路由间接命中。
//   - 复用顶层 auto 的选渠道 / 跨子分组重试 / 熔断能力（见 service/channel_select.go）。
//   - Groups 为有序子分组列表（顺序=优先级），每个子分组须已在 GroupRatio 注册并挂渠道。
type SecondaryAutoGroup struct {
	DisplayName string   `json:"display_name"`
	Groups      []string `json:"groups"`
}

var (
	// secondaryAutoGroups: 二级 auto 名 -> 定义
	secondaryAutoGroups = map[string]SecondaryAutoGroup{}
	// domainAutoRouting: 规范化域名 -> 二级 auto 名
	domainAutoRouting = map[string]string{}
	autoRoutingMu     sync.RWMutex
)

/* ----------------------------- SecondaryAutoGroups ----------------------------- */

func GetSecondaryAutoGroups() map[string]SecondaryAutoGroup {
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	out := make(map[string]SecondaryAutoGroup, len(secondaryAutoGroups))
	for k, v := range secondaryAutoGroups {
		out[k] = v
	}
	return out
}

func GetSecondaryAutoGroup(name string) (SecondaryAutoGroup, bool) {
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	g, ok := secondaryAutoGroups[name]
	return g, ok
}

func IsSecondaryAutoGroup(name string) bool {
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	_, ok := secondaryAutoGroups[name]
	return ok
}

// IsAutoTypeGroup 顶层 auto 或任一二级 auto。
func IsAutoTypeGroup(name string) bool {
	if name == "auto" {
		return true
	}
	return IsSecondaryAutoGroup(name)
}

// GetAutoGroupSubGroups 返回某个 auto 型分组的子分组列表（有序=优先级）。
// 顶层 "auto" -> 全局 AutoGroups；二级 auto -> 其 Groups 配置。
func GetAutoGroupSubGroups(name string) []string {
	if name == "auto" {
		return GetAutoGroups()
	}
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	if g, ok := secondaryAutoGroups[name]; ok {
		subs := make([]string, len(g.Groups))
		copy(subs, g.Groups)
		return subs
	}
	return []string{}
}

func SecondaryAutoGroups2JsonString() string {
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	b, err := common.Marshal(secondaryAutoGroups)
	if err != nil {
		return "{}"
	}
	return string(b)
}

func UpdateSecondaryAutoGroupsByJsonString(jsonStr string) error {
	m := map[string]SecondaryAutoGroup{}
	if strings.TrimSpace(jsonStr) != "" {
		if err := common.Unmarshal([]byte(jsonStr), &m); err != nil {
			return err
		}
	}
	autoRoutingMu.Lock()
	defer autoRoutingMu.Unlock()
	secondaryAutoGroups = m
	return nil
}

/* ----------------------------- DomainAutoRouting ----------------------------- */

func DomainAutoRouting2JsonString() string {
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	b, err := common.Marshal(domainAutoRouting)
	if err != nil {
		return "{}"
	}
	return string(b)
}

func UpdateDomainAutoRoutingByJsonString(jsonStr string) error {
	raw := map[string]string{}
	if strings.TrimSpace(jsonStr) != "" {
		if err := common.Unmarshal([]byte(jsonStr), &raw); err != nil {
			return err
		}
	}
	// 规范化域名 key（小写、去空白），与查询侧 NormalizeDomain 对齐
	norm := make(map[string]string, len(raw))
	for k, v := range raw {
		key := strings.ToLower(strings.TrimSpace(k))
		if key == "" {
			continue
		}
		norm[key] = strings.TrimSpace(v)
	}
	autoRoutingMu.Lock()
	defer autoRoutingMu.Unlock()
	domainAutoRouting = norm
	return nil
}

// ResolveDomainToSecondaryAuto 按规范化域名解析二级 auto 名。
func ResolveDomainToSecondaryAuto(normalizedDomain string) (string, bool) {
	if normalizedDomain == "" {
		return "", false
	}
	autoRoutingMu.RLock()
	defer autoRoutingMu.RUnlock()
	target, ok := domainAutoRouting[normalizedDomain]
	if !ok || target == "" {
		return "", false
	}
	return target, true
}

// NormalizeDomain 从 host 提取规范化域名：取首段（X-Forwarded-Host 可能多值）、去端口、小写。
func NormalizeDomain(host string) string {
	host = strings.TrimSpace(host)
	if host == "" {
		return ""
	}
	if idx := strings.IndexByte(host, ','); idx != -1 {
		host = strings.TrimSpace(host[:idx])
	}
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	return strings.ToLower(host)
}