From 132a6a6a2f45cd1db3dc1f0b756944b78adf4a24 Mon Sep 17 00:00:00 2001
From: Kaming Chan <hi@terrychan.me>
Date: Wed, 23 Feb 2022 11:22:46 +0800
Subject: [PATCH] Fix: listener tcp keepalive & reuse net.BufferedConn (#1987)

---
 listener/http/proxy.go  | 7 ++++++-
 listener/mixed/mixed.go | 2 ++
 listener/socks/tcp.go   | 7 +------
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/listener/http/proxy.go b/listener/http/proxy.go
index 23a73739a..becb36730 100644
--- a/listener/http/proxy.go
+++ b/listener/http/proxy.go
@@ -19,7 +19,12 @@ func HandleConn(c net.Conn, in chan<- C.ConnContext, cache *cache.Cache) {
 	client := newClient(c.RemoteAddr(), in)
 	defer client.CloseIdleConnections()
 
-	conn := N.NewBufferedConn(c)
+	var conn *N.BufferedConn
+	if bufConn, ok := c.(*N.BufferedConn); ok {
+		conn = bufConn
+	} else {
+		conn = N.NewBufferedConn(c)
+	}
 
 	keepAlive := true
 	trusted := cache == nil // disable authenticate if cache is nil
diff --git a/listener/mixed/mixed.go b/listener/mixed/mixed.go
index 8fd4f990e..57fd055e0 100644
--- a/listener/mixed/mixed.go
+++ b/listener/mixed/mixed.go
@@ -64,6 +64,8 @@ func New(addr string, in chan<- C.ConnContext) (*Listener, error) {
 }
 
 func handleConn(conn net.Conn, in chan<- C.ConnContext, cache *cache.Cache) {
+	conn.(*net.TCPConn).SetKeepAlive(true)
+
 	bufConn := N.NewBufferedConn(conn)
 	head, err := bufConn.Peek(1)
 	if err != nil {
diff --git a/listener/socks/tcp.go b/listener/socks/tcp.go
index 29016f5b1..7cce32eec 100644
--- a/listener/socks/tcp.go
+++ b/listener/socks/tcp.go
@@ -61,6 +61,7 @@ func New(addr string, in chan<- C.ConnContext) (*Listener, error) {
 }
 
 func handleSocks(conn net.Conn, in chan<- C.ConnContext) {
+	conn.(*net.TCPConn).SetKeepAlive(true)
 	bufConn := N.NewBufferedConn(conn)
 	head, err := bufConn.Peek(1)
 	if err != nil {
@@ -84,9 +85,6 @@ func HandleSocks4(conn net.Conn, in chan<- C.ConnContext) {
 		conn.Close()
 		return
 	}
-	if c, ok := conn.(*net.TCPConn); ok {
-		c.SetKeepAlive(true)
-	}
 	in <- inbound.NewSocket(socks5.ParseAddr(addr), conn, C.SOCKS4)
 }
 
@@ -96,9 +94,6 @@ func HandleSocks5(conn net.Conn, in chan<- C.ConnContext) {
 		conn.Close()
 		return
 	}
-	if c, ok := conn.(*net.TCPConn); ok {
-		c.SetKeepAlive(true)
-	}
 	if command == socks5.CmdUDPAssociate {
 		defer conn.Close()
 		io.Copy(io.Discard, conn)