WireGuard outbound: Enhance DNS resolution with family-specific lookups and caching

Author: kirilllivanov

proxy/wireguard/client.go

@@ -143,23 +143,39 @@ func (h *Handler) processWireGuard(ctx context.Context, dialer internet.Dialer)
 	return nil
 }
 
+func (h *Handler) lookupIPWithOption(domain string, option dns.IPOption) (net.Address, error) {
+	ips, _, err := h.dns.LookupIP(domain, option)
+	if err != nil {
+		return nil, err
+	} else if len(ips) == 0 {
+		return nil, dns.ErrEmptyResponse
+	}
+	return net.IPAddress(ips[dice.Roll(len(ips))]), nil
+}
+
 func (h *Handler) lookupIP(domain string) (net.Address, error) {
-	ips, _, err := h.dns.LookupIP(domain, dns.IPOption{
+	addr, err := h.lookupIPWithOption(domain, dns.IPOption{
 		IPv4Enable: h.hasIPv4 && h.conf.preferIP4(),
 		IPv6Enable: h.hasIPv6 && h.conf.preferIP6(),
 	})
-	if (len(ips) == 0 || err != nil) && h.conf.hasFallback() {
-		ips, _, err = h.dns.LookupIP(domain, dns.IPOption{
+	if err != nil && h.conf.hasFallback() {
+		return h.lookupIPWithOption(domain, dns.IPOption{
 			IPv4Enable: h.hasIPv4 && h.conf.fallbackIP4(),
 			IPv6Enable: h.hasIPv6 && h.conf.fallbackIP6(),
 		})
 	}
-	if err != nil {
-		return nil, err
-	} else if len(ips) == 0 {
-		return nil, dns.ErrEmptyResponse
+	return addr, err
+}
+
+func (h *Handler) lookupIPForFamily(domain string, family net.AddressFamily) (net.Address, error) {
+	switch family {
+	case net.AddressFamilyIPv4:
+		return h.lookupIPWithOption(domain, dns.IPOption{IPv4Enable: true})
+	case net.AddressFamilyIPv6:
+		return h.lookupIPWithOption(domain, dns.IPOption{IPv6Enable: true})
+	default:
+		return h.lookupIP(domain)
 	}
-	return net.IPAddress(ips[dice.Roll(len(ips))]), nil
 }
 
 // Process implements OutboundHandler.Dispatch().
@@ -185,8 +201,10 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 
 	// resolve dns
 	addr := destination.Address
+	destinationDomain := ""
 	if addr.Family().IsDomain() {
-		resolved, err := h.lookupIP(addr.Domain())
+		destinationDomain = addr.Domain()
+		resolved, err := h.lookupIP(destinationDomain)
 		if err != nil {
 			return errors.New("failed to lookup DNS").Base(err)
 		}
@@ -236,11 +254,18 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 		}
 		defer conn.Close()
 
+		resolvedUDPAddr := utils.NewTypedSyncMap[string, net.Address]()
+		if destinationDomain != "" {
+			resolvedUDPAddr.Store(strings.ToLower(destinationDomain), addr)
+		}
+		udpLookupIP := func(domain string) (net.Address, error) {
+			return h.lookupIPForFamily(domain, addr.Family())
+		}
 		conn = &udpConnClient{
 			Conn:            conn,
 			dest:            destination,
-			lookupIP:        h.lookupIP,
-			resolvedUDPAddr: utils.NewTypedSyncMap[string, net.Address](),
+			lookupIP:        udpLookupIP,
+			resolvedUDPAddr: resolvedUDPAddr,
 		}
 
 		requestFunc = func() error {

proxy/wireguard/client_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/xtls/xray-core/common/buf"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/utils"
+	feature_dns "github.com/xtls/xray-core/features/dns"
 	"github.com/xtls/xray-core/transport/internet"
 )
 
@@ -45,6 +46,33 @@ func (c *capturePacketConn) SetWriteDeadline(time.Time) error {
 	return nil
 }
 
+type recordingDNSClient struct {
+	options []feature_dns.IPOption
+}
+
+func (c *recordingDNSClient) Type() interface{} {
+	return feature_dns.ClientType()
+}
+
+func (c *recordingDNSClient) Start() error {
+	return nil
+}
+
+func (c *recordingDNSClient) Close() error {
+	return nil
+}
+
+func (c *recordingDNSClient) LookupIP(_ string, option feature_dns.IPOption) ([]net.IP, uint32, error) {
+	c.options = append(c.options, option)
+	if option.IPv4Enable {
+		return []net.IP{gonet.IPv4(162, 159, 36, 1)}, 0, nil
+	}
+	if option.IPv6Enable {
+		return []net.IP{gonet.ParseIP("2606:4700:4700::1111")}, 0, nil
+	}
+	return nil, 0, feature_dns.ErrEmptyResponse
+}
+
 func TestUDPConnClientWriteMultiBufferResolvesPacketDestination(t *testing.T) {
 	packetConn := &capturePacketConn{}
 	conn := &udpConnClient{
@@ -143,3 +171,56 @@ func TestUDPConnClientResolveDestinationCachesDomainsCaseInsensitive(t *testing.
 		t.Fatalf("unexpected cached addresses: %s != %s", first.Address, second.Address)
 	}
 }
+
+func TestUDPConnClientResolveDestinationUsesSeededDomainCache(t *testing.T) {
+	lookupCalls := 0
+	cache := utils.NewTypedSyncMap[string, net.Address]()
+	cache.Store("dns.cloudflare.internal", net.IPAddress([]byte{162, 159, 36, 1}))
+	conn := &udpConnClient{
+		resolvedUDPAddr: cache,
+		lookupIP: func(domain string) (net.Address, error) {
+			lookupCalls++
+			return net.IPAddress([]byte{162, 159, 46, 1}), nil
+		},
+	}
+
+	resolved, err := conn.resolveDestination(net.UDPDestination(net.DomainAddress("DNS.Cloudflare.Internal"), net.Port(53)))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if lookupCalls != 0 {
+		t.Fatalf("unexpected lookup calls: %d", lookupCalls)
+	}
+	if resolved.Address.String() != "162.159.36.1" {
+		t.Fatalf("unexpected seeded address: %s", resolved.Address)
+	}
+}
+
+func TestHandlerLookupIPForFamilyConstrainsLookupOption(t *testing.T) {
+	client := &recordingDNSClient{}
+	handler := &Handler{dns: client}
+
+	ipv4, err := handler.lookupIPForFamily("dns.cloudflare.internal", net.AddressFamilyIPv4)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !ipv4.Family().IsIPv4() {
+		t.Fatalf("unexpected IPv4 lookup result: %s", ipv4)
+	}
+	if len(client.options) != 1 || !client.options[0].IPv4Enable || client.options[0].IPv6Enable {
+		t.Fatalf("unexpected IPv4 lookup options: %+v", client.options)
+	}
+
+	client.options = nil
+	ipv6, err := handler.lookupIPForFamily("dns.cloudflare.internal", net.AddressFamilyIPv6)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !ipv6.Family().IsIPv6() {
+		t.Fatalf("unexpected IPv6 lookup result: %s", ipv6)
+	}
+	if len(client.options) != 1 || client.options[0].IPv4Enable || !client.options[0].IPv6Enable {
+		t.Fatalf("unexpected IPv6 lookup options: %+v", client.options)
+	}
+}