From 0f7f0a9b1a0bb48d9fb268a03701686576c00025 Mon Sep 17 00:00:00 2001
From: "X. Jason Lyu" <xjasonlyu@gmail.com>
Date: Mon, 12 Aug 2019 14:01:32 +0800
Subject: [PATCH] Optimization: socks UDP & fix typo (#261)

---
 .gitignore                  |  6 ++++++
 adapters/outbound/socks5.go | 10 ++++++++--
 tunnel/tunnel.go            | 18 +++++++++---------
 3 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 9a9bbe19..0593cfd0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,3 +14,9 @@ bin/*
 
 # dep
 vendor
+
+# GoLand
+.idea/*
+
+# macOS file
+.DS_Store
diff --git a/adapters/outbound/socks5.go b/adapters/outbound/socks5.go
index d3f7a3a9..a72c9386 100644
--- a/adapters/outbound/socks5.go
+++ b/adapters/outbound/socks5.go
@@ -116,7 +116,7 @@ func (ss *Socks5) DialUDP(metadata *C.Metadata) (_ C.PacketConn, _ net.Addr, err
 		pc.Close()
 	}()
 
-	return newPacketConn(&socksUDPConn{PacketConn: pc, rAddr: targetAddr}, ss), addr, nil
+	return newPacketConn(&socksUDPConn{PacketConn: pc, rAddr: targetAddr, tcpConn: c}, ss), addr, nil
 }
 
 func NewSocks5(option Socks5Option) *Socks5 {
@@ -146,7 +146,8 @@ func NewSocks5(option Socks5Option) *Socks5 {
 
 type socksUDPConn struct {
 	net.PacketConn
-	rAddr net.Addr
+	rAddr   net.Addr
+	tcpConn net.Conn
 }
 
 func (uc *socksUDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
@@ -166,3 +167,8 @@ func (uc *socksUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
 	copy(b, payload)
 	return n - len(addr) - 3, a, e
 }
+
+func (uc *socksUDPConn) Close() error {
+	uc.tcpConn.Close()
+	return uc.PacketConn.Close()
+}
diff --git a/tunnel/tunnel.go b/tunnel/tunnel.go
index 8681e5c2..837374ad 100644
--- a/tunnel/tunnel.go
+++ b/tunnel/tunnel.go
@@ -161,16 +161,16 @@ func (t *Tunnel) handleConn(localConn C.ServerAdapter) {
 func (t *Tunnel) handleUDPConn(localConn C.ServerAdapter, metadata *C.Metadata, proxy C.Proxy, rule C.Rule) {
 	pc, addr := natTable.Get(localConn.RemoteAddr())
 	if pc == nil {
-		rawpc, naddr, err := proxy.DialUDP(metadata)
-		addr = naddr
-		pc = rawpc
+		rawPc, nAddr, err := proxy.DialUDP(metadata)
+		addr = nAddr
+		pc = rawPc
 		if err != nil {
 			log.Warnln("dial %s error: %s", proxy.Name(), err.Error())
 			return
 		}
 
 		if rule != nil {
-			log.Infoln("%s --> %v match %s using %s", metadata.SrcIP.String(), metadata.String(), rule.RuleType().String(), rawpc.Chains().String())
+			log.Infoln("%s --> %v match %s using %s", metadata.SrcIP.String(), metadata.String(), rule.RuleType().String(), rawPc.Chains().String())
 		} else {
 			log.Infoln("%s --> %v doesn't match any rule using DIRECT", metadata.SrcIP.String(), metadata.String())
 		}
@@ -183,24 +183,24 @@ func (t *Tunnel) handleUDPConn(localConn C.ServerAdapter, metadata *C.Metadata,
 }
 
 func (t *Tunnel) handleTCPConn(localConn C.ServerAdapter, metadata *C.Metadata, proxy C.Proxy, rule C.Rule) {
-	remoConn, err := proxy.Dial(metadata)
+	remoteConn, err := proxy.Dial(metadata)
 	if err != nil {
 		log.Warnln("dial %s error: %s", proxy.Name(), err.Error())
 		return
 	}
-	defer remoConn.Close()
+	defer remoteConn.Close()
 
 	if rule != nil {
-		log.Infoln("%s --> %v match %s using %s", metadata.SrcIP.String(), metadata.String(), rule.RuleType().String(), remoConn.Chains().String())
+		log.Infoln("%s --> %v match %s using %s", metadata.SrcIP.String(), metadata.String(), rule.RuleType().String(), remoteConn.Chains().String())
 	} else {
 		log.Infoln("%s --> %v doesn't match any rule using DIRECT", metadata.SrcIP.String(), metadata.String())
 	}
 
 	switch adapter := localConn.(type) {
 	case *InboundAdapter.HTTPAdapter:
-		t.handleHTTP(adapter, remoConn)
+		t.handleHTTP(adapter, remoteConn)
 	case *InboundAdapter.SocketAdapter:
-		t.handleSocket(adapter, remoConn)
+		t.handleSocket(adapter, remoteConn)
 	}
 }