Mellaris/ruleset/builtins/geo/geo_matcher.go

package geo

import (
	"container/list"
	"net"
	"sort"
	"strings"
	"sync"
)

const (
	geoSiteResultCacheSize    = 1 << 16
	geoSiteSetResultCacheSize = 1 << 16
)

type GeoMatcher struct {
	geoLoader       GeoLoader
	geoSiteMatcher  map[string]hostMatcher
	siteMatcherLock sync.RWMutex
	geoSiteSets     map[string][]hostMatcher
	siteSetLock     sync.RWMutex
	geoIpMatcher    map[string]hostMatcher
	ipMatcherLock   sync.RWMutex
	geoSiteResult   *boolLRUCache
	geoSiteSetCache *boolLRUCache
}

func NewGeoMatcher(geoSiteFilename, geoIpFilename string) *GeoMatcher {
	return &GeoMatcher{
		geoLoader:       NewDefaultGeoLoader(geoSiteFilename, geoIpFilename),
		geoSiteMatcher:  make(map[string]hostMatcher),
		geoSiteSets:     make(map[string][]hostMatcher),
		geoIpMatcher:    make(map[string]hostMatcher),
		geoSiteResult:   newBoolLRUCache(geoSiteResultCacheSize),
		geoSiteSetCache: newBoolLRUCache(geoSiteSetResultCacheSize),
	}
}

func (g *GeoMatcher) MatchGeoIp(ip, condition string) bool {
	matcher, ok := g.getOrCreateGeoIPMatcher(condition)
	if !ok || matcher == nil {
		return false
	}
	parseIp := net.ParseIP(ip)
	if parseIp == nil {
		return false
	}
	ipv4 := parseIp.To4()
	if ipv4 != nil {
		return matcher.Match(HostInfo{IPv4: ipv4})
	}
	ipv6 := parseIp.To16()
	if ipv6 != nil {
		return matcher.Match(HostInfo{IPv6: ipv6})
	}
	return false
}

func (g *GeoMatcher) MatchGeoSite(site, condition string) bool {
	conditionKey := strings.TrimSpace(strings.ToLower(condition))
	if conditionKey == "" {
		return false
	}
	cacheKey := site + "\x1f" + conditionKey
	if v, ok := g.geoSiteResult.Get(cacheKey); ok {
		return v
	}
	matcher, ok := g.getOrCreateGeoSiteMatcher(condition)
	if !ok || matcher == nil {
		return false
	}
	result := matcher.Match(HostInfo{Name: site})
	g.geoSiteResult.Set(cacheKey, result)
	return result
}

func (g *GeoMatcher) MatchGeoSiteSet(site string, set *SiteConditionSet) bool {
	if set == nil {
		return false
	}
	conditions := normalizeGeoSiteSetConditions(set.Conditions)
	if len(conditions) == 0 {
		return false
	}
	key := strings.Join(conditions, "\x1f")
	cacheKey := site + "\x1e" + key
	if v, ok := g.geoSiteSetCache.Get(cacheKey); ok {
		return v
	}

	g.siteSetLock.RLock()
	matchers, ok := g.geoSiteSets[key]
	g.siteSetLock.RUnlock()
	if !ok {
		compiled := make([]hostMatcher, 0, len(conditions))
		for _, condition := range conditions {
			m, ok := g.getOrCreateGeoSiteMatcher(condition)
			if ok && m != nil {
				compiled = append(compiled, m)
			}
		}
		g.siteSetLock.Lock()
		if existing, exists := g.geoSiteSets[key]; exists {
			matchers = existing
		} else {
			g.geoSiteSets[key] = compiled
			matchers = compiled
		}
		g.siteSetLock.Unlock()
	}
	if len(matchers) == 0 {
		return false
	}
	host := HostInfo{Name: site}
	for _, matcher := range matchers {
		if matcher.Match(host) {
			g.geoSiteSetCache.Set(cacheKey, true)
			return true
		}
	}
	g.geoSiteSetCache.Set(cacheKey, false)
	return false
}

func (g *GeoMatcher) LoadGeoSite() error {
	_, err := g.geoLoader.LoadGeoSite()
	return err
}

func (g *GeoMatcher) LoadGeoIP() error {
	_, err := g.geoLoader.LoadGeoIP()
	return err
}

func parseGeoSiteName(s string) (string, []string) {
	parts := strings.Split(s, "@")
	base := strings.TrimSpace(parts[0])
	attrs := parts[1:]
	for i := range attrs {
		attrs[i] = strings.TrimSpace(attrs[i])
	}
	return base, attrs
}

func (g *GeoMatcher) getOrCreateGeoSiteMatcher(condition string) (hostMatcher, bool) {
	condition = strings.TrimSpace(strings.ToLower(condition))
	if condition == "" {
		return nil, false
	}
	g.siteMatcherLock.RLock()
	matcher, ok := g.geoSiteMatcher[condition]
	g.siteMatcherLock.RUnlock()
	if ok {
		return matcher, true
	}

	name, attrs := parseGeoSiteName(condition)
	if len(name) == 0 {
		return nil, false
	}
	gMap, err := g.geoLoader.LoadGeoSite()
	if err != nil {
		return nil, false
	}
	list, ok := gMap[name]
	if !ok || list == nil {
		return nil, false
	}
	matcher, err = newGeositeMatcher(list, attrs)
	if err != nil {
		return nil, false
	}
	g.siteMatcherLock.Lock()
	if existing, exists := g.geoSiteMatcher[condition]; exists {
		matcher = existing
	} else {
		g.geoSiteMatcher[condition] = matcher
	}
	g.siteMatcherLock.Unlock()
	return matcher, true
}

func (g *GeoMatcher) getOrCreateGeoIPMatcher(condition string) (hostMatcher, bool) {
	condition = strings.TrimSpace(strings.ToLower(condition))
	if condition == "" {
		return nil, false
	}
	g.ipMatcherLock.RLock()
	matcher, ok := g.geoIpMatcher[condition]
	g.ipMatcherLock.RUnlock()
	if ok {
		return matcher, true
	}

	gMap, err := g.geoLoader.LoadGeoIP()
	if err != nil {
		return nil, false
	}
	list, ok := gMap[condition]
	if !ok || list == nil {
		return nil, false
	}
	matcher, err = newGeoIPMatcher(list)
	if err != nil {
		return nil, false
	}
	g.ipMatcherLock.Lock()
	if existing, exists := g.geoIpMatcher[condition]; exists {
		matcher = existing
	} else {
		g.geoIpMatcher[condition] = matcher
	}
	g.ipMatcherLock.Unlock()
	return matcher, true
}

func normalizeGeoSiteSetConditions(in []string) []string {
	if len(in) == 0 {
		return nil
	}
	out := make([]string, 0, len(in))
	seen := make(map[string]struct{}, len(in))
	for _, v := range in {
		s := strings.TrimSpace(strings.ToLower(v))
		if s == "" {
			continue
		}
		if _, ok := seen[s]; ok {
			continue
		}
		seen[s] = struct{}{}
		out = append(out, s)
	}
	sort.Strings(out)
	return out
}

type boolLRUCache struct {
	mu    sync.Mutex
	cap   int
	ll    *list.List
	items map[string]*list.Element
}

type boolCacheEntry struct {
	key   string
	value bool
}

func newBoolLRUCache(capacity int) *boolLRUCache {
	if capacity <= 0 {
		capacity = 1
	}
	return &boolLRUCache{
		cap:   capacity,
		ll:    list.New(),
		items: make(map[string]*list.Element, capacity),
	}
}

func (c *boolLRUCache) Get(key string) (bool, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if ele, ok := c.items[key]; ok {
		c.ll.MoveToFront(ele)
		entry := ele.Value.(boolCacheEntry)
		return entry.value, true
	}
	return false, false
}

func (c *boolLRUCache) Set(key string, value bool) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if ele, ok := c.items[key]; ok {
		ele.Value = boolCacheEntry{key: key, value: value}
		c.ll.MoveToFront(ele)
		return
	}
	ele := c.ll.PushFront(boolCacheEntry{key: key, value: value})
	c.items[key] = ele
	if c.ll.Len() <= c.cap {
		return
	}
	back := c.ll.Back()
	if back == nil {
		return
	}
	entry := back.Value.(boolCacheEntry)
	delete(c.items, entry.key)
	c.ll.Remove(back)
}