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) }