feat: nameserver policy support multiple server

xishang0128 · Feb 3, 2023 · e1e1984 · e1e1984
1 parent 99662b6
commit e1e1984
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 12 deletions.
diff --git a/config/config.go b/config/config.go
@@ -8,6 +8,7 @@ import (
 	"net/netip"
 	"net/url"
 	"os"
+	"reflect"
 	"runtime"
 	"strconv"
 	"strings"
@@ -97,7 +98,7 @@ type DNS struct {
 	DefaultNameserver     []dns.NameServer `yaml:"default-nameserver"`
 	FakeIPRange           *fakeip.Pool
 	Hosts                 *trie.DomainTrie[netip.Addr]
-	NameServerPolicy      map[string]dns.NameServer
+	NameServerPolicy      map[string][]dns.NameServer
 	ProxyServerNameserver []dns.NameServer
 }
 
@@ -181,7 +182,7 @@ type RawDNS struct {
 	FakeIPRange           string            `yaml:"fake-ip-range"`
 	FakeIPFilter          []string          `yaml:"fake-ip-filter"`
 	DefaultNameserver     []string          `yaml:"default-nameserver"`
-	NameServerPolicy      map[string]string `yaml:"nameserver-policy"`
+	NameServerPolicy      map[string]any    `yaml:"nameserver-policy"`
 	ProxyServerNameserver []string          `yaml:"proxy-server-nameserver"`
 }
 
@@ -952,18 +953,35 @@ func parsePureDNSServer(server string) string {
 		}
 	}
 }
-func parseNameServerPolicy(nsPolicy map[string]string, preferH3 bool) (map[string]dns.NameServer, error) {
-	policy := map[string]dns.NameServer{}
+func parseNameServerPolicy(nsPolicy map[string]any, preferH3 bool) (map[string][]dns.NameServer, error) {
+	policy := map[string][]dns.NameServer{}
 
 	for domain, server := range nsPolicy {
-		nameservers, err := parseNameServer([]string{server}, preferH3)
+		var (
+			nameservers []dns.NameServer
+			err         error
+		)
+
+		switch reflect.TypeOf(server).Kind() {
+		case reflect.Slice, reflect.Array:
+			origin := reflect.ValueOf(server)
+			servers := make([]string, 0)
+			for i := 0; i < origin.Len(); i++ {
+				servers = append(servers, fmt.Sprintf("%v", origin.Index(i)))
+			}
+			nameservers, err = parseNameServer(servers, preferH3)
+		case reflect.String:
+			nameservers, err = parseNameServer([]string{fmt.Sprintf("%v", server)}, preferH3)
+		default:
+			return nil, errors.New("server format error, must be string or array")
+		}
 		if err != nil {
 			return nil, err
 		}
 		if _, valid := trie.ValidAndSplitDomain(domain); !valid {
 			return nil, fmt.Errorf("DNS ResoverRule invalid domain: %s", domain)
 		}
-		policy[domain] = nameservers[0]
+		policy[domain] = nameservers
 	}
 
 	return policy, nil

diff --git a/dns/enhancer.go b/dns/enhancer.go
@@ -109,7 +109,7 @@ func NewEnhancer(cfg Config) *ResolverEnhancer {
 
 	if cfg.EnhancedMode != C.DNSNormal {
 		fakePool = cfg.Pool
-		mapping = cache.New[netip.Addr, string](cache.WithSize[netip.Addr, string](4096), cache.WithStale[netip.Addr, string](true))
+		mapping = cache.New(cache.WithSize[netip.Addr, string](4096), cache.WithStale[netip.Addr, string](true))
 	}
 
 	return &ResolverEnhancer{

diff --git a/dns/resolver.go b/dns/resolver.go
@@ -423,19 +423,19 @@ type Config struct {
 	FallbackFilter FallbackFilter
 	Pool           *fakeip.Pool
 	Hosts          *trie.DomainTrie[netip.Addr]
-	Policy         map[string]NameServer
+	Policy         map[string][]NameServer
 }
 
 func NewResolver(config Config) *Resolver {
 	defaultResolver := &Resolver{
 		main:     transform(config.Default, nil),
-		lruCache: cache.New[string, *D.Msg](cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
+		lruCache: cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
 	}
 
 	r := &Resolver{
 		ipv6:     config.IPv6,
 		main:     transform(config.Main, defaultResolver),
-		lruCache: cache.New[string, *D.Msg](cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
+		lruCache: cache.New(cache.WithSize[string, *D.Msg](4096), cache.WithStale[string, *D.Msg](true)),
 		hosts:    config.Hosts,
 	}
 
@@ -464,11 +464,11 @@ func NewResolver(config Config) *Resolver {
 				}
 				r.geositePolicy = append(r.geositePolicy, geositePolicyRecord{
 					matcher:          matcher,
-					policy:           NewPolicy(transform([]NameServer{nameserver}, defaultResolver)),
+					policy:           NewPolicy(transform(nameserver, defaultResolver)),
 					inversedMatching: inverse,
 				})
 			} else {
-				_ = r.policy.Insert(domain, NewPolicy(transform([]NameServer{nameserver}, defaultResolver)))
+				_ = r.policy.Insert(domain, NewPolicy(transform(nameserver, defaultResolver)))
 			}
 		}
 		r.policy.Optimize()