From 02684a868f32ca05af9aa49d61a77fac7580a897 Mon Sep 17 00:00:00 2001
From: i40e <120920147+i40e@users.noreply.github.com>
Date: Sat, 21 Jan 2023 14:40:36 +0800
Subject: [PATCH] feature: geosite-based nameserver policy

---
 dns/filters.go  | 11 +++++++++++
 dns/resolver.go | 47 +++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/dns/filters.go b/dns/filters.go
index 0dbfa317..b51e6402 100644
--- a/dns/filters.go
+++ b/dns/filters.go
@@ -91,6 +91,17 @@ type geoSiteFilter struct {
 	matchers []*router.DomainMatcher
 }
 
+func NewGeoSite(group string) (fallbackDomainFilter, error) {
+	matcher, _, err := geodata.LoadGeoSiteMatcher(group)
+	if err != nil {
+		return nil, err
+	}
+	filter := &geoSiteFilter{
+		matchers: []*router.DomainMatcher{matcher},
+	}
+	return filter, nil
+}
+
 func (gsf *geoSiteFilter) Match(domain string) bool {
 	for _, matcher := range gsf.matchers {
 		if matcher.ApplyDomain(domain) {
diff --git a/dns/resolver.go b/dns/resolver.go
index d99a465d..65bcd6d3 100644
--- a/dns/resolver.go
+++ b/dns/resolver.go
@@ -4,17 +4,20 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"go.uber.org/atomic"
 	"math/rand"
 	"net/netip"
+	"strings"
 	"time"
 
+	"go.uber.org/atomic"
+
 	"github.com/Dreamacro/clash/common/cache"
 	"github.com/Dreamacro/clash/component/fakeip"
 	"github.com/Dreamacro/clash/component/geodata/router"
 	"github.com/Dreamacro/clash/component/resolver"
 	"github.com/Dreamacro/clash/component/trie"
 	C "github.com/Dreamacro/clash/constant"
+	"github.com/Dreamacro/clash/log"
 
 	D "github.com/miekg/dns"
 	"golang.org/x/sync/singleflight"
@@ -30,6 +33,12 @@ type result struct {
 	Error error
 }
 
+type geositePolicyRecord struct {
+	matcher          fallbackDomainFilter
+	policy           *Policy
+	inversedMatching bool
+}
+
 type Resolver struct {
 	ipv6                  bool
 	hosts                 *trie.DomainTrie[netip.Addr]
@@ -40,6 +49,7 @@ type Resolver struct {
 	group                 singleflight.Group
 	lruCache              *cache.LruCache[string, *D.Msg]
 	policy                *trie.DomainTrie[*Policy]
+	geositePolicy         []geositePolicyRecord
 	proxyServer           []dnsClient
 }
 
@@ -272,12 +282,18 @@ func (r *Resolver) matchPolicy(m *D.Msg) []dnsClient {
 	}
 
 	record := r.policy.Search(domain)
-	if record == nil {
-		return nil
+	if record != nil {
+		p := record.Data()
+		return p.GetData()
 	}
 
-	p := record.Data()
-	return p.GetData()
+	for _, geositeRecord := range r.geositePolicy {
+		matched := geositeRecord.matcher.Match(domain)
+		if matched != geositeRecord.inversedMatching {
+			return geositeRecord.policy.GetData()
+		}
+	}
+	return nil
 }
 
 func (r *Resolver) shouldOnlyQueryFallback(m *D.Msg) bool {
@@ -433,7 +449,26 @@ func NewResolver(config Config) *Resolver {
 	if len(config.Policy) != 0 {
 		r.policy = trie.New[*Policy]()
 		for domain, nameserver := range config.Policy {
-			_ = r.policy.Insert(domain, NewPolicy(transform([]NameServer{nameserver}, defaultResolver)))
+			if strings.HasPrefix(strings.ToLower(domain), "@geosite:") {
+				groupname := domain[9:]
+				inverse := false
+				if strings.HasPrefix(groupname, "!") {
+					inverse = true
+					groupname = groupname[1:]
+				}
+				log.Debugln("adding geosite policy: %s inversed %s", groupname, inverse)
+				matcher, err := NewGeoSite(groupname)
+				if err != nil {
+					continue
+				}
+				r.geositePolicy = append(r.geositePolicy, geositePolicyRecord{
+					matcher:          matcher,
+					policy:           NewPolicy(transform([]NameServer{nameserver}, defaultResolver)),
+					inversedMatching: inverse,
+				})
+			} else {
+				_ = r.policy.Insert(domain, NewPolicy(transform([]NameServer{nameserver}, defaultResolver)))
+			}
 		}
 		r.policy.Optimize()
 	}