From f5dbf4289a124ae59fadb8d3ad192fda25082b3a Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Mon, 31 Aug 2020 03:00:31 +0800
Subject: [PATCH] opt sth

---
 proxy/proxy.go | 12 +++++-------
 proxy/udp.go   | 12 ++++++------
 2 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index f84e5db..438fadd 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -18,11 +18,15 @@ var (
 	hijackDNS []string
 )
 
-// Register Monitor
 func RegisterMonitor(m S.Monitor) {
 	monitor = m
 }
 
+func RegisterFakeDNS(d D.FakeDNS, h string) {
+	fakeDNS = d
+	hijackDNS = append(hijackDNS, strings.Split(h, ",")...)
+}
+
 // Session Operation
 func addSession(key interface{}, session *S.Session) {
 	if monitor != nil {
@@ -36,12 +40,6 @@ func removeSession(key interface{}) {
 	}
 }
 
-// Register FakeDNS
-func RegisterFakeDNS(d D.FakeDNS, h string) {
-	fakeDNS = d
-	hijackDNS = append(hijackDNS, strings.Split(h, ",")...)
-}
-
 // Check target if is hijacked address.
 func isHijacked(target *net.UDPAddr) bool {
 	if fakeDNS == nil {
diff --git a/proxy/udp.go b/proxy/udp.go
index 2073ea7..705c538 100644
--- a/proxy/udp.go
+++ b/proxy/udp.go
@@ -120,13 +120,13 @@ func (h *udpHandler) ReceiveTo(conn core.UDPConn, data []byte, addr *net.UDPAddr
 		resp, err := fakeDNS.Resolve(data)
 		if err != nil {
 			return fmt.Errorf("hijack DNS request error: %v", err)
-		} else {
-			if _, err = conn.WriteFrom(resp, addr); err != nil {
-				return fmt.Errorf("write dns answer failed: %v", err)
-			}
-			h.Close(conn)
-			return nil
 		}
+
+		if _, err = conn.WriteFrom(resp, addr); err != nil {
+			return fmt.Errorf("write dns answer failed: %v", err)
+		}
+		h.Close(conn)
+		return nil
 	}
 
 	var remoteAddr net.Addr