From 71b31978b3685c355165f9bf1a3584036b8097e5 Mon Sep 17 00:00:00 2001 From: Mxmilu Date: Sat, 10 Oct 2026 10:45:34 +0800 Subject: [PATCH] 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 | 4 +- server/proxy/proxy.go | 4 +- server/proxy/udp.go | 4 +- test/e2e/v1/features/bandwidth_limit.go | 81 +++++++++++++++++++++++++ 4 files changed, 84 insertions(+), 9 deletions(-) diff --git a/server/proxy/http.go b/server/proxy/http.go index 05afc2e9..c3da8383 100644 --- a/server/proxy/http.go +++ b/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) diff --git a/server/proxy/proxy.go b/server/proxy/proxy.go index f4e75f0a..705aa9da 100644 --- a/server/proxy/proxy.go +++ b/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(), diff --git a/server/proxy/udp.go b/server/proxy/udp.go index 47bdfc8c..8a69c360 100644 --- a/server/proxy/udp.go +++ b/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) diff --git a/test/e2e/v1/features/bandwidth_limit.go b/test/e2e/v1/features/bandwidth_limit.go index 11503adb..a975bfe5 100644 --- a/test/e2e/v1/features/bandwidth_limit.go +++ b/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") + } + }) })