From 57648c145c59dd621fbae0a8f059f86599b0c865 Mon Sep 17 00:00:00 2001
From: Darien Raymond <admin@v2ray.com>
Date: Mon, 6 Nov 2017 22:30:56 +0100
Subject: [PATCH] cachable domain matcher: step 2

---
 app/router/condition.go      | 100 +++++++++++++++++++++++++++++++----
 app/router/condition_test.go |  55 +++++++++++++++++++
 2 files changed, 146 insertions(+), 9 deletions(-)

diff --git a/app/router/condition.go b/app/router/condition.go
index 749879289..d370f691c 100644
--- a/app/router/condition.go
+++ b/app/router/condition.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"regexp"
 	"strings"
+	"sync"
+	"time"
 
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
@@ -64,13 +66,22 @@ func (v *AnyCondition) Len() int {
 	return len(*v)
 }
 
+type timedResult struct {
+	timestamp time.Time
+	result    bool
+}
+
 type CachableDomainMatcher struct {
+	sync.Mutex
 	matchers []domainMatcher
+	cache    map[string]timedResult
+	lastScan time.Time
 }
 
 func NewCachableDomainMatcher() *CachableDomainMatcher {
 	return &CachableDomainMatcher{
 		matchers: make([]domainMatcher, 0, 64),
+		cache:    make(map[string]timedResult, 512),
 	}
 }
 
@@ -92,6 +103,85 @@ func (m *CachableDomainMatcher) Add(domain *Domain) error {
 	return nil
 }
 
+func (m *CachableDomainMatcher) applyInternal(domain string) bool {
+	for _, matcher := range m.matchers {
+		if matcher.Apply(domain) {
+			return true
+		}
+	}
+
+	return false
+}
+
+type cacheResult int
+
+const (
+	cacheMiss cacheResult = iota
+	cacheHitTrue
+	cacheHitFalse
+)
+
+func (m *CachableDomainMatcher) findInCache(domain string) cacheResult {
+	m.Lock()
+	defer m.Unlock()
+
+	r, f := m.cache[domain]
+	if !f {
+		return cacheMiss
+	}
+	r.timestamp = time.Now()
+	m.cache[domain] = r
+
+	if r.result {
+		return cacheHitTrue
+	}
+	return cacheHitFalse
+}
+
+func (m *CachableDomainMatcher) ApplyDomain(domain string) bool {
+	if len(m.matchers) < 64 {
+		return m.applyInternal(domain)
+	}
+
+	cr := m.findInCache(domain)
+
+	if cr == cacheHitTrue {
+		return true
+	}
+
+	if cr == cacheHitFalse {
+		return false
+	}
+
+	r := m.applyInternal(domain)
+	m.Lock()
+	defer m.Unlock()
+
+	m.cache[domain] = timedResult{
+		result:    r,
+		timestamp: time.Now(),
+	}
+
+	now := time.Now()
+	if len(m.cache) > 256 && now.Sub(m.lastScan)/time.Second > 5 {
+		remove := make([]string, 0, 128)
+
+		now := time.Now()
+
+		for k, v := range m.cache {
+			if now.Sub(v.timestamp)/time.Second > 60 {
+				remove = append(remove, k)
+			}
+		}
+		for _, v := range remove {
+			delete(m.cache, v)
+		}
+		m.lastScan = now
+	}
+
+	return r
+}
+
 func (m *CachableDomainMatcher) Apply(ctx context.Context) bool {
 	dest, ok := proxy.TargetFromContext(ctx)
 	if !ok {
@@ -101,15 +191,7 @@ func (m *CachableDomainMatcher) Apply(ctx context.Context) bool {
 	if !dest.Address.Family().IsDomain() {
 		return false
 	}
-	domain := dest.Address.Domain()
-
-	for _, matcher := range m.matchers {
-		if matcher.Apply(domain) {
-			return true
-		}
-	}
-
-	return false
+	return m.ApplyDomain(dest.Address.Domain())
 }
 
 type domainMatcher interface {
diff --git a/app/router/condition_test.go b/app/router/condition_test.go
index 2e35b7d7d..eaeac8861 100644
--- a/app/router/condition_test.go
+++ b/app/router/condition_test.go
@@ -2,13 +2,22 @@ package router_test
 
 import (
 	"context"
+	"os"
+	"path/filepath"
+	"strconv"
 	"testing"
+	"time"
 
+	proto "github.com/golang/protobuf/proto"
 	. "v2ray.com/core/app/router"
+	"v2ray.com/core/common"
+	"v2ray.com/core/common/errors"
 	"v2ray.com/core/common/net"
+	"v2ray.com/core/common/platform"
 	"v2ray.com/core/common/protocol"
 	"v2ray.com/core/proxy"
 	. "v2ray.com/ext/assert"
+	"v2ray.com/ext/sysio"
 )
 
 func TestSubDomainMatcher(t *testing.T) {
@@ -179,3 +188,49 @@ func TestRoutingRule(t *testing.T) {
 		}
 	}
 }
+
+func loadGeoSite(country string) ([]*Domain, error) {
+	geositeBytes, err := sysio.ReadAsset("geosite.dat")
+	if err != nil {
+		return nil, err
+	}
+	var geositeList GeoSiteList
+	if err := proto.Unmarshal(geositeBytes, &geositeList); err != nil {
+		return nil, err
+	}
+
+	for _, site := range geositeList.Entry {
+		if site.CountryCode == country {
+			return site.Domain, nil
+		}
+	}
+
+	return nil, errors.New("country not found: " + country)
+}
+
+func TestChinaSites(t *testing.T) {
+	assert := With(t)
+
+	common.Must(sysio.CopyFile(platform.GetAssetLocation("geosite.dat"), filepath.Join(os.Getenv("GOPATH"), "src", "v2ray.com", "core", "tools", "release", "config", "geosite.dat")))
+
+	domains, err := loadGeoSite("CN")
+	assert(err, IsNil)
+
+	matcher := NewCachableDomainMatcher()
+	for _, d := range domains {
+		assert(matcher.Add(d), IsNil)
+	}
+
+	assert(matcher.ApplyDomain("163.com"), IsTrue)
+	assert(matcher.ApplyDomain("163.com"), IsTrue)
+	assert(matcher.ApplyDomain("164.com"), IsFalse)
+	assert(matcher.ApplyDomain("164.com"), IsFalse)
+
+	for i := 0; i < 1024; i++ {
+		assert(matcher.ApplyDomain(strconv.Itoa(i)+".not-exists.com"), IsFalse)
+	}
+	time.Sleep(time.Second * 10)
+	for i := 0; i < 1024; i++ {
+		assert(matcher.ApplyDomain(strconv.Itoa(i)+".not-exists2.com"), IsFalse)
+	}
+}