OpenGFW/ruleset/builtins/geo/geo_loader.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

129 lines
2.7 KiB
Go

package geo
import (
"io"
"net/http"
"os"
"time"
"github.com/apernet/OpenGFW/ruleset/builtins/geo/v2geo"
)
const (
geoipFilename = "geoip.dat"
geoipURL = "https://cdn.jsdelivr.net/gh/Loyalsoldier/v2ray-rules-dat@release/geoip.dat"
geositeFilename = "geosite.dat"
geositeURL = "https://cdn.jsdelivr.net/gh/Loyalsoldier/v2ray-rules-dat@release/geosite.dat"
geoDefaultUpdateInterval = 7 * 24 * time.Hour // 7 days
)
var _ GeoLoader = (*V2GeoLoader)(nil)
// V2GeoLoader provides the on-demand GeoIP/MatchGeoSite database
// loading functionality required by the ACL engine.
// Empty filenames = automatic download from built-in URLs.
type V2GeoLoader struct {
GeoIPFilename string
GeoSiteFilename string
UpdateInterval time.Duration
DownloadFunc func(filename, url string)
DownloadErrFunc func(err error)
geoipMap map[string]*v2geo.GeoIP
geositeMap map[string]*v2geo.GeoSite
}
func NewDefaultGeoLoader(geoSiteFilename, geoIpFilename string) *V2GeoLoader {
return &V2GeoLoader{
GeoIPFilename: geoIpFilename,
GeoSiteFilename: geoSiteFilename,
DownloadFunc: func(filename, url string) {},
DownloadErrFunc: func(err error) {},
}
}
func (l *V2GeoLoader) shouldDownload(filename string) bool {
info, err := os.Stat(filename)
if os.IsNotExist(err) {
return true
}
dt := time.Now().Sub(info.ModTime())
if l.UpdateInterval == 0 {
return dt > geoDefaultUpdateInterval
} else {
return dt > l.UpdateInterval
}
}
func (l *V2GeoLoader) download(filename, url string) error {
l.DownloadFunc(filename, url)
resp, err := http.Get(url)
if err != nil {
l.DownloadErrFunc(err)
return err
}
defer resp.Body.Close()
f, err := os.Create(filename)
if err != nil {
l.DownloadErrFunc(err)
return err
}
defer f.Close()
_, err = io.Copy(f, resp.Body)
l.DownloadErrFunc(err)
return err
}
func (l *V2GeoLoader) LoadGeoIP() (map[string]*v2geo.GeoIP, error) {
if l.geoipMap != nil {
return l.geoipMap, nil
}
autoDL := false
filename := l.GeoIPFilename
if filename == "" {
autoDL = true
filename = geoipFilename
}
if autoDL && l.shouldDownload(filename) {
err := l.download(filename, geoipURL)
if err != nil {
return nil, err
}
}
m, err := v2geo.LoadGeoIP(filename)
if err != nil {
return nil, err
}
l.geoipMap = m
return m, nil
}
func (l *V2GeoLoader) LoadGeoSite() (map[string]*v2geo.GeoSite, error) {
if l.geositeMap != nil {
return l.geositeMap, nil
}
autoDL := false
filename := l.GeoSiteFilename
if filename == "" {
autoDL = true
filename = geositeFilename
}
if autoDL && l.shouldDownload(filename) {
err := l.download(filename, geositeURL)
if err != nil {
return nil, err
}
}
m, err := v2geo.LoadGeoSite(filename)
if err != nil {
return nil, err
}
l.geositeMap = m
return m, nil
}