diff --git a/pkg/util/vhost/vhost.go b/pkg/util/vhost/vhost.go index 637c8dcf..f29333f2 100644 --- a/pkg/util/vhost/vhost.go +++ b/pkg/util/vhost/vhost.go @@ -17,10 +17,9 @@ import ( "fmt" "net" "strings" + "sync" "time" - "github.com/fatedier/golib/errors" - "github.com/fatedier/frp/pkg/util/log" netpkg "github.com/fatedier/frp/pkg/util/net" "github.com/fatedier/frp/pkg/util/xlog" @@ -138,6 +137,7 @@ func (v *Muxer) Listen(ctx context.Context, cfg *RouteConfig) (l *Listener, err password: cfg.Password, mux: v, accept: make(chan net.Conn), + done: make(chan struct{}), ctx: ctx, } err = v.registryRouter.Add(cfg.Domain, cfg.Location, cfg.RouteByHTTPUser, l) @@ -215,11 +215,12 @@ func (v *Muxer) handle(c net.Conn) { c = sConn xl.Debugf("new request host [%s] path [%s] httpUser [%s]", name, path, httpUser) - err = errors.PanicToError(func() { - l.accept <- c - }) - if err != nil { + select { + case l.accept <- c: + // The connection is now owned by the Accept caller. + case <-l.done: xl.Warnf("listener is already closed, ignore this request") + _ = c.Close() } } @@ -232,13 +233,17 @@ type Listener struct { password string mux *Muxer // for closing Muxer accept chan net.Conn + done chan struct{} + closeOnce sync.Once ctx context.Context } func (l *Listener) Accept() (net.Conn, error) { xl := xlog.FromContextSafe(l.ctx) - conn, ok := <-l.accept - if !ok { + var conn net.Conn + select { + case conn = <-l.accept: + case <-l.done: return nil, fmt.Errorf("listener closed") } @@ -258,8 +263,10 @@ func (l *Listener) Accept() (net.Conn, error) { } func (l *Listener) Close() error { - l.mux.registryRouter.Del(l.name, l.location, l.routeByHTTPUser) - close(l.accept) + l.closeOnce.Do(func() { + l.mux.registryRouter.Del(l.name, l.location, l.routeByHTTPUser) + close(l.done) + }) return nil } diff --git a/pkg/util/vhost/vhost_close_test.go b/pkg/util/vhost/vhost_close_test.go new file mode 100644 index 00000000..1e32c6d9 --- /dev/null +++ b/pkg/util/vhost/vhost_close_test.go @@ -0,0 +1,116 @@ +// Copyright 2026 The frp Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package vhost + +import ( + "context" + "net" + "testing" + "time" +) + +func newCloseTestMuxer() *Muxer { + return &Muxer{ + timeout: time.Second, + registryRouter: NewRouters(), + vhostFunc: func(c net.Conn) (net.Conn, map[string]string, error) { + return c, map[string]string{"Host": "example.com"}, nil + }, + } +} + +func testPendingHandoffIsClosed(t *testing.T, client, server net.Conn) { + t.Helper() + + m := newCloseTestMuxer() + entered := make(chan struct{}) + release := make(chan struct{}) + m.SetSuccessHookFunc(func(net.Conn, map[string]string) error { + close(entered) + <-release + return nil + }) + l, err := m.Listen(context.Background(), &RouteConfig{Domain: "example.com"}) + if err != nil { + t.Fatal(err) + } + + handled := make(chan struct{}) + go func() { + defer close(handled) + m.handle(server) + }() + <-entered + if err := client.SetReadDeadline(time.Now().Add(300 * time.Millisecond)); err != nil { + t.Fatal(err) + } + + // Close the route after resolution but before the handoff. The pending + // connection must be closed, not abandoned. + if err := l.Close(); err != nil { + t.Fatal(err) + } + close(release) + <-handled + + if _, err := client.Read(make([]byte, 1)); err == nil { + t.Fatal("connection was not closed by the muxer") + } else if ne, ok := err.(net.Error); ok && ne.Timeout() { + t.Fatalf("connection was not closed by the muxer: %v", err) + } +} + +func TestListenerCloseReleasesPendingConn(t *testing.T) { + client, server := net.Pipe() + defer client.Close() + testPendingHandoffIsClosed(t, client, server) +} + +func TestListenerCloseReleasesPendingTCPConn(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + + client, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + defer client.Close() + + server, err := ln.Accept() + if err != nil { + t.Fatal(err) + } + testPendingHandoffIsClosed(t, client, server) +} + +func TestListenerCloseIsIdempotent(t *testing.T) { + m := newCloseTestMuxer() + l, err := m.Listen(context.Background(), &RouteConfig{Domain: "example.com"}) + if err != nil { + t.Fatal(err) + } + if err := l.Close(); err != nil { + t.Fatal(err) + } + if err := l.Close(); err != nil { + t.Fatal(err) + } + if _, err := l.Accept(); err == nil { + t.Fatal("Accept after Close should return an error") + } +}