Add address filter support for DNS rules

route/router_dns.go

@@ -2,13 +2,13 @@ package route
 
 import (
 	"context"
+	"errors"
 	"net/netip"
 	"strings"
 	"time"
 
 	"github.com/sagernet/sing-box/adapter"
 	C "github.com/sagernet/sing-box/constant"
-	"github.com/sagernet/sing-box/log"
 	"github.com/sagernet/sing-dns"
 	"github.com/sagernet/sing/common/cache"
 	E "github.com/sagernet/sing/common/exceptions"
@@ -37,41 +37,51 @@ func (m *DNSReverseMapping) Query(address netip.Addr) (string, bool) {
 	return domain, loaded
 }
 
-func (r *Router) matchDNS(ctx context.Context, allowFakeIP bool) (context.Context, dns.Transport, dns.DomainStrategy) {
+func (r *Router) matchDNS(ctx context.Context, allowFakeIP bool, index int) (context.Context, dns.Transport, dns.DomainStrategy, adapter.DNSRule, int) {
 	metadata := adapter.ContextFrom(ctx)
 	if metadata == nil {
 		panic("no context")
 	}
-	for i, rule := range r.dnsRules {
-		metadata.ResetRuleCache()
-		if rule.Match(metadata) {
-			detour := rule.Outbound()
-			transport, loaded := r.transportMap[detour]
-			if !loaded {
-				r.dnsLogger.ErrorContext(ctx, "transport not found: ", detour)
-				continue
-			}
-			if _, isFakeIP := transport.(adapter.FakeIPTransport); isFakeIP && !allowFakeIP {
-				continue
-			}
-			r.dnsLogger.DebugContext(ctx, "match[", i, "] ", rule.String(), " => ", detour)
-			if rule.DisableCache() {
-				ctx = dns.ContextWithDisableCache(ctx, true)
-			}
-			if rewriteTTL := rule.RewriteTTL(); rewriteTTL != nil {
-				ctx = dns.ContextWithRewriteTTL(ctx, *rewriteTTL)
-			}
-			if domainStrategy, dsLoaded := r.transportDomainStrategy[transport]; dsLoaded {
-				return ctx, transport, domainStrategy
-			} else {
-				return ctx, transport, r.defaultDomainStrategy
+	if index < len(r.dnsRules) {
+		dnsRules := r.dnsRules
+		if index != -1 {
+			dnsRules = dnsRules[index+1:]
+		}
+		for ruleIndex, rule := range dnsRules {
+			metadata.ResetRuleCache()
+			if rule.Match(metadata) {
+				detour := rule.Outbound()
+				transport, loaded := r.transportMap[detour]
+				if !loaded {
+					r.dnsLogger.ErrorContext(ctx, "transport not found: ", detour)
+					continue
+				}
+				if _, isFakeIP := transport.(adapter.FakeIPTransport); isFakeIP && !allowFakeIP {
+					continue
+				}
+				displayRuleIndex := ruleIndex
+				if index != -1 {
+					displayRuleIndex += index + 1
+				}
+				r.dnsLogger.DebugContext(ctx, "match[", displayRuleIndex, "] ", rule.String(), " => ", detour)
+				if rule.DisableCache() {
+					ctx = dns.ContextWithDisableCache(ctx, true)
+				}
+				if rewriteTTL := rule.RewriteTTL(); rewriteTTL != nil {
+					ctx = dns.ContextWithRewriteTTL(ctx, *rewriteTTL)
+				}
+				if domainStrategy, dsLoaded := r.transportDomainStrategy[transport]; dsLoaded {
+					return ctx, transport, domainStrategy, rule, ruleIndex
+				} else {
+					return ctx, transport, r.defaultDomainStrategy, rule, ruleIndex
+				}
 			}
 		}
 	}
 	if domainStrategy, dsLoaded := r.transportDomainStrategy[r.defaultTransport]; dsLoaded {
-		return ctx, r.defaultTransport, domainStrategy
+		return ctx, r.defaultTransport, domainStrategy, nil, -1
 	} else {
-		return ctx, r.defaultTransport, r.defaultDomainStrategy
+		return ctx, r.defaultTransport, r.defaultDomainStrategy, nil, -1
 	}
 }
 
@@ -86,7 +96,8 @@ func (r *Router) Exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, er
 	)
 	response, cached = r.dnsClient.ExchangeCache(ctx, message)
 	if !cached {
-		ctx, metadata := adapter.AppendContext(ctx)
+		var metadata *adapter.InboundContext
+		ctx, metadata = adapter.AppendContext(ctx)
 		if len(message.Question) > 0 {
 			metadata.QueryType = message.Question[0].Qtype
 			switch metadata.QueryType {
@@ -97,17 +108,47 @@ func (r *Router) Exchange(ctx context.Context, message *mDNS.Msg) (*mDNS.Msg, er
 			}
 			metadata.Domain = fqdnToDomain(message.Question[0].Name)
 		}
-		ctx, transport, strategy := r.matchDNS(ctx, true)
-		ctx, cancel := context.WithTimeout(ctx, C.DNSTimeout)
-		defer cancel()
-		response, err = r.dnsClient.Exchange(ctx, transport, message, strategy)
-		if err != nil && len(message.Question) > 0 {
-			r.dnsLogger.ErrorContext(ctx, E.Cause(err, "exchange failed for ", formatQuestion(message.Question[0].String())))
+		var (
+			transport dns.Transport
+			strategy  dns.DomainStrategy
+			rule      adapter.DNSRule
+			ruleIndex int
+		)
+		ruleIndex = -1
+		for {
+			var (
+				dnsCtx       context.Context
+				cancel       context.CancelFunc
+				addressLimit bool
+			)
+
+			dnsCtx, transport, strategy, rule, ruleIndex = r.matchDNS(ctx, true, ruleIndex)
+			dnsCtx, cancel = context.WithTimeout(dnsCtx, C.DNSTimeout)
+			if rule != nil && rule.WithAddressLimit() && isAddressQuery(message) {
+				addressLimit = true
+				response, err = r.dnsClient.ExchangeWithResponseCheck(dnsCtx, transport, message, strategy, func(response *mDNS.Msg) bool {
+					metadata.DestinationAddresses, _ = dns.MessageToAddresses(response)
+					return rule.MatchAddressLimit(metadata)
+				})
+			} else {
+				addressLimit = false
+				response, err = r.dnsClient.Exchange(dnsCtx, transport, message, strategy)
+			}
+			cancel()
+			if err != nil {
+				if errors.Is(err, dns.ErrResponseRejected) {
+					r.dnsLogger.DebugContext(ctx, E.Cause(err, "response rejected for ", formatQuestion(message.Question[0].String())))
+				} else if len(message.Question) > 0 {
+					r.dnsLogger.ErrorContext(ctx, E.Cause(err, "exchange failed for ", formatQuestion(message.Question[0].String())))
+				} else {
+					r.dnsLogger.ErrorContext(ctx, E.Cause(err, "exchange failed for <empty query>"))
+				}
+			}
+			if !addressLimit || err == nil {
+				break
+			}
 		}
 	}
-	if len(message.Question) > 0 && response != nil {
-		LogDNSAnswers(r.dnsLogger, ctx, message.Question[0].Name, response.Answer)
-	}
 	if r.dnsReverseMapping != nil && len(message.Question) > 0 && response != nil && len(response.Answer) > 0 {
 		for _, answer := range response.Answer {
 			switch record := answer.(type) {
@@ -125,22 +166,57 @@ func (r *Router) Lookup(ctx context.Context, domain string, strategy dns.DomainS
 	r.dnsLogger.DebugContext(ctx, "lookup domain ", domain)
 	ctx, metadata := adapter.AppendContext(ctx)
 	metadata.Domain = domain
-	ctx, transport, transportStrategy := r.matchDNS(ctx, false)
-	if strategy == dns.DomainStrategyAsIS {
-		strategy = transportStrategy
+	var (
+		transport         dns.Transport
+		transportStrategy dns.DomainStrategy
+		rule              adapter.DNSRule
+		ruleIndex         int
+		resultAddrs       []netip.Addr
+		err               error
+	)
+	ruleIndex = -1
+	for {
+		var (
+			dnsCtx       context.Context
+			cancel       context.CancelFunc
+			addressLimit bool
+		)
+		metadata.ResetRuleCache()
+		metadata.DestinationAddresses = nil
+		dnsCtx, transport, transportStrategy, rule, ruleIndex = r.matchDNS(ctx, false, ruleIndex)
+		if strategy == dns.DomainStrategyAsIS {
+			strategy = transportStrategy
+		}
+		dnsCtx, cancel = context.WithTimeout(dnsCtx, C.DNSTimeout)
+		if rule != nil && rule.WithAddressLimit() {
+			addressLimit = true
+			resultAddrs, err = r.dnsClient.LookupWithResponseCheck(dnsCtx, transport, domain, strategy, func(responseAddrs []netip.Addr) bool {
+				metadata.DestinationAddresses = responseAddrs
+				return rule.MatchAddressLimit(metadata)
+			})
+		} else {
+			addressLimit = false
+			resultAddrs, err = r.dnsClient.Lookup(dnsCtx, transport, domain, strategy)
+		}
+		cancel()
+		if err != nil {
+			if errors.Is(err, dns.ErrResponseRejected) {
+				r.dnsLogger.DebugContext(ctx, "response rejected for ", domain)
+			} else {
+				r.dnsLogger.ErrorContext(ctx, E.Cause(err, "lookup failed for ", domain))
+			}
+		} else if len(resultAddrs) == 0 {
+			r.dnsLogger.ErrorContext(ctx, "lookup failed for ", domain, ": empty result")
+			err = dns.RCodeNameError
+		}
+		if !addressLimit || err == nil {
+			break
+		}
 	}
-	ctx, cancel := context.WithTimeout(ctx, C.DNSTimeout)
-	defer cancel()
-	addrs, err := r.dnsClient.Lookup(ctx, transport, domain, strategy)
-	if len(addrs) > 0 {
-		r.dnsLogger.InfoContext(ctx, "lookup succeed for ", domain, ": ", strings.Join(F.MapToString(addrs), " "))
-	} else if err != nil {
-		r.dnsLogger.ErrorContext(ctx, E.Cause(err, "lookup failed for ", domain))
-	} else {
-		r.dnsLogger.ErrorContext(ctx, "lookup failed for ", domain, ": empty result")
-		err = dns.RCodeNameError
+	if len(resultAddrs) > 0 {
+		r.dnsLogger.InfoContext(ctx, "lookup succeed for ", domain, ": ", strings.Join(F.MapToString(resultAddrs), " "))
 	}
-	return addrs, err
+	return resultAddrs, err
 }
 
 func (r *Router) LookupDefault(ctx context.Context, domain string) ([]netip.Addr, error) {
@@ -154,10 +230,13 @@ func (r *Router) ClearDNSCache() {
 	}
 }
 
-func LogDNSAnswers(logger log.ContextLogger, ctx context.Context, domain string, answers []mDNS.RR) {
-	for _, answer := range answers {
-		logger.InfoContext(ctx, "exchanged ", domain, " ", mDNS.Type(answer.Header().Rrtype).String(), " ", formatQuestion(answer.String()))
+func isAddressQuery(message *mDNS.Msg) bool {
+	for _, question := range message.Question {
+		if question.Qtype == mDNS.TypeA || question.Qtype == mDNS.TypeAAAA {
+			return true
+		}
 	}
+	return false
 }
 
 func fqdnToDomain(fqdn string) string {