From 6b0e861afa09d7f2d2c3937d07d460409971da11 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Wed, 7 May 2025 13:27:53 +0800
Subject: [PATCH] Fix set edns0 client subnet

---
 dns/client.go                 |  2 +-
 dns/extension_edns0_subnet.go | 15 ++++++++-------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/dns/client.go b/dns/client.go
index 5d7f109d..0c5490f9 100644
--- a/dns/client.go
+++ b/dns/client.go
@@ -105,7 +105,7 @@ func (c *Client) Exchange(ctx context.Context, transport adapter.DNSTransport, m
 	}
 	question := message.Question[0]
 	if options.ClientSubnet.IsValid() {
-		message = SetClientSubnet(message, options.ClientSubnet, true)
+		message = SetClientSubnet(message, options.ClientSubnet)
 	}
 	isSimpleRequest := len(message.Question) == 1 &&
 		len(message.Ns) == 0 &&
diff --git a/dns/extension_edns0_subnet.go b/dns/extension_edns0_subnet.go
index 1c4033d3..e804fb6c 100644
--- a/dns/extension_edns0_subnet.go
+++ b/dns/extension_edns0_subnet.go
@@ -6,7 +6,11 @@ import (
 	"github.com/miekg/dns"
 )
 
-func SetClientSubnet(message *dns.Msg, clientSubnet netip.Prefix, override bool) *dns.Msg {
+func SetClientSubnet(message *dns.Msg, clientSubnet netip.Prefix) *dns.Msg {
+	return setClientSubnet(message, clientSubnet, true)
+}
+
+func setClientSubnet(message *dns.Msg, clientSubnet netip.Prefix, clone bool) *dns.Msg {
 	var (
 		optRecord    *dns.OPT
 		subnetOption *dns.EDNS0_SUBNET
@@ -19,9 +23,6 @@ findExists:
 				var isEDNS0Subnet bool
 				subnetOption, isEDNS0Subnet = option.(*dns.EDNS0_SUBNET)
 				if isEDNS0Subnet {
-					if !override {
-						return message
-					}
 					break findExists
 				}
 			}
@@ -37,14 +38,14 @@ findExists:
 			},
 		}
 		message.Extra = append(message.Extra, optRecord)
-	} else {
-		message = message.Copy()
+	} else if clone {
+		return setClientSubnet(message.Copy(), clientSubnet, false)
 	}
 	if subnetOption == nil {
 		subnetOption = new(dns.EDNS0_SUBNET)
+		subnetOption.Code = dns.EDNS0SUBNET
 		optRecord.Option = append(optRecord.Option, subnetOption)
 	}
-	subnetOption.Code = dns.EDNS0SUBNET
 	if clientSubnet.Addr().Is4() {
 		subnetOption.Family = 1
 	} else {