From 8674ed5a0db0627d0f55803ebc685397a1847bb9 Mon Sep 17 00:00:00 2001
From: xiaorouji <xiaoroujixxoo@gmail.com>
Date: Tue, 27 Aug 2024 22:19:33 +0800
Subject: [PATCH] Support DNS queryStrategy config for UDP NameServer (#3728)

---
 app/dns/nameserver.go     |  2 +-
 app/dns/nameserver_udp.go | 34 ++++++++++++++++++++--------------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/app/dns/nameserver.go b/app/dns/nameserver.go
index 8dfa9dc2..cf79103f 100644
--- a/app/dns/nameserver.go
+++ b/app/dns/nameserver.go
@@ -62,7 +62,7 @@ func NewServer(dest net.Destination, dispatcher routing.Dispatcher, queryStrateg
 		dest.Network = net.Network_UDP
 	}
 	if dest.Network == net.Network_UDP { // UDP classic DNS mode
-		return NewClassicNameServer(dest, dispatcher), nil
+		return NewClassicNameServer(dest, dispatcher, queryStrategy), nil
 	}
 	return nil, errors.New("No available name server could be created from ", dest).AtWarning()
 }
diff --git a/app/dns/nameserver_udp.go b/app/dns/nameserver_udp.go
index 67e1d73a..24a99593 100644
--- a/app/dns/nameserver_udp.go
+++ b/app/dns/nameserver_udp.go
@@ -24,29 +24,31 @@ import (
 // ClassicNameServer implemented traditional UDP DNS.
 type ClassicNameServer struct {
 	sync.RWMutex
-	name      string
-	address   *net.Destination
-	ips       map[string]*record
-	requests  map[uint16]*dnsRequest
-	pub       *pubsub.Service
-	udpServer *udp.Dispatcher
-	cleanup   *task.Periodic
-	reqID     uint32
+	name          string
+	address       *net.Destination
+	ips           map[string]*record
+	requests      map[uint16]*dnsRequest
+	pub           *pubsub.Service
+	udpServer     *udp.Dispatcher
+	cleanup       *task.Periodic
+	reqID         uint32
+	queryStrategy QueryStrategy
 }
 
 // NewClassicNameServer creates udp server object for remote resolving.
-func NewClassicNameServer(address net.Destination, dispatcher routing.Dispatcher) *ClassicNameServer {
+func NewClassicNameServer(address net.Destination, dispatcher routing.Dispatcher, queryStrategy QueryStrategy) *ClassicNameServer {
 	// default to 53 if unspecific
 	if address.Port == 0 {
 		address.Port = net.Port(53)
 	}
 
 	s := &ClassicNameServer{
-		address:  &address,
-		ips:      make(map[string]*record),
-		requests: make(map[uint16]*dnsRequest),
-		pub:      pubsub.NewService(),
-		name:     strings.ToUpper(address.String()),
+		address:       &address,
+		ips:           make(map[string]*record),
+		requests:      make(map[uint16]*dnsRequest),
+		pub:           pubsub.NewService(),
+		name:          strings.ToUpper(address.String()),
+		queryStrategy: queryStrategy,
 	}
 	s.cleanup = &task.Periodic{
 		Interval: time.Minute,
@@ -239,6 +241,10 @@ func (s *ClassicNameServer) findIPsForDomain(domain string, option dns_feature.I
 // QueryIP implements Server.
 func (s *ClassicNameServer) QueryIP(ctx context.Context, domain string, clientIP net.IP, option dns_feature.IPOption, disableCache bool) ([]net.IP, error) {
 	fqdn := Fqdn(domain)
+	option = ResolveIpOptionOverride(s.queryStrategy, option)
+	if !option.IPv4Enable && !option.IPv6Enable {
+		return nil, dns_feature.ErrEmptyResponse
+	}
 
 	if disableCache {
 		errors.LogDebug(ctx, "DNS cache is disabled. Querying IP for ", domain, " at ", s.name)