From 23f154659190e86715392cc89854527bca9e04e2 Mon Sep 17 00:00:00 2001
From: yinyue200 <yinyue200@outlook.com>
Date: Sun, 12 Feb 2023 21:19:28 +0800
Subject: [PATCH] Remove unnecessary DNS queries in case of SOCKS5 outbound

---
 core/cs/server_client.go | 34 ++++++++++++++++++++++++++++------
 core/transport/server.go | 12 ++++++++++--
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/core/cs/server_client.go b/core/cs/server_client.go
index aaa02a5..db109b1 100644
--- a/core/cs/server_client.go
+++ b/core/cs/server_client.go
@@ -140,10 +140,13 @@ func (c *serverClient) handleMessage(msg []byte) {
 		var err error
 		if c.ACLEngine != nil {
 			action, arg, isDomain, ipAddr, err = c.ACLEngine.ResolveAndMatch(dfMsg.Host, dfMsg.Port, true)
+		} else if c.Transport.ProxyEnabled() { // Case for SOCKS5 outbound
+			ipAddr, isDomain = c.Transport.ParseIPAddr(dfMsg.Host) // It is safe to leave ipAddr as nil since addrExToSOCKS5Addr will ignore it when there is a domain
+			err = nil
 		} else {
 			ipAddr, isDomain, err = c.Transport.ResolveIPAddr(dfMsg.Host)
 		}
-		if err != nil && !(isDomain && c.Transport.ProxyEnabled()) { // Special case for domain requests + SOCKS5 outbound
+		if err != nil {
 			return
 		}
 		switch action {
@@ -162,8 +165,16 @@ func (c *serverClient) handleMessage(msg []byte) {
 		case acl.ActionBlock:
 			// Do nothing
 		case acl.ActionHijack:
-			hijackIPAddr, isDomain, err := c.Transport.ResolveIPAddr(arg)
-			if err == nil || (isDomain && c.Transport.ProxyEnabled()) { // Special case for domain requests + SOCKS5 outbound
+			var isDomain bool
+			var hijackIPAddr *net.IPAddr
+			var err error
+			if c.Transport.ProxyEnabled() { // Case for domain requests + SOCKS5 outbound
+				hijackIPAddr, isDomain = c.Transport.ParseIPAddr(arg) // It is safe to leave ipAddr as nil since addrExToSOCKS5Addr will ignore it when there is a domain
+				err = nil
+			} else {
+				hijackIPAddr, isDomain, err = c.Transport.ResolveIPAddr(arg)
+			}
+			if err == nil {
 				addrEx := &transport.AddrEx{
 					IPAddr: hijackIPAddr,
 					Port:   int(dfMsg.Port),
@@ -190,10 +201,13 @@ func (c *serverClient) handleTCP(stream quic.Stream, host string, port uint16) {
 	var err error
 	if c.ACLEngine != nil {
 		action, arg, isDomain, ipAddr, err = c.ACLEngine.ResolveAndMatch(host, port, false)
+	} else if c.Transport.ProxyEnabled() { // Case for domain requests + SOCKS5 outbound
+		ipAddr, isDomain = c.Transport.ParseIPAddr(host) // It is safe to leave ipAddr as nil since addrExToSOCKS5Addr will ignore it when there is a domain
+		err = nil
 	} else {
 		ipAddr, isDomain, err = c.Transport.ResolveIPAddr(host)
 	}
-	if err != nil && !(isDomain && c.Transport.ProxyEnabled()) { // Special case for domain requests + SOCKS5 outbound
+	if err != nil {
 		_ = struc.Pack(stream, &serverResponse{
 			OK:      false,
 			Message: "host resolution failure",
@@ -229,8 +243,16 @@ func (c *serverClient) handleTCP(stream quic.Stream, host string, port uint16) {
 		})
 		return
 	case acl.ActionHijack:
-		hijackIPAddr, isDomain, err := c.Transport.ResolveIPAddr(arg)
-		if err != nil && !(isDomain && c.Transport.ProxyEnabled()) { // Special case for domain requests + SOCKS5 outbound
+		var isDomain bool
+		var hijackIPAddr *net.IPAddr
+		var err error
+		if c.Transport.ProxyEnabled() { // Case for domain requests + SOCKS5 outbound
+			hijackIPAddr, isDomain = c.Transport.ParseIPAddr(arg) // It is safe to leave ipAddr as nil since addrExToSOCKS5Addr will ignore it when there is a domain
+			err = nil
+		} else {
+			hijackIPAddr, isDomain, err = c.Transport.ResolveIPAddr(arg)
+		}
+		if err != nil {
 			_ = struc.Pack(stream, &serverResponse{
 				OK:      false,
 				Message: err.Error(),
diff --git a/core/transport/server.go b/core/transport/server.go
index 42a5025..3b9f79c 100644
--- a/core/transport/server.go
+++ b/core/transport/server.go
@@ -69,10 +69,18 @@ var DefaultServerTransport = &ServerTransport{
 	ResolvePreference: ResolvePreferenceDefault,
 }
 
-func (st *ServerTransport) ResolveIPAddr(address string) (*net.IPAddr, bool, error) {
+func (st *ServerTransport) ParseIPAddr(address string) (*net.IPAddr, bool) {
 	ip, zone := utils.ParseIPZone(address)
 	if ip != nil {
-		return &net.IPAddr{IP: ip, Zone: zone}, false, nil
+		return &net.IPAddr{IP: ip, Zone: zone}, false
+	}
+	return nil, true
+}
+
+func (st *ServerTransport) ResolveIPAddr(address string) (*net.IPAddr, bool, error) {
+	ip, isDomain := st.ParseIPAddr(address)
+	if !isDomain {
+		return ip, false, nil
 	}
 	ipAddr, err := resolveIPAddrWithPreference(address, st.ResolvePreference)
 	return ipAddr, true, err