Refactor UDP DNS hijack

route/dns.go

@@ -10,12 +10,9 @@ import (
 	"github.com/sagernet/sing-box/dns"
 	dnsOutbound "github.com/sagernet/sing-box/protocol/dns"
 	R "github.com/sagernet/sing-box/route/rule"
-	"github.com/sagernet/sing/common/buf"
 	E "github.com/sagernet/sing/common/exceptions"
-	"github.com/sagernet/sing/common/logger"
 	M "github.com/sagernet/sing/common/metadata"
 	N "github.com/sagernet/sing/common/network"
-	"github.com/sagernet/sing/common/udpnat2"
 
 	mDNS "github.com/miekg/dns"
 )
@@ -36,24 +33,6 @@ func (r *Router) hijackDNSStream(ctx context.Context, conn net.Conn, metadata ad
 }
 
 func (r *Router) hijackDNSPacket(ctx context.Context, conn N.PacketConn, packetBuffers []*N.PacketBuffer, metadata adapter.InboundContext, onClose N.CloseHandlerFunc) error {
-	if natConn, isNatConn := conn.(udpnat.Conn); isNatConn {
-		metadata.Destination = M.Socksaddr{}
-		for _, packet := range packetBuffers {
-			buffer := packet.Buffer
-			destination := packet.Destination
-			N.PutPacketBuffer(packet)
-			go ExchangeDNSPacket(ctx, r.dns, r.logger, natConn, buffer, metadata, destination)
-		}
-		natConn.SetHandler(&dnsHijacker{
-			router:   r.dns,
-			logger:   r.logger,
-			conn:     conn,
-			ctx:      ctx,
-			metadata: metadata,
-			onClose:  onClose,
-		})
-		return nil
-	}
 	err := dnsOutbound.NewDNSPacketConnection(ctx, r.dns, conn, packetBuffers, metadata)
 	N.CloseOnHandshakeFailure(conn, onClose, err)
 	if err != nil && !E.IsClosedOrCanceled(err) {
@@ -62,48 +41,31 @@ func (r *Router) hijackDNSPacket(ctx context.Context, conn N.PacketConn, packetB
 	return nil
 }
 
-func ExchangeDNSPacket(ctx context.Context, router adapter.DNSRouter, logger logger.ContextLogger, conn N.PacketConn, buffer *buf.Buffer, metadata adapter.InboundContext, destination M.Socksaddr) {
-	err := exchangeDNSPacket(ctx, router, conn, buffer, metadata, destination)
-	if err != nil && !R.IsRejected(err) && !E.IsClosedOrCanceled(err) {
-		logger.ErrorContext(ctx, E.Cause(err, "process DNS packet"))
-	}
-}
-
-func exchangeDNSPacket(ctx context.Context, router adapter.DNSRouter, conn N.PacketConn, buffer *buf.Buffer, metadata adapter.InboundContext, destination M.Socksaddr) error {
+func (r *Router) HijackDNSPacket(ctx context.Context, payload []byte, writer N.PacketWriter, metadata adapter.InboundContext) {
 	var message mDNS.Msg
-	err := message.Unpack(buffer.Bytes())
-	buffer.Release()
+	err := message.Unpack(payload)
 	if err != nil {
-		return E.Cause(err, "unpack request")
+		r.logger.ErrorContext(ctx, E.Cause(err, "process DNS packet: unpack request"))
+		return
 	}
-	response, err := router.Exchange(adapter.WithContext(ctx, &metadata), &message, adapter.DNSQueryOptions{})
+	destination := metadata.Destination
+	metadata.Destination = M.Socksaddr{}
+	go func() {
+		exchangeErr := r.exchangeDNSPacket(ctx, &message, writer, metadata, destination)
+		if exchangeErr != nil && !R.IsRejected(exchangeErr) && !E.IsClosedOrCanceled(exchangeErr) {
+			r.logger.ErrorContext(ctx, E.Cause(exchangeErr, "process DNS packet"))
+		}
+	}()
+}
+
+func (r *Router) exchangeDNSPacket(ctx context.Context, message *mDNS.Msg, writer N.PacketWriter, metadata adapter.InboundContext, destination M.Socksaddr) error {
+	response, err := r.dns.Exchange(adapter.WithContext(ctx, &metadata), message, adapter.DNSQueryOptions{})
 	if err != nil {
 		return err
 	}
-	responseBuffer, err := dns.TruncateDNSMessage(&message, response, 1024)
+	responseBuffer, err := dns.TruncateDNSMessage(message, response, 1024)
 	if err != nil {
 		return err
 	}
-	err = conn.WritePacket(responseBuffer, destination)
-	return err
-}
-
-type dnsHijacker struct {
-	router   adapter.DNSRouter
-	logger   logger.ContextLogger
-	conn     N.PacketConn
-	ctx      context.Context
-	metadata adapter.InboundContext
-	onClose  N.CloseHandlerFunc
-}
-
-func (h *dnsHijacker) NewPacketEx(buffer *buf.Buffer, destination M.Socksaddr) {
-	go ExchangeDNSPacket(h.ctx, h.router, h.logger, h.conn, buffer, h.metadata, destination)
-}
-
-func (h *dnsHijacker) Close() error {
-	if h.onClose != nil {
-		h.onClose(nil)
-	}
-	return nil
+	return writer.WritePacket(responseBuffer, destination)
 }

route/route.go

@@ -391,6 +391,11 @@ func (r *Router) PreMatch(metadata adapter.InboundContext, firstPacket []byte) a
 				return adapter.PreMatchResult{Action: adapter.PreMatchDrop}
 			}
 			return adapter.PreMatchResult{Action: adapter.PreMatchReject}
+		case *R.RuleActionHijackDNS:
+			if metadata.Network != N.NetworkUDP {
+				return continueResult
+			}
+			return adapter.PreMatchResult{Action: adapter.PreMatchHijackDNS}
 		case *R.RuleActionResolve:
 			resolveErr := r.actionResolve(adapter.WithContext(ctx, &metadata), &metadata, action)
 			if resolveErr != nil {

route/rule/rule_item_rule_set_test.go

@@ -24,6 +24,9 @@ func (r *ruleSetItemTestRouter) PreMatch(adapter.InboundContext, []byte) adapter
 	return adapter.PreMatchResult{}
 }
 
+func (r *ruleSetItemTestRouter) HijackDNSPacket(context.Context, []byte, N.PacketWriter, adapter.InboundContext) {
+}
+
 func (r *ruleSetItemTestRouter) RouteConnection(context.Context, net.Conn, adapter.InboundContext) error {
 	return nil
 }