From d42a7e2eb83e64952d145abb13c261c84cc1bdb1 Mon Sep 17 00:00:00 2001
From: WaterLemons2k <62788816+WaterLemons2k@users.noreply.github.com>
Date: Mon, 24 Apr 2023 20:52:33 +0800
Subject: [PATCH] feat: custom DNS server (#682)

* feat: custom DNS server
---
 main.go                  |  9 ++++++---
 util/dns.go              | 23 +++++++++++++++++++++++
 util/dns_test.go         | 16 ++++++++++++++++
 util/http_client_util.go |  2 ++
 4 files changed, 47 insertions(+), 3 deletions(-)
 create mode 100644 util/dns.go
 create mode 100644 util/dns_test.go

diff --git a/main.go b/main.go
index 37b3e1604..4b81f0d88 100644
--- a/main.go
+++ b/main.go
@@ -38,6 +38,9 @@ var noWebService = flag.Bool("noweb", false, "不启动 web 服务")
 // 跳过验证证书
 var skipVerify = flag.Bool("skipVerify", false, "跳过验证证书, 适合不能升级的老系统")
 
+// 自定义 DNS 服务器
+var customDNSServer = flag.String("dns", "", "自定义 DNS 服务器（例如 1.1.1.1）")
+
 //go:embed static
 var staticEmbededFiles embed.FS
 
@@ -47,9 +50,6 @@ var faviconEmbededFile embed.FS
 // version
 var version = "DEV"
 
-// buildTime
-var buildTime = ""
-
 func main() {
 	flag.Parse()
 	if _, err := net.ResolveTCPAddr("tcp", *listen); err != nil {
@@ -63,6 +63,9 @@ func main() {
 	if *skipVerify {
 		os.Setenv(util.SkipVerifyENV, "true")
 	}
+	if *customDNSServer != "" {
+		os.Setenv(util.DNSServerEnv, *customDNSServer+":53")
+	}
 	switch *serviceType {
 	case "install":
 		installService()
diff --git a/util/dns.go b/util/dns.go
new file mode 100644
index 000000000..4d88f5636
--- /dev/null
+++ b/util/dns.go
@@ -0,0 +1,23 @@
+package util
+
+import (
+	"context"
+	"net"
+	"os"
+)
+
+const DNSServerEnv = "DDNS_GO_DNS_SERVER"
+
+// customDNSResolver 当 DNSServerEnv 值不为空时，使用 Go 内置 DNS 解析器来解析其 DNS 服务器。
+func customDNSResolver() *net.Resolver {
+	s := os.Getenv(DNSServerEnv)
+	if s != "" {
+		return &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				return net.Dial("udp", s)
+			},
+		}
+	}
+	return &net.Resolver{}
+}
diff --git a/util/dns_test.go b/util/dns_test.go
new file mode 100644
index 000000000..50b5a12f0
--- /dev/null
+++ b/util/dns_test.go
@@ -0,0 +1,16 @@
+package util
+
+import (
+	"context"
+	"os"
+	"testing"
+)
+
+// TestCustomDNSResolver 测试能否通过 DNSServerEnv 值的 DNS 服务器解析域名的 IP。
+func TestCustomDNSResolver(t *testing.T) {
+	os.Setenv(DNSServerEnv, "1.1.1.1:53")
+	_, err := customDNSResolver().LookupIP(context.Background(), "ip", "cloudflare.com")
+	if err != nil {
+		t.Errorf("Failed to lookup IP, err: %v", err)
+	}
+}
diff --git a/util/http_client_util.go b/util/http_client_util.go
index 1425229e0..b607692ac 100644
--- a/util/http_client_util.go
+++ b/util/http_client_util.go
@@ -31,6 +31,7 @@ var defaultTransport = &http.Transport{
 
 // CreateHTTPClient Create Default HTTP Client
 func CreateHTTPClient() *http.Client {
+	dialer.Resolver = customDNSResolver()
 	// SkipVerfiry
 	defaultTransport.TLSClientConfig = &tls.Config{InsecureSkipVerify: os.Getenv(SkipVerifyENV) == "true"}
 	return &http.Client{
@@ -73,6 +74,7 @@ var noProxyTcp6Transport = &http.Transport{
 
 // CreateNoProxyHTTPClient Create NoProxy HTTP Client
 func CreateNoProxyHTTPClient(network string) *http.Client {
+	dialer.Resolver = customDNSResolver()
 	if network == "tcp6" {
 		// SkipVerfiry
 		noProxyTcp6Transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: os.Getenv(SkipVerifyENV) == "true"}