From ecbbf9d2201d4e4c3d1f2f11fa91fb19878f94d0 Mon Sep 17 00:00:00 2001
From: wwqgtxx <wwqgtxx@gmail.com>
Date: Wed, 11 Sep 2024 10:28:34 +0800
Subject: [PATCH] feat: doh client support `ecs` and `ecs-override`

---
 dns/doh.go          | 29 ++++++++++++++++++++++++++
 dns/edns0_subnet.go | 51 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+)
 create mode 100644 dns/edns0_subnet.go

diff --git a/dns/doh.go b/dns/doh.go
index 97e01ea720..ffb65fcef0 100644
--- a/dns/doh.go
+++ b/dns/doh.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"net/netip"
 	"net/url"
 	"runtime"
 	"strconv"
@@ -67,6 +68,8 @@ type dnsOverHTTPS struct {
 	dialer         *dnsDialer
 	addr           string
 	skipCertVerify bool
+	ecsPrefix      netip.Prefix
+	ecsOverride    bool
 }
 
 // type check
@@ -99,6 +102,28 @@ func newDoHClient(urlString string, r *Resolver, preferH3 bool, params map[strin
 		doh.skipCertVerify = true
 	}
 
+	if ecs := params["ecs"]; ecs != "" {
+		prefix, err := netip.ParsePrefix(ecs)
+		if err != nil {
+			addr, err := netip.ParseAddr(ecs)
+			if err != nil {
+				log.Warnln("DOH [%s] config with invalid ecs: %s", doh.addr, ecs)
+			} else {
+				doh.ecsPrefix = netip.PrefixFrom(addr, addr.BitLen())
+			}
+		} else {
+			doh.ecsPrefix = prefix
+		}
+	}
+
+	if doh.ecsPrefix.IsValid() {
+		log.Debugln("DOH [%s] config with ecs: %s", doh.addr, doh.ecsPrefix)
+	}
+
+	if params["ecs-override"] == "true" {
+		doh.ecsOverride = true
+	}
+
 	runtime.SetFinalizer(doh, (*dnsOverHTTPS).Close)
 
 	return doh
@@ -126,6 +151,10 @@ func (doh *dnsOverHTTPS) ExchangeContext(ctx context.Context, m *D.Msg) (msg *D.
 		}
 	}()
 
+	if doh.ecsPrefix.IsValid() {
+		setEdns0Subnet(m, doh.ecsPrefix, doh.ecsOverride)
+	}
+
 	// Check if there was already an active client before sending the request.
 	// We'll only attempt to re-connect if there was one.
 	client, isCached, err := doh.getClient(ctx)
diff --git a/dns/edns0_subnet.go b/dns/edns0_subnet.go
new file mode 100644
index 0000000000..2ed4f140b5
--- /dev/null
+++ b/dns/edns0_subnet.go
@@ -0,0 +1,51 @@
+package dns
+
+import (
+	"net/netip"
+
+	"github.com/miekg/dns"
+)
+
+func setEdns0Subnet(message *dns.Msg, clientSubnet netip.Prefix, override bool) bool {
+	var (
+		optRecord    *dns.OPT
+		subnetOption *dns.EDNS0_SUBNET
+	)
+findExists:
+	for _, record := range message.Extra {
+		var isOPTRecord bool
+		if optRecord, isOPTRecord = record.(*dns.OPT); isOPTRecord {
+			for _, option := range optRecord.Option {
+				var isEDNS0Subnet bool
+				if subnetOption, isEDNS0Subnet = option.(*dns.EDNS0_SUBNET); isEDNS0Subnet {
+					if !override {
+						return false
+					}
+					break findExists
+				}
+			}
+		}
+	}
+	if optRecord == nil {
+		optRecord = &dns.OPT{
+			Hdr: dns.RR_Header{
+				Name:   ".",
+				Rrtype: dns.TypeOPT,
+			},
+		}
+		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.Addr().Is4() {
+		subnetOption.Family = 1
+	} else {
+		subnetOption.Family = 2
+	}
+	subnetOption.SourceNetmask = uint8(clientSubnet.Bits())
+	subnetOption.Address = clientSubnet.Addr().AsSlice()
+	return true
+}