From 0499b5cb02ee565b2da252cce8335617a1556d7c Mon Sep 17 00:00:00 2001
From: rui0572 <125641819+rui0572@users.noreply.github.com>
Date: Sun, 7 May 2023 01:04:47 +0800
Subject: [PATCH] =?UTF-8?q?=E5=BE=AE=E8=B0=83.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 proxy/wireguard/config.go    | 25 +++++++++++++++++++++++++
 proxy/wireguard/wireguard.go | 15 +++++----------
 2 files changed, 30 insertions(+), 10 deletions(-)
 create mode 100644 proxy/wireguard/config.go

diff --git a/proxy/wireguard/config.go b/proxy/wireguard/config.go
new file mode 100644
index 000000000000..756227539991
--- /dev/null
+++ b/proxy/wireguard/config.go
@@ -0,0 +1,25 @@
+package wireguard
+
+func (c *DeviceConfig) preferIP4() bool {
+	return c.DomainStrategy == DeviceConfig_FORCE_IP ||
+		c.DomainStrategy == DeviceConfig_FORCE_IP4 ||
+		c.DomainStrategy == DeviceConfig_FORCE_IP46
+}
+
+func (c *DeviceConfig) preferIP6() bool {
+	return c.DomainStrategy == DeviceConfig_FORCE_IP ||
+		c.DomainStrategy == DeviceConfig_FORCE_IP6 ||
+		c.DomainStrategy == DeviceConfig_FORCE_IP64
+}
+
+func (c *DeviceConfig) hasFallback() bool {
+	return c.DomainStrategy == DeviceConfig_FORCE_IP46 || c.DomainStrategy == DeviceConfig_FORCE_IP64
+}
+
+func (c *DeviceConfig) fallbackIP4() bool {
+	return c.DomainStrategy == DeviceConfig_FORCE_IP64
+}
+
+func (c *DeviceConfig) fallbackIP6() bool {
+	return c.DomainStrategy == DeviceConfig_FORCE_IP46
+}
diff --git a/proxy/wireguard/wireguard.go b/proxy/wireguard/wireguard.go
index b2a363216146..48e2ace38423 100644
--- a/proxy/wireguard/wireguard.go
+++ b/proxy/wireguard/wireguard.go
@@ -160,19 +160,14 @@ func (h *Handler) Process(ctx context.Context, link *transport.Link, dialer inte
 	addr := destination.Address
 	if addr.Family().IsDomain() {
 		ips, err := h.dns.LookupIP(addr.Domain(), dns.IPOption{
-			IPv4Enable: h.hasIPv4 && (h.conf.DomainStrategy == DeviceConfig_FORCE_IP ||
-				h.conf.DomainStrategy == DeviceConfig_FORCE_IP4 ||
-				h.conf.DomainStrategy == DeviceConfig_FORCE_IP46),
-
-			IPv6Enable: h.hasIPv6 && (h.conf.DomainStrategy == DeviceConfig_FORCE_IP ||
-				h.conf.DomainStrategy == DeviceConfig_FORCE_IP6 ||
-				h.conf.DomainStrategy == DeviceConfig_FORCE_IP64),
+			IPv4Enable: h.hasIPv4 && h.conf.preferIP4(),
+			IPv6Enable: h.hasIPv6 && h.conf.preferIP6(),
 		})
 		{ // Resolve fallback
-			if (len(ips) == 0 || err != nil) && (h.conf.DomainStrategy == DeviceConfig_FORCE_IP46 || h.conf.DomainStrategy == DeviceConfig_FORCE_IP64) {
+			if (len(ips) == 0 || err != nil) && h.conf.hasFallback() {
 				ips, err = h.dns.LookupIP(addr.Domain(), dns.IPOption{
-					IPv4Enable: h.hasIPv4 && h.conf.DomainStrategy != DeviceConfig_FORCE_IP46,
-					IPv6Enable: h.hasIPv6 && h.conf.DomainStrategy != DeviceConfig_FORCE_IP64,
+					IPv4Enable: h.hasIPv4 && h.conf.fallbackIP4(),
+					IPv6Enable: h.hasIPv6 && h.conf.fallbackIP6(),
 				})
 			}
 		}