ssh: serialize tunnel channel writes (#5473)

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 {

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")
+	}
+}