OpenGFW/ruleset/builtins/geo/geo_matcher.go
Rinka f07a38bc47
Add GeoIP and GeoSite support for expr (#38)
* feat: copy something from hysteria/extras/outbounds/acl

* feat: add geoip and geosite support for expr

* refactor: geo matcher

* fix: typo

* refactor: geo matcher

* feat: expose config options to specify local geoip/geosite db files

* refactor: engine.Config should not contains geo

* feat: make geosite and geoip lazy downloaded

* chore: minor code improvement

* docs: add geoip/geosite usage

---------

Co-authored-by: Toby <tobyxdd@gmail.com>
2024-01-30 17:30:35 -08:00

116 lines
2.4 KiB
Go

package geo
import (
"net"
"strings"
"sync"
)
type GeoMatcher struct {
geoLoader GeoLoader
geoSiteMatcher map[string]hostMatcher
siteMatcherLock sync.Mutex
geoIpMatcher map[string]hostMatcher
ipMatcherLock sync.Mutex
}
func NewGeoMatcher(geoSiteFilename, geoIpFilename string) (*GeoMatcher, error) {
geoLoader := NewDefaultGeoLoader(geoSiteFilename, geoIpFilename)
return &GeoMatcher{
geoLoader: geoLoader,
geoSiteMatcher: make(map[string]hostMatcher),
geoIpMatcher: make(map[string]hostMatcher),
}, nil
}
func (g *GeoMatcher) MatchGeoIp(ip, condition string) bool {
g.ipMatcherLock.Lock()
defer g.ipMatcherLock.Unlock()
matcher, ok := g.geoIpMatcher[condition]
if !ok {
// GeoIP matcher
condition = strings.ToLower(condition)
country := condition
if len(country) == 0 {
return false
}
gMap, err := g.geoLoader.LoadGeoIP()
if err != nil {
return false
}
list, ok := gMap[country]
if !ok || list == nil {
return false
}
matcher, err = newGeoIPMatcher(list)
if err != nil {
return false
}
g.geoIpMatcher[condition] = matcher
}
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 {
g.siteMatcherLock.Lock()
defer g.siteMatcherLock.Unlock()
matcher, ok := g.geoSiteMatcher[condition]
if !ok {
// MatchGeoSite matcher
condition = strings.ToLower(condition)
name, attrs := parseGeoSiteName(condition)
if len(name) == 0 {
return false
}
gMap, err := g.geoLoader.LoadGeoSite()
if err != nil {
return false
}
list, ok := gMap[name]
if !ok || list == nil {
return false
}
matcher, err = newGeositeMatcher(list, attrs)
if err != nil {
return false
}
g.geoSiteMatcher[condition] = matcher
}
return matcher.Match(HostInfo{Name: site})
}
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
}