From bcc6d7c50481c83471be6358703f4fcb1d36759a Mon Sep 17 00:00:00 2001
From: wweir <wweir@foxmail.com>
Date: Wed, 19 Feb 2020 05:58:05 +0800
Subject: [PATCH] Fix the last dot in domain match

---
 conf/dynamic_rule.go        | 14 ++++++++------
 internal/http/tgt_parser.go |  7 ++++++-
 proxy/proxy.go              |  3 ++-
 3 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/conf/dynamic_rule.go b/conf/dynamic_rule.go
index 7c1759c..e799c33 100644
--- a/conf/dynamic_rule.go
+++ b/conf/dynamic_rule.go
@@ -27,8 +27,14 @@ var dynamicMu = sync.Mutex{}
 
 // ShouldProxy check if the domain shoule request though proxy
 func ShouldProxy(domain string) bool {
+	// break deadlook, for wildcard
+	if strings.Count(domain, ".") > 4 {
+		return false
+	}
+	domain = strings.TrimSuffix(domain, ".")
+
 	if domain == Client.Address {
-		return true
+		return false
 	}
 	if Client.Router.directRules.Match(domain) {
 		return false
@@ -43,11 +49,7 @@ func ShouldProxy(domain string) bool {
 }
 
 func (d *dynamic) Get(key interface{}) (err error) {
-	// break deadloop, for ugly wildcard setting dns setting
-	domain := strings.TrimSuffix(key.(string), ".")
-	if strings.Count(domain, ".") > 10 {
-		return nil
-	}
+	domain := key.(string)
 	domainUnderscore := strings.ReplaceAll(domain, ".", "_")
 	var score int
 
diff --git a/internal/http/tgt_parser.go b/internal/http/tgt_parser.go
index 9dafbce..4e0eb86 100644
--- a/internal/http/tgt_parser.go
+++ b/internal/http/tgt_parser.go
@@ -64,7 +64,12 @@ func (c *conn) Write(b []byte) (n int, err error) {
 			return 0, err
 		}
 
-		n, err := c.Conn.Write(append(c.domain, b...))
+		switch c.typ {
+		case TGT_OTHER:
+			n, err = c.Conn.Write(append(c.domain, b...))
+		default:
+			n, err = c.Conn.Write(b)
+		}
 		return n - len(c.domain), err
 	}
 
diff --git a/proxy/proxy.go b/proxy/proxy.go
index 091ec47..5ee680f 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -133,12 +133,13 @@ func StartServer(relayTarget, password, certFile, keyFile, email string) {
 		}
 
 		go func(conn net.Conn) {
+			defer conn.Close()
+
 			conn, domain, port, err := _http.ParseAddr(conn, passwordData)
 			if err != nil {
 				log.Errorw("parse relay target", "err", err)
 				return
 			}
-			defer conn.Close()
 
 			addr := relayTarget
 			if domain != "" {