From 5fe3ef94dd6ba973d8b937a82d0127a36a6ac626 Mon Sep 17 00:00:00 2001
From: fsj2009yx <289137665@qq.com>
Date: Sun, 31 Aug 2025 11:20:41 +0800
Subject: [PATCH] use buffer pool on two-way relay

---
 tcp.go   | 24 ++++++++++++++++++++----
 utils.go | 11 +++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/tcp.go b/tcp.go
index 2a297ca..cdd196d 100644
--- a/tcp.go
+++ b/tcp.go
@@ -13,8 +13,15 @@ import (
 )
 
 func tcpCopyData(dst net.Conn, src net.Conn, ch chan<- error) {
-	_, err := io.Copy(dst, src)
-	ch <- err
+	buf := GetBuffer()
+	defer PutBuffer(buf)
+
+	_, err := io.CopyBuffer(dst, src, buf)
+	if err != nil && err != io.EOF {
+		ch <- err
+	} else {
+		ch <- nil
+	}
 }
 
 func tcpHandleConnection(conn net.Conn, logger *slog.Logger) {
@@ -44,6 +51,7 @@ func tcpHandleConnection(conn net.Conn, logger *slog.Logger) {
 		return
 	}
 
+	//parse PROXY header
 	saddr, _, restBytes, err := PROXYReadRemoteAddr(buffer[:n], TCP)
 	if err != nil {
 		logger.Debug("failed to parse PROXY header", "error", err, slog.Bool("dropConnection", true))
@@ -83,6 +91,7 @@ func tcpHandleConnection(conn net.Conn, logger *slog.Logger) {
 		logger.Debug("successfully established upstream connection")
 	}
 
+	//Disable Nagle's Algorithm
 	if err := conn.(*net.TCPConn).SetNoDelay(true); err != nil {
 		logger.Debug("failed to set nodelay on downstream connection", "error", err, slog.Bool("dropConnection", true))
 	} else if Opts.Verbose > 1 {
@@ -109,8 +118,15 @@ func tcpHandleConnection(conn net.Conn, logger *slog.Logger) {
 	buffer = nil
 
 	outErr := make(chan error, 2)
-	go tcpCopyData(upstreamConn, conn, outErr)
-	go tcpCopyData(conn, upstreamConn, outErr)
+	go func() {
+		tcpCopyData(conn, upstreamConn, outErr)
+		conn.Close()
+	}()
+
+	go func() {
+		tcpCopyData(upstreamConn, conn, outErr)
+		upstreamConn.Close()
+	}()
 
 	err = <-outErr
 	if err != nil {
diff --git a/utils.go b/utils.go
index 432c152..179c90f 100644
--- a/utils.go
+++ b/utils.go
@@ -37,23 +37,33 @@ func DialUpstreamControl(sport int) func(string, string, syscall.RawConn) error
 			if Opts.Protocol == "tcp" {
 				syscallErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_SYNCNT, 2)
 				if syscallErr != nil {
+					//TCP_SYNCNT:SYN retransmission count set to 2, speed up the detection of connection failures
 					syscallErr = fmt.Errorf("setsockopt(IPPROTO_TCP, TCP_SYNCTNT, 2): %w", syscallErr)
 					return
 				}
 			}
 
+			// IP_TRANSPARENT allows a socket to bind to an IP address that is not
+			// configured on the local machine. This is essential for transparent
+			// This option requires appropriate capabilities (CAP_NET_ADMIN and CAP_NET_RAW)
 			syscallErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IP, syscall.IP_TRANSPARENT, 1)
 			if syscallErr != nil {
 				syscallErr = fmt.Errorf("setsockopt(IPPROTO_IP, IP_TRANSPARENT, 1): %w", syscallErr)
 				return
 			}
 
+			// SO_REUSEADDR：allows multiple sockets to bind to the same address and port
+			// This is useful for restarting the proxy without waiting for old sockets to time out
 			syscallErr = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1)
 			if syscallErr != nil {
 				syscallErr = fmt.Errorf("setsockopt(SOL_SOCKET, SO_REUSEADDR, 1): %w", syscallErr)
 				return
 			}
 
+			//IP_BIND_ADDRESS_NO_PORT（Linux 4.4+）
+			// Allows binding to an IP address without specifying a port
+			// This is useful for transparent proxying scenarios where the original
+			// destination port is not known at bind time
 			if sport == 0 {
 				ipBindAddressNoPort := 24
 				syscallErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_IP, ipBindAddressNoPort, 1)
@@ -63,6 +73,7 @@ func DialUpstreamControl(sport int) func(string, string, syscall.RawConn) error
 				}
 			}
 
+			//SO_MARK used by iptables to mark packets for routing
 			if Opts.Mark != 0 {
 				syscallErr = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_MARK, Opts.Mark)
 				if syscallErr != nil {