From 223b44336cba7cd21e6dd05ea3b991768c8139fa Mon Sep 17 00:00:00 2001 From: fatedier Date: Sun, 9 Aug 2026 19:12:12 +0800 Subject: [PATCH] ssh: serialize tunnel channel writes (#5473) --- pkg/ssh/server.go | 5 +++++ pkg/ssh/server_test.go | 44 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/pkg/ssh/server.go b/pkg/ssh/server.go index c042b9c1..8a69fc30 100644 --- a/pkg/ssh/server.go +++ b/pkg/ssh/server.go @@ -70,6 +70,7 @@ type TunnelServer struct { sshConn *ssh.ServerConn sc *ssh.ServerConfig firstChannel ssh.Channel + firstChannelMu sync.Mutex vc *virtual.Client peerServerListener *netpkg.InternalListener @@ -191,6 +192,8 @@ func (s *TunnelServer) Run() error { } func (s *TunnelServer) writeToClient(data string) { + s.firstChannelMu.Lock() + defer s.firstChannelMu.Unlock() if s.firstChannel == nil { return } @@ -304,9 +307,11 @@ func (s *TunnelServer) handleNewChannel(channel ssh.NewChannel, extraPayloadCh c if err != nil { return } + s.firstChannelMu.Lock() if s.firstChannel == nil { s.firstChannel = ch } + s.firstChannelMu.Unlock() go s.keepAlive(ch) for req := range reqs { diff --git a/pkg/ssh/server_test.go b/pkg/ssh/server_test.go index 78421a55..56b1ba5b 100644 --- a/pkg/ssh/server_test.go +++ b/pkg/ssh/server_test.go @@ -16,7 +16,11 @@ package ssh import ( "encoding/binary" + "io" + "sync" + "sync/atomic" "testing" + "time" "github.com/stretchr/testify/require" cryptossh "golang.org/x/crypto/ssh" @@ -69,3 +73,43 @@ func TestParseExecPayloadRejectsMalformedPayloads(t *testing.T) { }) } } + +type trackingChannel struct { + active atomic.Int32 + concurrent atomic.Bool +} + +func (c *trackingChannel) Read([]byte) (int, error) { return 0, io.EOF } + +func (c *trackingChannel) Write(p []byte) (int, error) { + if c.active.Add(1) != 1 { + c.concurrent.Store(true) + } + time.Sleep(time.Millisecond) + c.active.Add(-1) + return len(p), nil +} + +func (c *trackingChannel) Close() error { return nil } +func (c *trackingChannel) CloseWrite() error { return nil } +func (c *trackingChannel) SendRequest(string, bool, []byte) (bool, error) { return false, nil } +func (c *trackingChannel) Stderr() io.ReadWriter { return nil } + +func TestWriteToClientSerializesChannelWrites(t *testing.T) { + channel := &trackingChannel{} + s := &TunnelServer{firstChannel: channel} + start := make(chan struct{}) + var wg sync.WaitGroup + for range 8 { + wg.Go(func() { + <-start + s.writeToClient("message") + }) + } + close(start) + wg.Wait() + + if channel.concurrent.Load() { + t.Fatal("channel writes were concurrent") + } +}