Feat: refine geodata loader (#323)

This commit is contained in:
Loyalsoldier
2021-05-11 18:18:40 +08:00
committed by GitHub
parent a9ce6d4ada
commit 531bc2b1d6
5 changed files with 29 additions and 26 deletions
+10 -10
View File
@@ -11,27 +11,27 @@ import (
"github.com/p4gefau1t/trojan-go/log"
)
type GeoIPCache map[string]*v2router.GeoIP
type geoipCache map[string]*v2router.GeoIP
func (g GeoIPCache) Has(key string) bool {
func (g geoipCache) Has(key string) bool {
return !(g.Get(key) == nil)
}
func (g GeoIPCache) Get(key string) *v2router.GeoIP {
func (g geoipCache) Get(key string) *v2router.GeoIP {
if g == nil {
return nil
}
return g[key]
}
func (g GeoIPCache) Set(key string, value *v2router.GeoIP) {
func (g geoipCache) Set(key string, value *v2router.GeoIP) {
if g == nil {
g = make(map[string]*v2router.GeoIP)
}
g[key] = value
}
func (g GeoIPCache) Unmarshal(filename, code string) (*v2router.GeoIP, error) {
func (g geoipCache) Unmarshal(filename, code string) (*v2router.GeoIP, error) {
asset := common.GetAssetLocation(filename)
idx := strings.ToLower(asset + ":" + code)
if g.Has(idx) {
@@ -77,27 +77,27 @@ func (g GeoIPCache) Unmarshal(filename, code string) (*v2router.GeoIP, error) {
return nil, common.NewError("country code " + code + " not found in " + filename)
}
type GeoSiteCache map[string]*v2router.GeoSite
type geositeCache map[string]*v2router.GeoSite
func (g GeoSiteCache) Has(key string) bool {
func (g geositeCache) Has(key string) bool {
return !(g.Get(key) == nil)
}
func (g GeoSiteCache) Get(key string) *v2router.GeoSite {
func (g geositeCache) Get(key string) *v2router.GeoSite {
if g == nil {
return nil
}
return g[key]
}
func (g GeoSiteCache) Set(key string, value *v2router.GeoSite) {
func (g geositeCache) Set(key string, value *v2router.GeoSite) {
if g == nil {
g = make(map[string]*v2router.GeoSite)
}
g[key] = value
}
func (g GeoSiteCache) Unmarshal(filename, code string) (*v2router.GeoSite, error) {
func (g geositeCache) Unmarshal(filename, code string) (*v2router.GeoSite, error) {
asset := common.GetAssetLocation(filename)
idx := strings.ToLower(asset + ":" + code)
if g.Has(idx) {
+2 -2
View File
@@ -72,7 +72,7 @@ func BenchmarkLoadGeoIP(b *testing.B) {
m1 := runtime.MemStats{}
m2 := runtime.MemStats{}
loader := geodata.GetGeodataLoader()
loader := geodata.NewGeodataLoader()
runtime.ReadMemStats(&m1)
cn, _ := loader.LoadGeoIP("cn")
@@ -89,7 +89,7 @@ func BenchmarkLoadGeoSite(b *testing.B) {
m3 := runtime.MemStats{}
m4 := runtime.MemStats{}
loader := geodata.GetGeodataLoader()
loader := geodata.NewGeodataLoader()
runtime.ReadMemStats(&m3)
cn, _ := loader.LoadGeoSite("cn")
+10
View File
@@ -0,0 +1,10 @@
package geodata
import v2router "github.com/v2fly/v2ray-core/v4/app/router"
type GeodataLoader interface {
LoadIP(filename, country string) ([]*v2router.CIDR, error)
LoadSite(filename, list string) ([]*v2router.Domain, error)
LoadGeoIP(country string) ([]*v2router.CIDR, error)
LoadGeoSite(list string) ([]*v2router.Domain, error)
}
+6 -13
View File
@@ -6,27 +6,20 @@ import (
v2router "github.com/v2fly/v2ray-core/v4/app/router"
)
type geodataLoader interface {
LoadIP(filename, country string) ([]*v2router.CIDR, error)
LoadSite(filename, list string) ([]*v2router.Domain, error)
LoadGeoIP(country string) ([]*v2router.CIDR, error)
LoadGeoSite(list string) ([]*v2router.Domain, error)
type geodataCache struct {
geoipCache
geositeCache
}
func GetGeodataLoader() geodataLoader {
func NewGeodataLoader() GeodataLoader {
return &geodataCache{
make(map[string]*v2router.GeoIP),
make(map[string]*v2router.GeoSite),
}
}
type geodataCache struct {
GeoIPCache
GeoSiteCache
}
func (g *geodataCache) LoadIP(filename, country string) ([]*v2router.CIDR, error) {
geoip, err := g.GeoIPCache.Unmarshal(filename, country)
geoip, err := g.geoipCache.Unmarshal(filename, country)
if err != nil {
return nil, err
}
@@ -35,7 +28,7 @@ func (g *geodataCache) LoadIP(filename, country string) ([]*v2router.CIDR, error
}
func (g *geodataCache) LoadSite(filename, list string) ([]*v2router.Domain, error) {
geosite, err := g.GeoSiteCache.Unmarshal(filename, list)
geosite, err := g.geositeCache.Unmarshal(filename, list)
if err != nil {
return nil, err
}
+1 -1
View File
@@ -312,7 +312,7 @@ func NewClient(ctx context.Context, underlay tunnel.Client) (*Client, error) {
runtime.ReadMemStats(&m1)
geodataLoader := geodata.GetGeodataLoader()
geodataLoader := geodata.NewGeodataLoader()
ipCode := loadCode(cfg, "geoip:")
for _, c := range ipCode {