From dcf97ff5b4ccd088ed83889ae2bc0b016b8cda78 Mon Sep 17 00:00:00 2001
From: Dreamacro <305009791@qq.com>
Date: Fri, 7 Feb 2020 20:53:43 +0800
Subject: [PATCH] Fix: should prehandle metadata before resolve

---
 constant/metadata.go |  4 ++++
 tunnel/tunnel.go     | 47 ++++++++++++++++++++++++++++++++------------
 2 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/constant/metadata.go b/constant/metadata.go
index 78ac9580..61ce4343 100644
--- a/constant/metadata.go
+++ b/constant/metadata.go
@@ -75,6 +75,10 @@ func (m *Metadata) SourceAddress() string {
 	return net.JoinHostPort(m.SrcIP.String(), m.SrcPort)
 }
 
+func (m *Metadata) Resolved() bool {
+	return m.DstIP != nil
+}
+
 func (m *Metadata) UDPAddr() *net.UDPAddr {
 	if m.NetWork != UDP || m.DstIP == nil {
 		return nil
diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
index 061a82e5..f572399a 100644
--- a/tunnel/tunnel.go
+++ b/tunnel/tunnel.go
@@ -132,8 +132,8 @@ func (t *Tunnel) needLookupIP(metadata *C.Metadata) bool {
 	return dns.DefaultResolver != nil && (dns.DefaultResolver.IsMapping() || dns.DefaultResolver.FakeIPEnabled()) && metadata.Host == "" && metadata.DstIP != nil
 }
 
-func (t *Tunnel) resolveMetadata(metadata *C.Metadata) (C.Proxy, C.Rule, error) {
-	// handle host equal IP string
+func (t *Tunnel) preHandleMetadata(metadata *C.Metadata) error {
+	// handle IP string on host
 	if ip := net.ParseIP(metadata.Host); ip != nil {
 		metadata.DstIP = ip
 	}
@@ -147,9 +147,15 @@ func (t *Tunnel) resolveMetadata(metadata *C.Metadata) (C.Proxy, C.Rule, error)
 			if dns.DefaultResolver.FakeIPEnabled() {
 				metadata.DstIP = nil
 			}
+		} else if dns.DefaultResolver.IsFakeIP(metadata.DstIP) {
+			return fmt.Errorf("fake DNS record %s missing", metadata.DstIP)
 		}
 	}
 
+	return nil
+}
+
+func (t *Tunnel) resolveMetadata(metadata *C.Metadata) (C.Proxy, C.Rule, error) {
 	var proxy C.Proxy
 	var rule C.Rule
 	switch t.mode {
@@ -175,21 +181,24 @@ func (t *Tunnel) handleUDPConn(packet *inbound.PacketAdapter) {
 		return
 	}
 
-	if metadata.DstIP == nil {
-		ip, err := t.resolveIP(metadata.Host)
-		if err != nil {
-			log.Warnln("[UDP] Resolve %s failed: %s, %#v", metadata.Host, err.Error(), metadata)
-			return
-		}
-		metadata.DstIP = ip
+	if err := t.preHandleMetadata(metadata); err != nil {
+		log.Debugln("[Metadata PreHandle] error: %s", err)
+		return
 	}
 
 	key := packet.LocalAddr().String()
-
 	pc := t.natTable.Get(key)
-	addr := metadata.UDPAddr()
 	if pc != nil {
-		t.handleUDPToRemote(packet, pc, addr)
+		if !metadata.Resolved() {
+			ip, err := t.resolveIP(metadata.Host)
+			if err != nil {
+				log.Warnln("[UDP] Resolve %s failed: %s, %#v", metadata.Host, err.Error(), metadata)
+				return
+			}
+			metadata.DstIP = ip
+		}
+
+		t.handleUDPToRemote(packet, pc, metadata.UDPAddr())
 		return
 	}
 
@@ -236,7 +245,14 @@ func (t *Tunnel) handleUDPConn(packet *inbound.PacketAdapter) {
 		wg.Wait()
 		pc := t.natTable.Get(key)
 		if pc != nil {
-			t.handleUDPToRemote(packet, pc, addr)
+			if !metadata.Resolved() {
+				ip, err := dns.ResolveIP(metadata.Host)
+				if err != nil {
+					return
+				}
+				metadata.DstIP = ip
+			}
+			t.handleUDPToRemote(packet, pc, metadata.UDPAddr())
 		}
 	}()
 }
@@ -250,6 +266,11 @@ func (t *Tunnel) handleTCPConn(localConn C.ServerAdapter) {
 		return
 	}
 
+	if err := t.preHandleMetadata(metadata); err != nil {
+		log.Debugln("[Metadata PreHandle] error: %s", err)
+		return
+	}
+
 	proxy, rule, err := t.resolveMetadata(metadata)
 	if err != nil {
 		log.Warnln("Parse metadata failed: %v", err)