Fix: listener tcp keepalive & reuse net.BufferedConn (#1987)

lchenc3 · Feb 23, 2022 · 132a6a6 · 132a6a6
1 parent 03e4b5d
commit 132a6a6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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)