From 4801b6f05772f2217a913efae6c654acf9988fbc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=96=E7=95=8C?= <i@sekai.icu>
Date: Sun, 4 Sep 2022 12:39:43 +0800
Subject: [PATCH] Fix DNS routing

---
 outbound/dns.go     | 14 ++++----------
 route/router_dns.go |  3 +++
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/outbound/dns.go b/outbound/dns.go
index 83d96a82..81133e42 100644
--- a/outbound/dns.go
+++ b/outbound/dns.go
@@ -77,12 +77,9 @@ func (d *DNS) handleConnection(ctx context.Context, conn net.Conn, metadata adap
 	if err != nil {
 		return err
 	}
-	if len(message.Questions) > 0 {
-		question := message.Questions[0]
-		metadata.Domain = string(question.Name.Data[:question.Name.Length-1])
-	}
+	metadataInQuery := metadata
 	go func() error {
-		response, err := d.router.Exchange(ctx, &message)
+		response, err := d.router.Exchange(adapter.WithContext(ctx, &metadataInQuery), &message)
 		if err != nil {
 			return err
 		}
@@ -125,13 +122,10 @@ func (d *DNS) NewPacketConnection(ctx context.Context, conn N.PacketConn, metada
 			if err != nil {
 				return err
 			}
-			if len(message.Questions) > 0 {
-				question := message.Questions[0]
-				metadata.Domain = string(question.Name.Data[:question.Name.Length-1])
-			}
 			timeout.Update()
+			metadataInQuery := metadata
 			go func() error {
-				response, err := d.router.Exchange(ctx, &message)
+				response, err := d.router.Exchange(adapter.WithContext(ctx, &metadataInQuery), &message)
 				if err != nil {
 					return err
 				}
diff --git a/route/router_dns.go b/route/router_dns.go
index 795e5c79..373d5dc9 100644
--- a/route/router_dns.go
+++ b/route/router_dns.go
@@ -52,6 +52,7 @@ func (r *Router) Exchange(ctx context.Context, message *dnsmessage.Message) (*dn
 		case dnsmessage.TypeAAAA:
 			metadata.IPVersion = 6
 		}
+		metadata.Domain = string(message.Questions[0].Name.Data[:message.Questions[0].Name.Length-1])
 	}
 	ctx, transport, strategy := r.matchDNS(ctx)
 	ctx, cancel := context.WithTimeout(ctx, C.DNSTimeout)
@@ -68,6 +69,8 @@ func (r *Router) Exchange(ctx context.Context, message *dnsmessage.Message) (*dn
 
 func (r *Router) Lookup(ctx context.Context, domain string, strategy dns.DomainStrategy) ([]netip.Addr, error) {
 	r.dnsLogger.DebugContext(ctx, "lookup domain ", domain)
+	ctx, metadata := adapter.AppendContext(ctx)
+	metadata.Domain = domain
 	ctx, transport, transportStrategy := r.matchDNS(ctx)
 	if strategy == dns.DomainStrategyAsIS {
 		strategy = transportStrategy