DNS outbound: Add rules (matches qtype and domain, then action) (#5981)

https://github.com/XTLS/Xray-core/pull/5981#issuecomment-4279809648

Example: https://github.com/XTLS/Xray-core/pull/5981#issuecomment-4283200236

Closes https://github.com/XTLS/Xray-core/issues/5218

proxy/dns/dns.go

@@ -11,6 +11,7 @@ import (
 	"github.com/xtls/xray-core/common"
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/errors"
+	"github.com/xtls/xray-core/common/geodata"
 	"github.com/xtls/xray-core/common/net"
 	dns_proto "github.com/xtls/xray-core/common/protocol/dns"
 	"github.com/xtls/xray-core/common/session"
@@ -40,6 +41,31 @@ func init() {
 	}))
 }
 
+type DNSRule struct {
+	action  RuleAction
+	qTypes  []uint16
+	domains geodata.DomainMatcher
+}
+
+func (r *DNSRule) matchQType(qType uint16) bool {
+	if len(r.qTypes) == 0 {
+		return true
+	}
+	for _, t := range r.qTypes {
+		if t == qType {
+			return true
+		}
+	}
+	return false
+}
+
+func (r *DNSRule) Apply(qType uint16, domain string) bool {
+	if !r.matchQType(qType) {
+		return false
+	}
+	return r.domains == nil || r.domains.MatchAny(strings.TrimSuffix(strings.ToLower(domain), "."))
+}
+
 type ownLinkVerifier interface {
 	IsOwnLink(ctx context.Context) bool
 }
@@ -50,8 +76,7 @@ type Handler struct {
 	ownLinkVerifier ownLinkVerifier
 	server          net.Destination
 	timeout         time.Duration
-	nonIPQuery      string
-	blockTypes      []int32
+	rules           []*DNSRule
 }
 
 func (h *Handler) Init(config *Config, dnsClient dns.Client, policyManager policy.Manager) error {
@@ -65,11 +90,26 @@ func (h *Handler) Init(config *Config, dnsClient dns.Client, policyManager polic
 	if config.Server != nil {
 		h.server = config.Server.AsDestination()
 	}
-	h.nonIPQuery = config.Non_IPQuery
-	if h.nonIPQuery == "" {
-		h.nonIPQuery = "reject"
+
+	h.rules = make([]*DNSRule, 0, len(config.Rule))
+	for _, r := range config.Rule {
+		rule := &DNSRule{
+			action: r.Action,
+			qTypes: make([]uint16, 0, len(r.Qtype)),
+		}
+		for _, t := range r.Qtype {
+			rule.qTypes = append(rule.qTypes, uint16(t))
+		}
+		if len(r.Domain) > 0 {
+			m, err := geodata.DomainReg.BuildDomainMatcher(r.Domain)
+			if err != nil {
+				return err
+			}
+			rule.domains = m
+		}
+		h.rules = append(h.rules, rule)
 	}
-	h.blockTypes = config.BlockTypes
+
 	return nil
 }
 
@@ -77,30 +117,38 @@ func (h *Handler) isOwnLink(ctx context.Context) bool {
 	return h.ownLinkVerifier != nil && h.ownLinkVerifier.IsOwnLink(ctx)
 }
 
-func parseIPQuery(b []byte) (r bool, domain string, id uint16, qType dnsmessage.Type) {
+func parseQuery(b []byte) (id uint16, qType dnsmessage.Type, domain string, ok bool) {
 	var parser dnsmessage.Parser
 	header, err := parser.Start(b)
 	if err != nil {
 		errors.LogInfoInner(context.Background(), err, "parser start")
 		return
 	}
-
 	id = header.ID
 	q, err := parser.Question()
 	if err != nil {
 		errors.LogInfoInner(context.Background(), err, "question")
 		return
 	}
-	domain = q.Name.String()
 	qType = q.Type
-	if qType != dnsmessage.TypeA && qType != dnsmessage.TypeAAAA {
-		return
-	}
-
-	r = true
+	domain = q.Name.String()
+	ok = true
 	return
 }
 
+func (h *Handler) applyRules(qType dnsmessage.Type, domain string) RuleAction {
+	qCode := uint16(qType)
+	for _, r := range h.rules {
+		if r.Apply(qCode, domain) {
+			return r.action
+		}
+	}
+	if qType == dnsmessage.TypeA || qType == dnsmessage.TypeAAAA {
+		return RuleAction_Hijack
+	}
+	return RuleAction_Reject
+}
+
 // Process implements proxy.Outbound.
 func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.Dialer) error {
 	outbounds := session.OutboundsFromContext(ctx)
@@ -183,51 +231,51 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, d internet.
 			if err == io.EOF {
 				return nil
 			}
-
 			if err != nil {
 				return err
 			}
 
 			timer.Update()
 
-			if !h.isOwnLink(ctx) {
-				isIPQuery, domain, id, qType := parseIPQuery(b.Bytes())
-				if len(h.blockTypes) > 0 {
-					for _, blocktype := range h.blockTypes {
-						if blocktype == int32(qType) {
-							b.Release()
-							errors.LogInfo(ctx, "blocked type ", qType, " query for domain ", domain)
-							if h.nonIPQuery == "reject" {
-								err := h.rejectNonIPQuery(id, qType, domain, writer)
-								if err != nil {
-									return err
-								}
-							}
-							return nil
-						}
-					}
-				}
-				if isIPQuery {
-					b.Release()
-					go h.handleIPQuery(id, qType, domain, writer, timer)
-					continue
-				}
-				if h.nonIPQuery == "drop" {
-					b.Release()
-					continue
-				}
-				if h.nonIPQuery == "reject" {
-					b.Release()
-					err := h.rejectNonIPQuery(id, qType, domain, writer)
-					if err != nil {
-						return err
-					}
-					continue
+			if h.isOwnLink(ctx) {
+				if err := connWriter.WriteMessage(b); err != nil {
+					return err
 				}
+				continue
 			}
 
-			if err := connWriter.WriteMessage(b); err != nil {
-				return err
+			id, qType, domain, ok := parseQuery(b.Bytes())
+			if !ok {
+				b.Release()
+				continue
+			}
+
+			switch h.applyRules(qType, domain) {
+			case RuleAction_Drop:
+				b.Release()
+				errors.LogInfo(ctx, "blocked type ", qType, " query for domain ", domain)
+			case RuleAction_Reject:
+				b.Release()
+				errors.LogInfo(ctx, "rejected type ", qType, " query for domain ", domain)
+				if err := h.rejectNonIPQuery(id, qType, domain, writer); err != nil {
+					return err
+				}
+			case RuleAction_Hijack:
+				b.Release()
+				if qType != dnsmessage.TypeA && qType != dnsmessage.TypeAAAA {
+					errors.LogError(ctx, "can only hijack A/AAAA records")
+					if err := h.rejectNonIPQuery(id, qType, domain, writer); err != nil {
+						return err
+					}
+				} else {
+					go h.handleIPQuery(id, qType, domain, writer, timer)
+				}
+			case RuleAction_Direct:
+				if err := connWriter.WriteMessage(b); err != nil {
+					return err
+				}
+			default:
+				panic("unknown rule action")
 			}
 		}
 	}