From e316cd4c66c5d01044d652866280162fda335b06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 1 Sep 2021 09:37:40 +0800
Subject: [PATCH] Hook to replace DOHL dialer

---
 app/dns/nameserver_doh.go           |  2 +-
 transport/internet/dialer.go        | 29 +++++++++++++++++++++++++++++
 transport/internet/system_dialer.go | 12 +++++++++++-
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/app/dns/nameserver_doh.go b/app/dns/nameserver_doh.go
index 6286c3e5..3e8b0f53 100644
--- a/app/dns/nameserver_doh.go
+++ b/app/dns/nameserver_doh.go
@@ -114,7 +114,7 @@ func NewDoHLocalNameServer(url *url.URL) *DoHNameServer {
 			if err != nil {
 				return nil, err
 			}
-			conn, err := internet.DialSystem(ctx, dest, nil)
+			conn, err := internet.DialSystemDNS(ctx, dest, nil)
 			log.Record(&log.AccessMessage{
 				From:   "DoH",
 				To:     s.dohURL,
diff --git a/transport/internet/dialer.go b/transport/internet/dialer.go
index fe676b3b..ae9d99b1 100644
--- a/transport/internet/dialer.go
+++ b/transport/internet/dialer.go
@@ -161,6 +161,35 @@ func DialSystem(ctx context.Context, dest net.Destination, sockopt *SocketConfig
 	return effectiveSystemDialer.Dial(ctx, src, dest, sockopt)
 }
 
+func DialSystemDNS(ctx context.Context, dest net.Destination, sockopt *SocketConfig) (net.Conn, error) {
+	var src net.Address
+	if outbound := session.OutboundFromContext(ctx); outbound != nil {
+		src = outbound.Gateway
+	}
+	if sockopt == nil {
+		return effectiveSystemDNSDialer.Dial(ctx, src, dest, sockopt)
+	}
+
+	if canLookupIP(ctx, dest, sockopt) {
+		ips, err := lookupIP(dest.Address.String(), sockopt.DomainStrategy, src)
+		if err == nil && len(ips) > 0 {
+			dest.Address = net.IPAddress(ips[dice.Roll(len(ips))])
+			newError("replace destination with " + dest.String()).AtInfo().WriteToLog()
+		} else if err != nil {
+			newError("failed to resolve ip").Base(err).AtWarning().WriteToLog()
+		}
+	}
+
+	if obm != nil && len(sockopt.DialerProxy) > 0 {
+		nc := redirect(ctx, dest, sockopt.DialerProxy)
+		if nc != nil {
+			return nc, nil
+		}
+	}
+
+	return effectiveSystemDNSDialer.Dial(ctx, src, dest, sockopt)
+}
+
 func InitSystemDialer(dc dns.Client, om outbound.Manager) {
 	dnsClient = dc
 	obm = om
diff --git a/transport/internet/system_dialer.go b/transport/internet/system_dialer.go
index 8058fa1b..9feab5c4 100644
--- a/transport/internet/system_dialer.go
+++ b/transport/internet/system_dialer.go
@@ -11,7 +11,10 @@ import (
 	"github.com/xtls/xray-core/features/outbound"
 )
 
-var effectiveSystemDialer SystemDialer = &DefaultSystemDialer{}
+var (
+	effectiveSystemDialer    SystemDialer = &DefaultSystemDialer{}
+	effectiveSystemDNSDialer SystemDialer = &DefaultSystemDialer{}
+)
 
 type SystemDialer interface {
 	Dial(ctx context.Context, source net.Address, destination net.Destination, sockopt *SocketConfig) (net.Conn, error)
@@ -176,6 +179,13 @@ func UseAlternativeSystemDialer(dialer SystemDialer) {
 	effectiveSystemDialer = dialer
 }
 
+func UseAlternativeSystemDNSDialer(dialer SystemDialer) {
+	if dialer == nil {
+		effectiveSystemDNSDialer = &DefaultSystemDialer{}
+	}
+	effectiveSystemDNSDialer = dialer
+}
+
 // RegisterDialerController adds a controller to the effective system dialer.
 // The controller can be used to operate on file descriptors before they are put into use.
 // It only works when effective dialer is the default dialer.