From 3257d0776b17e79c477a3befcdee542fb993f377 Mon Sep 17 00:00:00 2001
From: Zerorigin <36463444+Zerorigin@users.noreply.github.com>
Date: Fri, 9 Aug 2024 09:55:50 +0800
Subject: [PATCH] feat: add DNS-over-TCP support to customize DNS servers
 (#1217)

* feat: add DNS-over-TCP & custom Port support to customize DNS servers

* feat: add DNS-over-TCP support to customize DNS servers

* fix: wrong dns string concatenation operator.
---
 util/net_resolver.go | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/util/net_resolver.go b/util/net_resolver.go
index 42d7f87ff..ba85a294d 100644
--- a/util/net_resolver.go
+++ b/util/net_resolver.go
@@ -3,6 +3,8 @@ package util
 import (
 	"context"
 	"net"
+	"net/url"
+	"strings"
 
 	"golang.org/x/text/language"
 )
@@ -24,14 +26,29 @@ func InitBackupDNS(customDNS, lang string) {
 
 // SetDNS sets the dialer.Resolver to use the given DNS server.
 func SetDNS(dns string) {
-	// Error means that the given DNS doesn't have a port. Add it.
-	if _, _, err := net.SplitHostPort(dns); err != nil {
-		dns = net.JoinHostPort(dns, "53")
+
+	if !strings.Contains(dns, "://") {
+		dns = "udp://" + dns
+	}
+	svrParse, _ := url.Parse(dns)
+
+	var network string
+	switch strings.ToLower(svrParse.Scheme) {
+	case "tcp":
+		network = "tcp"
+	default:
+		network = "udp"
+	}
+
+	if svrParse.Port() == "" {
+		dns = net.JoinHostPort(svrParse.Host, "53")
+	} else {
+		dns = svrParse.Host
 	}
 
 	dialer.Resolver = &net.Resolver{
 		PreferGo: true,
-		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+		Dial: func(ctx context.Context, _, address string) (net.Conn, error) {
 			return net.Dial(network, dns)
 		},
 	}