From 29c23c4ba3929791e6f230328299874088ec0a14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Thu, 29 Feb 2024 12:46:26 +0800
Subject: [PATCH] Improve set client subnet

---
 extension_edns0_subnet.go | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/extension_edns0_subnet.go b/extension_edns0_subnet.go
index fcbf370..7481698 100644
--- a/extension_edns0_subnet.go
+++ b/extension_edns0_subnet.go
@@ -18,10 +18,14 @@ func (t *edns0SubnetTransportWrapper) Exchange(ctx context.Context, message *dns
 }
 
 func SetClientSubnet(message *dns.Msg, clientSubnet netip.Addr, override bool) {
-	var subnetOption *dns.EDNS0_SUBNET
+	var (
+		optRecord    *dns.OPT
+		subnetOption *dns.EDNS0_SUBNET
+	)
 findExists:
 	for _, record := range message.Extra {
-		if optRecord, isOPTRecord := record.(*dns.OPT); isOPTRecord {
+		var isOPTRecord bool
+		if optRecord, isOPTRecord = record.(*dns.OPT); isOPTRecord {
 			for _, option := range optRecord.Option {
 				var isEDNS0Subnet bool
 				subnetOption, isEDNS0Subnet = option.(*dns.EDNS0_SUBNET)
@@ -34,15 +38,18 @@ findExists:
 			}
 		}
 	}
-	if subnetOption == nil {
-		subnetOption = new(dns.EDNS0_SUBNET)
-		message.Extra = append(message.Extra, &dns.OPT{
+	if optRecord == nil {
+		optRecord = &dns.OPT{
 			Hdr: dns.RR_Header{
 				Name:   ".",
 				Rrtype: dns.TypeOPT,
 			},
-			Option: []dns.EDNS0{subnetOption},
-		})
+		}
+		message.Extra = append(message.Extra, optRecord)
+	}
+	if subnetOption == nil {
+		subnetOption = new(dns.EDNS0_SUBNET)
+		optRecord.Option = append(optRecord.Option, subnetOption)
 	}
 	subnetOption.Code = dns.EDNS0SUBNET
 	if clientSubnet.Is4() {