From b745074449439478554dc02435353d6236b741b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Fri, 10 May 2024 20:16:27 +0800
Subject: [PATCH] Add ECS source netmark option

---
 extension_edns0_subnet.go | 10 +++++-----
 extensions.go             |  6 +++---
 transport.go              |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/extension_edns0_subnet.go b/extension_edns0_subnet.go
index 7481698..a2c3935 100644
--- a/extension_edns0_subnet.go
+++ b/extension_edns0_subnet.go
@@ -9,7 +9,7 @@ import (
 
 type edns0SubnetTransportWrapper struct {
 	Transport
-	clientSubnet netip.Addr
+	clientSubnet netip.Prefix
 }
 
 func (t *edns0SubnetTransportWrapper) Exchange(ctx context.Context, message *dns.Msg) (*dns.Msg, error) {
@@ -17,7 +17,7 @@ func (t *edns0SubnetTransportWrapper) Exchange(ctx context.Context, message *dns
 	return t.Transport.Exchange(ctx, message)
 }
 
-func SetClientSubnet(message *dns.Msg, clientSubnet netip.Addr, override bool) {
+func SetClientSubnet(message *dns.Msg, clientSubnet netip.Prefix, override bool) {
 	var (
 		optRecord    *dns.OPT
 		subnetOption *dns.EDNS0_SUBNET
@@ -52,11 +52,11 @@ findExists:
 		optRecord.Option = append(optRecord.Option, subnetOption)
 	}
 	subnetOption.Code = dns.EDNS0SUBNET
-	if clientSubnet.Is4() {
+	if clientSubnet.Addr().Is4() {
 		subnetOption.Family = 1
 	} else {
 		subnetOption.Family = 2
 	}
-	subnetOption.SourceNetmask = uint8(clientSubnet.BitLen())
-	subnetOption.Address = clientSubnet.AsSlice()
+	subnetOption.SourceNetmask = uint8(clientSubnet.Bits())
+	subnetOption.Address = clientSubnet.Addr().AsSlice()
 }
diff --git a/extensions.go b/extensions.go
index 6b89562..86bbbd5 100644
--- a/extensions.go
+++ b/extensions.go
@@ -46,11 +46,11 @@ func transportNameFromContext(ctx context.Context) (string, bool) {
 
 type clientSubnetKey struct{}
 
-func ContextWithClientSubnet(ctx context.Context, clientSubnet netip.Addr) context.Context {
+func ContextWithClientSubnet(ctx context.Context, clientSubnet netip.Prefix) context.Context {
 	return context.WithValue(ctx, clientSubnetKey{}, clientSubnet)
 }
 
-func ClientSubnetFromContext(ctx context.Context) (netip.Addr, bool) {
-	clientSubnet, ok := ctx.Value(clientSubnetKey{}).(netip.Addr)
+func ClientSubnetFromContext(ctx context.Context) (netip.Prefix, bool) {
+	clientSubnet, ok := ctx.Value(clientSubnetKey{}).(netip.Prefix)
 	return clientSubnet, ok
 }
diff --git a/transport.go b/transport.go
index 29d5c58..b3edd0b 100644
--- a/transport.go
+++ b/transport.go
@@ -30,7 +30,7 @@ type TransportOptions struct {
 	Name         string
 	Dialer       N.Dialer
 	Address      string
-	ClientSubnet netip.Addr
+	ClientSubnet netip.Prefix
 }
 
 var transports map[string]TransportConstructor