Refactor and add support for the client-subnet option

client.go

@@ -117,6 +117,10 @@ func (c *Client) ExchangeWithResponseCheck(ctx context.Context, transport Transp
 		return nil, E.New("DNS query loopback in transport[", contextTransport, "]")
 	}
 	ctx = contextWithTransportName(ctx, transport.Name())
+	clientSubnet, loaded := ClientSubnetFromContext(ctx)
+	if loaded {
+		SetClientSubnet(message, clientSubnet, true)
+	}
 	response, err := transport.Exchange(ctx, message)
 	if err != nil {
 		return nil, err

extension_edns0_subnet.go

@@ -0,0 +1,55 @@
+package dns
+
+import (
+	"context"
+	"net/netip"
+
+	"github.com/miekg/dns"
+)
+
+type edns0SubnetTransportWrapper struct {
+	Transport
+	clientSubnet netip.Addr
+}
+
+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) {
+	var subnetOption *dns.EDNS0_SUBNET
+findExists:
+	for _, record := range message.Extra {
+		if optRecord, isOPTRecord := record.(*dns.OPT); isOPTRecord {
+			for _, option := range optRecord.Option {
+				var isEDNS0Subnet bool
+				subnetOption, isEDNS0Subnet = option.(*dns.EDNS0_SUBNET)
+				if isEDNS0Subnet {
+					if !override {
+						return
+					}
+					break findExists
+				}
+			}
+		}
+	}
+	if subnetOption == nil {
+		subnetOption = new(dns.EDNS0_SUBNET)
+		message.Extra = append(message.Extra, &dns.OPT{
+			Hdr: dns.RR_Header{
+				Name:   ".",
+				Rrtype: dns.TypeOPT,
+			},
+			Option: []dns.EDNS0{subnetOption},
+		})
+	}
+	subnetOption.Code = dns.EDNS0SUBNET
+	if clientSubnet.Is4() {
+		subnetOption.Family = 1
+	} else {
+		subnetOption.Family = 2
+	}
+	subnetOption.SourceNetmask = uint8(clientSubnet.BitLen())
+	subnetOption.Address = clientSubnet.AsSlice()
+}

extensions.go

@@ -0,0 +1,56 @@
+package dns
+
+import (
+	"context"
+	"net/netip"
+)
+
+type disableCacheKey struct{}
+
+func ContextWithDisableCache(ctx context.Context, val bool) context.Context {
+	return context.WithValue(ctx, (*disableCacheKey)(nil), val)
+}
+
+func DisableCacheFromContext(ctx context.Context) bool {
+	val := ctx.Value((*disableCacheKey)(nil))
+	if val == nil {
+		return false
+	}
+	return val.(bool)
+}
+
+type rewriteTTLKey struct{}
+
+func ContextWithRewriteTTL(ctx context.Context, val uint32) context.Context {
+	return context.WithValue(ctx, (*rewriteTTLKey)(nil), val)
+}
+
+func RewriteTTLFromContext(ctx context.Context) (uint32, bool) {
+	val := ctx.Value((*rewriteTTLKey)(nil))
+	if val == nil {
+		return 0, false
+	}
+	return val.(uint32), true
+}
+
+type transportKey struct{}
+
+func contextWithTransportName(ctx context.Context, transportName string) context.Context {
+	return context.WithValue(ctx, transportKey{}, transportName)
+}
+
+func transportNameFromContext(ctx context.Context) (string, bool) {
+	value, loaded := ctx.Value(transportKey{}).(string)
+	return value, loaded
+}
+
+type clientSubnetKey struct{}
+
+func ContextWithClientSubnet(ctx context.Context, clientSubnet netip.Addr) context.Context {
+	return context.WithValue(ctx, clientSubnetKey{}, clientSubnet)
+}
+
+func ClientSubnetFromContext(ctx context.Context) (netip.Addr, bool) {
+	clientSubnet, ok := ctx.Value(clientSubnetKey{}).(netip.Addr)
+	return clientSubnet, ok
+}

quic/transport_http3.go

@@ -14,7 +14,6 @@ import (
 	"github.com/sagernet/quic-go/http3"
 	"github.com/sagernet/sing-dns"
 	"github.com/sagernet/sing/common/bufio"
-	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
 
@@ -24,16 +23,9 @@ import (
 var _ dns.Transport = (*HTTP3Transport)(nil)
 
 func init() {
-	dns.RegisterTransport([]string{"h3"}, CreateHTTP3Transport)
-}
-
-func CreateHTTP3Transport(name string, ctx context.Context, logger logger.ContextLogger, dialer N.Dialer, link string) (dns.Transport, error) {
-	linkURL, err := url.Parse(link)
-	if err != nil {
-		return nil, err
-	}
-	linkURL.Scheme = "https"
-	return NewHTTP3Transport(name, dialer, linkURL.String()), nil
+	dns.RegisterTransport([]string{"h3"}, func(options dns.TransportOptions) (dns.Transport, error) {
+		return NewHTTP3Transport(options)
+	})
 }
 
 type HTTP3Transport struct {
@@ -42,24 +34,29 @@ type HTTP3Transport struct {
 	transport   *http3.RoundTripper
 }
 
-func NewHTTP3Transport(name string, dialer N.Dialer, serverURL string) *HTTP3Transport {
+func NewHTTP3Transport(options dns.TransportOptions) (*HTTP3Transport, error) {
+	serverURL, err := url.Parse(options.Address)
+	if err != nil {
+		return nil, err
+	}
+	serverURL.Scheme = "https"
 	return &HTTP3Transport{
-		name:        name,
-		destination: serverURL,
+		name:        options.Name,
+		destination: options.Address,
 		transport: &http3.RoundTripper{
 			Dial: func(ctx context.Context, addr string, tlsCfg *tls.Config, cfg *quic.Config) (quic.EarlyConnection, error) {
 				destinationAddr := M.ParseSocksaddr(addr)
-				conn, err := dialer.DialContext(ctx, N.NetworkUDP, destinationAddr)
-				if err != nil {
-					return nil, err
+				conn, dialErr := options.Dialer.DialContext(ctx, N.NetworkUDP, destinationAddr)
+				if dialErr != nil {
+					return nil, dialErr
 				}
 				return quic.DialEarly(ctx, bufio.NewUnbindPacketConn(conn), conn.RemoteAddr(), tlsCfg, cfg)
 			},
 			TLSClientConfig: &tls.Config{
 				NextProtos: []string{"dns"},
 			},
 		},
-	}
+	}, nil
 }
 
 func (t *HTTP3Transport) Name() string {

transport.go

@@ -12,7 +12,7 @@ import (
 	"github.com/miekg/dns"
 )
 
-type TransportConstructor = func(name string, ctx context.Context, logger logger.ContextLogger, dialer N.Dialer, link string) (Transport, error)
+type TransportConstructor = func(options TransportOptions) (Transport, error)
 
 type Transport interface {
 	Name() string
@@ -24,6 +24,15 @@ type Transport interface {
 	Lookup(ctx context.Context, domain string, strategy DomainStrategy) ([]netip.Addr, error)
 }
 
+type TransportOptions struct {
+	Context      context.Context
+	Logger       logger.ContextLogger
+	Name         string
+	Dialer       N.Dialer
+	Address      string
+	ClientSubnet netip.Addr
+}
+
 var transports map[string]TransportConstructor
 
 func RegisterTransport(schemes []string, constructor TransportConstructor) {
@@ -35,18 +44,26 @@ func RegisterTransport(schemes []string, constructor TransportConstructor) {
 	}
 }
 
-func CreateTransport(name string, ctx context.Context, logger logger.ContextLogger, dialer N.Dialer, address string) (Transport, error) {
-	constructor := transports[address]
+func CreateTransport(options TransportOptions) (Transport, error) {
+	constructor := transports[options.Address]
 	if constructor == nil {
-		serverURL, _ := url.Parse(address)
+		serverURL, _ := url.Parse(options.Address)
 		var scheme string
 		if serverURL != nil {
 			scheme = serverURL.Scheme
 		}
 		constructor = transports[scheme]
 	}
 	if constructor == nil {
-		return nil, E.New("unknown DNS server format: " + address)
+		return nil, E.New("unknown DNS server format: " + options.Address)
+	}
+	options.Context = contextWithTransportName(options.Context, options.Name)
+	transport, err := constructor(options)
+	if err != nil {
+		return nil, err
+	}
+	if options.ClientSubnet.IsValid() {
+		transport = &edns0SubnetTransportWrapper{transport, options.ClientSubnet}
 	}
-	return constructor(name, contextWithTransportName(ctx, name), logger, dialer, address)
+	return transport, nil
 }

transport_base.go

@@ -28,19 +28,21 @@ type myTransportAdapter struct {
 	cancel     context.CancelFunc
 	dialer     N.Dialer
 	serverAddr M.Socksaddr
+	clientAddr netip.Addr
 	handler    myTransportHandler
 	access     sync.Mutex
 	conn       *dnsConnection
 }
 
-func newAdapter(name string, ctx context.Context, dialer N.Dialer, serverAddr M.Socksaddr) myTransportAdapter {
-	ctx, cancel := context.WithCancel(ctx)
+func newAdapter(options TransportOptions, serverAddr M.Socksaddr) myTransportAdapter {
+	ctx, cancel := context.WithCancel(options.Context)
 	return myTransportAdapter{
-		name:       name,
+		name:       options.Name,
 		ctx:        ctx,
 		cancel:     cancel,
-		dialer:     dialer,
+		dialer:     options.Dialer,
 		serverAddr: serverAddr,
+		clientAddr: options.ClientSubnet,
 	}
 }
 

transport_https.go

@@ -10,9 +10,7 @@ import (
 	"net/netip"
 	"os"
 
-	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
-	N "github.com/sagernet/sing/common/network"
 
 	"github.com/miekg/dns"
 )
@@ -28,21 +26,19 @@ type HTTPSTransport struct {
 }
 
 func init() {
-	RegisterTransport([]string{"https"}, CreateHTTPSTransport)
+	RegisterTransport([]string{"https"}, func(options TransportOptions) (Transport, error) {
+		return NewHTTPSTransport(options), nil
+	})
 }
 
-func CreateHTTPSTransport(name string, ctx context.Context, logger logger.ContextLogger, dialer N.Dialer, link string) (Transport, error) {
-	return NewHTTPSTransport(name, dialer, link), nil
-}
-
-func NewHTTPSTransport(name string, dialer N.Dialer, serverURL string) *HTTPSTransport {
+func NewHTTPSTransport(options TransportOptions) *HTTPSTransport {
 	return &HTTPSTransport{
-		name:        name,
-		destination: serverURL,
+		name:        options.Name,
+		destination: options.Address,
 		transport: &http.Transport{
 			ForceAttemptHTTP2: true,
 			DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
-				return dialer.DialContext(ctx, network, M.ParseSocksaddr(addr))
+				return options.Dialer.DialContext(ctx, network, M.ParseSocksaddr(addr))
 			},
 			TLSClientConfig: &tls.Config{
 				NextProtos: []string{"dns"},

transport_local.go

@@ -8,19 +8,16 @@ import (
 	"sort"
 
 	"github.com/sagernet/sing/common"
-	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
 
 	"github.com/miekg/dns"
 )
 
 func init() {
-	RegisterTransport([]string{"local"}, CreateLocalTransport)
-}
-
-func CreateLocalTransport(name string, ctx context.Context, logger logger.ContextLogger, dialer N.Dialer, link string) (Transport, error) {
-	return NewLocalTransport(name, dialer), nil
+	RegisterTransport([]string{"local"}, func(options TransportOptions) (Transport, error) {
+		return NewLocalTransport(options), nil
+	})
 }
 
 var _ Transport = (*LocalTransport)(nil)
@@ -30,12 +27,12 @@ type LocalTransport struct {
 	resolver net.Resolver
 }
 
-func NewLocalTransport(name string, dialer N.Dialer) *LocalTransport {
+func NewLocalTransport(options TransportOptions) *LocalTransport {
 	return &LocalTransport{
-		name: name,
+		name: options.Name,
 		resolver: net.Resolver{
 			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-				return dialer.DialContext(ctx, N.NetworkName(network), M.ParseSocksaddr(address))
+				return options.Dialer.DialContext(ctx, N.NetworkName(network), M.ParseSocksaddr(address))
 			},
 		},
 	}