Feature: dns resolve domain through nameserver-policy (#1406)

dns/resolver.go

@@ -43,6 +43,7 @@ type Resolver struct {
 	fallbackIPFilters     []fallbackIPFilter
 	group                 singleflight.Group
 	lruCache              *cache.LruCache
+	policy                *trie.DomainTrie
 }
 
 // ResolveIP request with TypeA and TypeAAAA, priority return TypeA
@@ -131,6 +132,9 @@ func (r *Resolver) exchangeWithoutCache(m *D.Msg) (msg *D.Msg, err error) {
 			return r.ipExchange(m)
 		}
 
+		if matched := r.matchPolicy(m); len(matched) != 0 {
+			return r.batchExchange(matched, m)
+		}
 		return r.batchExchange(r.main, m)
 	})
 
@@ -172,6 +176,24 @@ func (r *Resolver) batchExchange(clients []dnsClient, m *D.Msg) (msg *D.Msg, err
 	return
 }
 
+func (r *Resolver) matchPolicy(m *D.Msg) []dnsClient {
+	if r.policy == nil {
+		return nil
+	}
+
+	domain := r.msgToDomain(m)
+	if domain == "" {
+		return nil
+	}
+
+	record := r.policy.Search(domain)
+	if record == nil {
+		return nil
+	}
+
+	return record.Data.([]dnsClient)
+}
+
 func (r *Resolver) shouldOnlyQueryFallback(m *D.Msg) bool {
 	if r.fallback == nil || len(r.fallbackDomainFilters) == 0 {
 		return false
@@ -194,6 +216,11 @@ func (r *Resolver) shouldOnlyQueryFallback(m *D.Msg) bool {
 
 func (r *Resolver) ipExchange(m *D.Msg) (msg *D.Msg, err error) {
 
+	if matched := r.matchPolicy(m); len(matched) != 0 {
+		res := <-r.asyncExchange(matched, m)
+		return res.Msg, res.Error
+	}
+
 	onlyFallback := r.shouldOnlyQueryFallback(m)
 
 	if onlyFallback {
@@ -293,6 +320,7 @@ type Config struct {
 	FallbackFilter FallbackFilter
 	Pool           *fakeip.Pool
 	Hosts          *trie.DomainTrie
+	Policy         map[string]NameServer
 }
 
 func NewResolver(config Config) *Resolver {
@@ -312,6 +340,13 @@ func NewResolver(config Config) *Resolver {
 		r.fallback = transform(config.Fallback, defaultResolver)
 	}
 
+	if len(config.Policy) != 0 {
+		r.policy = trie.New()
+		for domain, nameserver := range config.Policy {
+			r.policy.Insert(domain, transform([]NameServer{nameserver}, defaultResolver))
+		}
+	}
+
 	fallbackIPFilters := []fallbackIPFilter{}
 	if config.FallbackFilter.GeoIP {
 		fallbackIPFilters = append(fallbackIPFilters, &geoipFilter{})