diff --git a/pkg/msg/handler.go b/pkg/msg/handler.go index 9e437198..b0ac187a 100644 --- a/pkg/msg/handler.go +++ b/pkg/msg/handler.go @@ -20,6 +20,7 @@ import ( "net" "reflect" "sync" + "time" "github.com/fatedier/frp/pkg/proto/wire" ) @@ -117,9 +118,11 @@ type Dispatcher struct { msgHandlers map[reflect.Type]func(Message) } -// NewDispatcher uses Close, when rw implements io.Closer, to interrupt ReadMsg -// after a write failure. Otherwise, the caller must arrange for ReadMsg to return -// before Done can close. +// NewDispatcher uses SetReadDeadline, when supported, and Close, when rw +// implements io.Closer, to interrupt ReadMsg after a write failure. Without a +// working read deadline, Close or the caller must arrange for ReadMsg to return. +// These methods must return in bounded time and must not wait for Done: Done +// closes only after ReadMsg, any synchronous handler, and stop have returned. func NewDispatcher(rw ReadWriter) *Dispatcher { return &Dispatcher{ rw: rw, @@ -136,11 +139,14 @@ func (d *Dispatcher) Run() { go d.readLoop() } -// stop rejects new sends and interrupts ReadMsg when rw is closable. Done remains +// stop rejects new sends and interrupts ReadMsg when supported by rw. Done remains // owned by readLoop so the control owner cannot clean up while a handler is running. func (d *Dispatcher) stop() { d.stopOnce.Do(func() { close(d.stopCh) + if conn, ok := d.rw.(interface{ SetReadDeadline(time.Time) error }); ok { + _ = conn.SetReadDeadline(time.Now()) + } if closer, ok := d.rw.(io.Closer); ok { _ = closer.Close() } diff --git a/pkg/msg/handler_shutdown_test.go b/pkg/msg/handler_shutdown_test.go index ee6f193a..3f74c06d 100644 --- a/pkg/msg/handler_shutdown_test.go +++ b/pkg/msg/handler_shutdown_test.go @@ -31,18 +31,13 @@ func TestDispatcherEnqueueRejectsStoppedDispatcher(t *testing.T) { // Model shutdown after Send's initial check: the stop and buffered send // cases are both ready. A send selected here must still return EOF. - enqueued := false - for range 100 { - if err := d.enqueue(&ReqWorkConn{}); !errors.Is(err, io.EOF) { - t.Fatalf("enqueue racing with shutdown: %v, want EOF", err) - } - if len(d.sendCh) > 0 { - enqueued = true - <-d.sendCh - } + if err := d.enqueue(&ReqWorkConn{}); !errors.Is(err, io.EOF) { + t.Fatalf("enqueue racing with shutdown: %v, want EOF", err) } - if !enqueued { - t.Fatal("did not exercise the ready send case") + // This single call checks that a stopped dispatcher's enqueue returns EOF; it + // does not guarantee exercising the send case and post-enqueue check on every run. + if len(d.sendCh) > 0 { + t.Log("selected the send case; post-enqueue check returned EOF") } } @@ -65,7 +60,7 @@ func TestDispatcherAcceptsReadWriterWithoutClose(t *testing.T) { }) c.transport.arm() - if err := d.Send(&ReqWorkConn{}); err != nil { + if err := d.Send(&ReqWorkConn{}); err != nil && !errors.Is(err, io.EOF) { t.Fatal(err) } if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { diff --git a/pkg/msg/handler_write_error_test.go b/pkg/msg/handler_write_error_test.go index 0e59fc7b..17142297 100644 --- a/pkg/msg/handler_write_error_test.go +++ b/pkg/msg/handler_write_error_test.go @@ -207,7 +207,7 @@ func TestDispatcherWriteErrorClosesControlAndFreshSessionWorks(t *testing.T) { d := NewDispatcher(c) d.Run() c.transport.arm() - if err := d.Send(&ReqWorkConn{}); err != nil { + if err := d.Send(&ReqWorkConn{}); err != nil && !errors.Is(err, io.EOF) { t.Fatal(err) } if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { @@ -215,10 +215,8 @@ func TestDispatcherWriteErrorClosesControlAndFreshSessionWorks(t *testing.T) { } regressionAwait(t, c.transport.closed, "failed control connection close") regressionAwait(t, d.Done(), "failed dispatcher completion") - for range 100 { - if err := d.Send(&ReqWorkConn{}); !errors.Is(err, io.EOF) { - t.Fatalf("Send after completed shutdown: %v", err) - } + if err := d.Send(&ReqWorkConn{}); !errors.Is(err, io.EOF) { + t.Fatalf("Send after completed shutdown: %v", err) } if c.closeCalls.Load() != 1 { t.Fatalf("connection closed %d times", c.closeCalls.Load()) @@ -255,40 +253,42 @@ func TestDispatcherWriteErrorClosesControlAndFreshSessionWorks(t *testing.T) { } func TestDispatcherWriteErrorWaitsForSynchronousHandler(t *testing.T) { - c := newRegressionControl(t) - regressionWarm(t, c) - defer c.transport.close() - d := NewDispatcher(c) - entered, release := make(chan struct{}), make(chan struct{}) - var releaseOnce sync.Once - defer releaseOnce.Do(func() { close(release) }) - d.RegisterHandler(&ReqWorkConn{}, func(Message) { - close(entered) - <-release - }) - d.Run() - c.incoming <- &ReqWorkConn{} - regressionAwait(t, entered, "synchronous handler entry") - c.transport.arm() - if err := d.Send(&ReqWorkConn{}); err != nil { - t.Fatal(err) - } - if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { - t.Fatal(err) - } - regressionAwait(t, c.transport.closed, "write failure closes control during handler") - for range 100 { + synctest.Test(t, func(t *testing.T) { + c := newRegressionControl(t) + regressionWarm(t, c) + defer c.transport.close() + d := NewDispatcher(c) + entered, release := make(chan struct{}), make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + d.RegisterHandler(&ReqWorkConn{}, func(Message) { + close(entered) + <-release + }) + d.Run() + c.incoming <- &ReqWorkConn{} + regressionAwait(t, entered, "synchronous handler entry") + c.transport.arm() + if err := d.Send(&ReqWorkConn{}); err != nil && !errors.Is(err, io.EOF) { + t.Fatal(err) + } + if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { + t.Fatal(err) + } + regressionAwait(t, c.transport.closed, "write failure closes control during handler") if err := d.Send(&ReqWorkConn{}); !errors.Is(err, io.EOF) { t.Fatalf("Send after stop while handler is active: %v", err) } - } - select { - case <-d.Done(): - t.Fatal("Done closed before synchronous handler returned; owner cleanup could race proxy registration") - case <-time.After(100 * time.Millisecond): - } - releaseOnce.Do(func() { close(release) }) - regressionAwait(t, d.Done(), "completion after handler returned") + // Let the writer finish stopping while the handler remains blocked. + synctest.Wait() + select { + case <-d.Done(): + t.Fatal("Done closed before synchronous handler returned; owner cleanup could race proxy registration") + default: + } + releaseOnce.Do(func() { close(release) }) + regressionAwait(t, d.Done(), "completion after handler returned") + }) } func TestDispatcherWriteErrorUnblocksSenderWithFullQueue(t *testing.T) { @@ -335,53 +335,49 @@ func TestDispatcherWriteErrorUnblocksSenderWithFullQueue(t *testing.T) { } func TestDispatcherConcurrentReadWriteFailureClosesOnce(t *testing.T) { - for range 100 { - func() { - c := newRegressionControl(t) - regressionWarm(t, c) - defer c.transport.close() - readEntered, writeEntered := make(chan struct{}), make(chan struct{}) - release := make(chan struct{}) - var releaseOnce sync.Once - defer releaseOnce.Do(func() { close(release) }) - c.beforeRead = func() error { - close(readEntered) - <-release - return io.EOF - } - c.beforeWrite = func() { - close(writeEntered) - <-release - } - d := NewDispatcher(c) - c.transport.arm() - if err := d.Send(&ReqWorkConn{}); err != nil { - t.Fatal(err) - } - readExited, writeExited := make(chan struct{}), make(chan struct{}) - go func() { - defer close(readExited) - d.readLoop() - }() - go func() { - defer close(writeExited) - d.sendLoop() - }() - regressionAwait(t, readEntered, "read path before injected failure") - regressionAwait(t, writeEntered, "write path before injected failure") - releaseOnce.Do(func() { close(release) }) - regressionAwait(t, readExited, "failed read loop exit") - regressionAwait(t, writeExited, "failed write loop exit") - regressionAwait(t, d.Done(), "concurrent read/write shutdown") - if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { - t.Fatalf("concurrent WriteMsg failure: %v", err) - } - if _, failures, _ := c.transport.snapshot(); failures != 1 { - t.Fatalf("injected transport failed %d times", failures) - } - if c.closeCalls.Load() != 1 { - t.Fatalf("connection closed %d times", c.closeCalls.Load()) - } - }() + c := newRegressionControl(t) + regressionWarm(t, c) + defer c.transport.close() + readEntered, writeEntered := make(chan struct{}), make(chan struct{}) + release := make(chan struct{}) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(release) }) + c.beforeRead = func() error { + close(readEntered) + <-release + return io.EOF + } + c.beforeWrite = func() { + close(writeEntered) + <-release + } + d := NewDispatcher(c) + c.transport.arm() + if err := d.Send(&ReqWorkConn{}); err != nil { + t.Fatal(err) + } + readExited, writeExited := make(chan struct{}), make(chan struct{}) + go func() { + defer close(readExited) + d.readLoop() + }() + go func() { + defer close(writeExited) + d.sendLoop() + }() + regressionAwait(t, readEntered, "read path before injected failure") + regressionAwait(t, writeEntered, "write path before injected failure") + releaseOnce.Do(func() { close(release) }) + regressionAwait(t, readExited, "failed read loop exit") + regressionAwait(t, writeExited, "failed write loop exit") + regressionAwait(t, d.Done(), "concurrent read/write shutdown") + if err := regressionReceiveWrite(t, c); !errors.Is(err, errRegressionWriteTimeout) { + t.Fatalf("concurrent WriteMsg failure: %v", err) + } + if _, failures, _ := c.transport.snapshot(); failures != 1 { + t.Fatalf("injected transport failed %d times", failures) + } + if c.closeCalls.Load() != 1 { + t.Fatalf("connection closed %d times", c.closeCalls.Load()) } } diff --git a/pkg/msg/handler_yamux_test.go b/pkg/msg/handler_yamux_test.go new file mode 100644 index 00000000..25d91a9e --- /dev/null +++ b/pkg/msg/handler_yamux_test.go @@ -0,0 +1,225 @@ +// 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 msg + +import ( + "errors" + "io" + "net" + "sync" + "sync/atomic" + "testing" + "time" + + fmux "github.com/fatedier/yamux" + + "github.com/fatedier/frp/pkg/proto/wire" + netpkg "github.com/fatedier/frp/pkg/util/net" +) + +// Delay one real TCP write without injecting an error. The stream's write +// failure must come from yamux's ConnectionWriteTimeout. +type delayedTCPConn struct { + net.Conn + armed atomic.Bool + entered chan struct{} + release chan struct{} +} + +func (c *delayedTCPConn) Write(p []byte) (int, error) { + if c.armed.Swap(false) { + close(c.entered) + <-c.release + } + return c.Conn.Write(p) +} + +type observedYamuxConn struct { + net.Conn + writes chan error + readEntered chan struct{} + readOnce sync.Once + closed chan struct{} + closeCalls atomic.Int32 +} + +func (c *observedYamuxConn) Read(p []byte) (int, error) { + c.readOnce.Do(func() { close(c.readEntered) }) + return c.Conn.Read(p) +} + +func (c *observedYamuxConn) Write(p []byte) (int, error) { + n, err := c.Conn.Write(p) + c.writes <- err + return n, err +} + +func (c *observedYamuxConn) Close() error { + c.closeCalls.Add(1) + err := c.Conn.Close() + close(c.closed) + return err +} + +func newLoopbackYamuxPair(t *testing.T) (*fmux.Session, *fmux.Session, *delayedTCPConn) { + t.Helper() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer listener.Close() + clientTCP, err := net.Dial("tcp", listener.Addr().String()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { clientTCP.Close() }) + serverTCP, err := listener.Accept() + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { serverTCP.Close() }) + gate := &delayedTCPConn{Conn: clientTCP, entered: make(chan struct{}), release: make(chan struct{})} + config := fmux.DefaultConfig() + config.ConnectionWriteTimeout = 80 * time.Millisecond + config.LogOutput = io.Discard + client, err := fmux.Client(gate, config) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { client.Close() }) + server, err := fmux.Server(serverTCP, config) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { server.Close() }) + return client, server, gate +} + +func loopbackCryptoConn(t *testing.T, conn net.Conn, protocol string, role netpkg.AEADCryptoRole) *Conn { + t.Helper() + var rw io.ReadWriter + var err error + if protocol == wire.ProtocolV2 { + rw, err = netpkg.NewAEADCryptoReadWriter( + conn, regressionKey, role, wire.AEADAlgorithmAES256GCM, []byte("isolated-test-transcript")) + } else { + rw, err = netpkg.NewCryptoReadWriter(conn, regressionKey) + } + if err != nil { + t.Fatal(err) + } + return NewConn(conn, NewReadWriter(rw, protocol)) +} + +func loopbackControlPair(t *testing.T, client, server *fmux.Session, protocol string) (*Conn, *Conn, *observedYamuxConn) { + t.Helper() + local, err := client.OpenStream() + if err != nil { + t.Fatal(err) + } + remote, err := server.AcceptStream() + if err != nil { + t.Fatal(err) + } + observed := &observedYamuxConn{ + Conn: local, writes: make(chan error, 16), readEntered: make(chan struct{}), closed: make(chan struct{}), + } + return loopbackCryptoConn(t, observed, protocol, netpkg.AEADCryptoRoleClient), + loopbackCryptoConn(t, remote, protocol, netpkg.AEADCryptoRoleServer), observed +} + +func loopbackReadReq(t *testing.T, conn *Conn) { + t.Helper() + if err := conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatal(err) + } + defer conn.SetReadDeadline(time.Time{}) + m, err := conn.ReadMsg() + if err != nil { + t.Fatal(err) + } + if _, ok := m.(*ReqWorkConn); !ok { + t.Fatalf("decoded %T, want ReqWorkConn", m) + } +} + +func TestDispatcherYamuxWriteTimeoutInterruptsRead(t *testing.T) { + for _, protocol := range []string{wire.ProtocolV1, wire.ProtocolV2} { + t.Run(protocol, func(t *testing.T) { + client, server, gate := newLoopbackYamuxPair(t) + var releaseOnce sync.Once + defer releaseOnce.Do(func() { close(gate.release) }) + control, peer, observed := loopbackControlPair(t, client, server, protocol) + if err := control.WriteMsg(&ReqWorkConn{}); err != nil { + t.Fatal(err) + } + loopbackReadReq(t, peer) + // Warm-up has completed, including all crypto header writes. + for len(observed.writes) > 0 { + <-observed.writes + } + d := NewDispatcher(control) + readExited := make(chan struct{}) + go func() { + defer close(readExited) + d.readLoop() + }() + regressionAwait(t, observed.readEntered, "control reader entry") + gate.armed.Store(true) + // Queue before starting the writer so this Send cannot race its failure. + if err := d.Send(&ReqWorkConn{}); err != nil { + t.Fatal(err) + } + go d.sendLoop() + regressionAwait(t, gate.entered, "TCP write delay barrier") + select { + case err := <-observed.writes: + if !errors.Is(err, fmux.ErrConnectionWriteTimeout) { + t.Fatalf("stream write error: %v, want yamux connection write timeout", err) + } + t.Logf("production yamux write failed: %v", err) + case <-time.After(time.Second): + t.Fatal("yamux write did not time out") + } + regressionAwait(t, d.stopCh, "write failure stops sends") + if err := d.Send(&ReqWorkConn{}); !errors.Is(err, io.EOF) { + t.Fatalf("Send after stop: %v, want EOF", err) + } + releaseOnce.Do(func() { close(gate.release) }) + regressionAwait(t, observed.closed, "control Close returns") + // peer never closes its write half: local yamux Close alone cannot + // interrupt Read. Neither a peer FIN nor session cleanup may help. + regressionAwait(t, d.Done(), "Done without peer FIN") + regressionAwait(t, readExited, "control readLoop exit without peer FIN") + if observed.closeCalls.Load() != 1 { + t.Fatalf("control closed %d times, want once", observed.closeCalls.Load()) + } + if client.IsClosed() || server.IsClosed() { + t.Fatal("yamux session unexpectedly closed") + } + fresh, freshPeer, freshObserved := loopbackControlPair(t, client, server, protocol) + for range 5 { + if err := fresh.WriteMsg(&ReqWorkConn{}); err != nil { + t.Fatal(err) + } + loopbackReadReq(t, freshPeer) + for len(freshObserved.writes) > 0 { + <-freshObserved.writes + } + } + t.Log("Done/readLoop exited without peer FIN; Close once; same yamux session decrypted five fresh messages") + }) + } +}