Browse Source

optimize multi-geoip matcher

Darien Raymond 7 years ago
parent
commit
58221ebae1
2 changed files with 68 additions and 14 deletions
  1. 15 14
      app/router/condition.go
  2. 53 0
      app/router/condition_test.go

+ 15 - 14
app/router/condition.go

@@ -113,7 +113,7 @@ func targetFromContent(ctx context.Context) net.Destination {
 
 type MultiGeoIPMatcher struct {
 	matchers []*GeoIPMatcher
-	onSource bool
+	destFunc func(context.Context) net.Destination
 }
 
 func NewMultiGeoIPMatcher(geoips []*GeoIP, onSource bool) (*MultiGeoIPMatcher, error) {
@@ -126,30 +126,31 @@ func NewMultiGeoIPMatcher(geoips []*GeoIP, onSource bool) (*MultiGeoIPMatcher, e
 		matchers = append(matchers, matcher)
 	}
 
+	var destFunc func(context.Context) net.Destination
+	if onSource {
+		destFunc = sourceFromContext
+	} else {
+		destFunc = targetFromContent
+	}
+
 	return &MultiGeoIPMatcher{
 		matchers: matchers,
-		onSource: onSource,
+		destFunc: destFunc,
 	}, nil
 }
 
 func (m *MultiGeoIPMatcher) Apply(ctx context.Context) bool {
 	ips := make([]net.IP, 0, 4)
-	if resolver, ok := ResolvedIPsFromContext(ctx); ok {
-		resolvedIPs := resolver.Resolve()
-		for _, rip := range resolvedIPs {
-			ips = append(ips, rip.IP())
-		}
-	}
 
-	var dest net.Destination
-	if m.onSource {
-		dest = sourceFromContext(ctx)
-	} else {
-		dest = targetFromContent(ctx)
-	}
+	dest := m.destFunc(ctx)
 
 	if dest.IsValid() && (dest.Address.Family().IsIPv4() || dest.Address.Family().IsIPv6()) {
 		ips = append(ips, dest.Address.IP())
+	} else if resolver, ok := ResolvedIPsFromContext(ctx); ok {
+		resolvedIPs := resolver.Resolve()
+		for _, rip := range resolvedIPs {
+			ips = append(ips, rip.IP())
+		}
 	}
 
 	for _, ip := range ips {

+ 53 - 0
app/router/condition_test.go

@@ -285,3 +285,56 @@ func TestChinaSites(t *testing.T) {
 		assert(matcher.ApplyDomain(strconv.Itoa(i)+".not-exists.com"), IsFalse)
 	}
 }
+
+func BenchmarkMultiGeoIPMatcher(b *testing.B) {
+	common.Must(sysio.CopyFile(platform.GetAssetLocation("geoip.dat"), filepath.Join(os.Getenv("GOPATH"), "src", "v2ray.com", "core", "release", "config", "geoip.dat")))
+
+	var geoips []*GeoIP
+
+	{
+		ips, err := loadGeoIP("CN")
+		common.Must(err)
+		geoips = append(geoips, &GeoIP{
+			CountryCode: "CN",
+			Cidr:        ips,
+		})
+	}
+
+	{
+		ips, err := loadGeoIP("JP")
+		common.Must(err)
+		geoips = append(geoips, &GeoIP{
+			CountryCode: "JP",
+			Cidr:        ips,
+		})
+	}
+
+	{
+		ips, err := loadGeoIP("CA")
+		common.Must(err)
+		geoips = append(geoips, &GeoIP{
+			CountryCode: "CA",
+			Cidr:        ips,
+		})
+	}
+
+	{
+		ips, err := loadGeoIP("US")
+		common.Must(err)
+		geoips = append(geoips, &GeoIP{
+			CountryCode: "US",
+			Cidr:        ips,
+		})
+	}
+
+	matcher, err := NewMultiGeoIPMatcher(geoips, false)
+	common.Must(err)
+
+	ctx := withOutbound(&session.Outbound{Target: net.TCPDestination(net.ParseAddress("8.8.8.8"), 80)})
+
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_ = matcher.Apply(ctx)
+	}
+}