Add ECS source netmark option

SagerNet · May 10, 2024 · f52ce53 · f52ce53
1 parent 9fb3613
commit f52ce53
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 10 deletions.
diff --git a/extension_edns0_subnet.go b/extension_edns0_subnet.go
@@ -9,15 +9,15 @@ import (
 
 type edns0SubnetTransportWrapper struct {
 	Transport
-	clientSubnet netip.Addr
+	clientSubnet netip.Prefix
 }
 
 func (t *edns0SubnetTransportWrapper) Exchange(ctx context.Context, message *dns.Msg) (*dns.Msg, error) {
 	SetClientSubnet(message, t.clientSubnet, false)
 	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
@@ -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
@@ -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

diff --git a/transport_base.go b/transport_base.go
@@ -28,7 +28,7 @@ type myTransportAdapter struct {
 	cancel     context.CancelFunc
 	dialer     N.Dialer
 	serverAddr M.Socksaddr
-	clientAddr netip.Addr
+	clientAddr netip.Prefix
 	handler    myTransportHandler
 	access     sync.Mutex
 	conn       *dnsConnection