fix(server): close underlying work conn when server bandwidth limiter is enabled (#5553)

When bandwidthLimitMode is "server", the work connection is wrapped by
libio.WrapReadWriteCloser with a close callback that captured the
variable being reassigned (local/rwc). At call time the callback
referenced the wrapper itself, whose closed guard was already set, so
Close returned nil without closing the underlying work connection.

For TCP proxies this left libio.Join blocked on the other copy
goroutine after the user disconnected, keeping the work connection and
the backend connection open. HTTP and UDP proxies had the same pattern.

Bind the Close method value before reassignment, and add an e2e test
verifying that the local connection is closed after the user
disconnects.

server/proxy/http.go

@@ -139,9 +139,7 @@ func (pxy *HTTPProxy) GetRealConn(remoteAddr string) (workConn net.Conn, err err
 	}
 
 	if pxy.GetLimiter() != nil {
-		rwc = libio.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), func() error {
-			return rwc.Close()
-		})
+		rwc = libio.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), rwc.Close)
 	}
 
 	workConn = netpkg.WrapReadWriteCloserToConn(rwc, tmpConn)

server/proxy/proxy.go

@@ -308,9 +308,7 @@ func (pxy *BaseProxy) handleUserTCPConnection(userConn net.Conn) {
 	}
 
 	if pxy.GetLimiter() != nil {
-		local = libio.WrapReadWriteCloser(limit.NewReader(local, pxy.GetLimiter()), limit.NewWriter(local, pxy.GetLimiter()), func() error {
-			return local.Close()
-		})
+		local = libio.WrapReadWriteCloser(limit.NewReader(local, pxy.GetLimiter()), limit.NewWriter(local, pxy.GetLimiter()), local.Close)
 	}
 
 	xl.Debugf("join connections, workConn(l[%s] r[%s]) userConn(l[%s] r[%s])", workConn.LocalAddr().String(),

server/proxy/udp.go

@@ -217,9 +217,7 @@ func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
 			}
 
 			if pxy.GetLimiter() != nil {
-				rwc = libio.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), func() error {
-					return rwc.Close()
-				})
+				rwc = libio.WrapReadWriteCloser(limit.NewReader(rwc, pxy.GetLimiter()), limit.NewWriter(rwc, pxy.GetLimiter()), rwc.Close)
 			}
 
 			pxy.workConn = netpkg.WrapReadWriteCloserToConn(rwc, workConn)

test/e2e/v1/features/bandwidth_limit.go

@@ -2,6 +2,9 @@ package features
 
 import (
 	"fmt"
+	"io"
+	"net"
+	"strconv"
 	"strings"
 	"time"
 
@@ -105,4 +108,82 @@ var _ = ginkgo.Describe("[Feature: Bandwidth Limit]", func() {
 
 		framework.ExpectTrue(duration.Seconds() > 8, "100Kb with 10KB limit, want > 8 seconds, but got %s", duration.String())
 	})
+
+	ginkgo.It("Server Bandwidth Limit closes work connection when user disconnects", func() {
+		newFunc := func() *plugin.Request {
+			var r plugin.Request
+			r.Content = &plugin.NewProxyContent{}
+			return &r
+		}
+		pluginPort := f.AllocPort()
+		handler := func(req *plugin.Request) *plugin.Response {
+			var ret plugin.Response
+			content := req.Content.(*plugin.NewProxyContent)
+			content.BandwidthLimit = "10MB"
+			content.BandwidthLimitMode = "server"
+			ret.Content = content
+			return &ret
+		}
+		pluginServer := pluginpkg.NewHTTPPluginServer(pluginPort, newFunc, handler, nil)
+
+		f.RunServer("", pluginServer)
+
+		serverConf := consts.DefaultServerConfig + fmt.Sprintf(`
+		[[httpPlugins]]
+		name = "test"
+		addr = "127.0.0.1:%d"
+		path = "/handler"
+		ops = ["NewProxy"]
+		`, pluginPort)
+		clientConf := consts.DefaultClientConfig
+
+		// The local server echoes the first message, then stays idle until the
+		// connection from frpc is closed.
+		localClosed := make(chan struct{})
+		localPort := f.AllocPort()
+		localServer := streamserver.New(streamserver.TCP, streamserver.WithBindPort(localPort),
+			streamserver.WithCustomHandler(func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 1024)
+				n, err := c.Read(buf)
+				if err != nil {
+					return
+				}
+				if _, err := c.Write(buf[:n]); err != nil {
+					return
+				}
+				_, _ = io.Copy(io.Discard, c)
+				close(localClosed)
+			}))
+		f.RunServer("", localServer)
+
+		remotePort := f.AllocPort()
+		clientConf += fmt.Sprintf(`
+			[[proxies]]
+			name = "tcp"
+			type = "tcp"
+			localPort = %d
+			remotePort = %d
+			`, localPort, remotePort)
+
+		f.RunProcesses(serverConf, []string{clientConf})
+
+		conn, err := net.Dial("tcp", net.JoinHostPort("127.0.0.1", strconv.Itoa(remotePort)))
+		framework.ExpectNoError(err)
+		_ = conn.SetDeadline(time.Now().Add(5 * time.Second))
+		_, err = conn.Write([]byte("ping"))
+		framework.ExpectNoError(err)
+		buf := make([]byte, 4)
+		_, err = io.ReadFull(conn, buf)
+		framework.ExpectNoError(err)
+		framework.ExpectEqual(string(buf), "ping")
+
+		conn.Close()
+
+		select {
+		case <-localClosed:
+		case <-time.After(5 * time.Second):
+			framework.Failf("local connection is not closed after user connection is closed")
+		}
+	})
 })