k8s-operator,sessionrecording: fixing race condition between resize (#16454)

messages and cast headers when recording `kubectl attach` sessions

Updates #16490

Signed-off-by: chaosinthecrd <tom@tmlabs.co.uk>

k8s-operator/sessionrecording/ws/conn.go

@@ -3,12 +3,13 @@
 
 //go:build !plan9
 
-// package ws has functionality to parse 'kubectl exec' sessions streamed using
+// package ws has functionality to parse 'kubectl exec/attach' sessions streamed using
 // WebSocket protocol.
 package ws
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -24,31 +25,53 @@ import (
 )
 
 // New wraps the provided network connection and returns a connection whose reads and writes will get triggered as data is received on the hijacked connection.
-// The connection must be a hijacked connection for a 'kubectl exec' session using WebSocket protocol and a *.channel.k8s.io subprotocol.
+// The connection must be a hijacked connection for a 'kubectl exec/attach' session using WebSocket protocol and a *.channel.k8s.io subprotocol.
 // The hijacked connection is used to transmit *.channel.k8s.io streams between Kubernetes client ('kubectl') and the destination proxy controlled by Kubernetes.
 // Data read from the underlying network connection is data sent via one of the streams from the client to the container.
 // Data written to the underlying connection is data sent from the container to the client.
 // We parse the data and send everything for the stdout/stderr streams to the configured tsrecorder as an asciinema recording with the provided header.
 // https://github.com/kubernetes/enhancements/tree/master/keps/sig-api-machinery/4006-transition-spdy-to-websockets#proposal-new-remotecommand-sub-protocol-version---v5channelk8sio
-func New(c net.Conn, rec *tsrecorder.Client, ch sessionrecording.CastHeader, hasTerm bool, log *zap.SugaredLogger) net.Conn {
-	return &conn{
-		Conn:               c,
-		rec:                rec,
-		ch:                 ch,
-		hasTerm:            hasTerm,
-		log:                log,
-		initialTermSizeSet: make(chan struct{}, 1),
+func New(ctx context.Context, c net.Conn, rec *tsrecorder.Client, ch sessionrecording.CastHeader, hasTerm bool, log *zap.SugaredLogger) (net.Conn, error) {
+	lc := &conn{
+		Conn:                  c,
+		ctx:                   ctx,
+		rec:                   rec,
+		ch:                    ch,
+		hasTerm:               hasTerm,
+		log:                   log,
+		initialCastHeaderSent: make(chan struct{}, 1),
 	}
+
+	// if there is no term, we don't need to wait for a resize message
+	if !hasTerm {
+		var err error
+		lc.writeCastHeaderOnce.Do(func() {
+			// If this is a session with a terminal attached,
+			// we must wait for the terminal width and
+			// height to be parsed from a resize message
+			// before sending CastHeader, else tsrecorder
+			// will not be able to play this recording.
+			err = lc.rec.WriteCastHeader(ch)
+			close(lc.initialCastHeaderSent)
+		})
+		if err != nil {
+			return nil, fmt.Errorf("error writing CastHeader: %w", err)
+		}
+	}
+
+	return lc, nil
 }
 
 // conn is a wrapper around net.Conn. It reads the bytestream
-// for a 'kubectl exec' session, sends session recording data to the configured
+// for a 'kubectl exec/attach' session, sends session recording data to the configured
 // recorder and forwards the raw bytes to the original destination.
 // A new conn is created per session.
-// conn only knows to how to read a 'kubectl exec' session that is streamed using WebSocket protocol.
+// conn only knows to how to read a 'kubectl exec/attach' session that is streamed using WebSocket protocol.
 // https://www.rfc-editor.org/rfc/rfc6455
 type conn struct {
 	net.Conn
+
+	ctx context.Context
 	// rec knows how to send data to a tsrecorder instance.
 	rec *tsrecorder.Client
 
@@ -56,7 +79,7 @@ type conn struct {
 	// CastHeader must be sent before any payload. If the session has a
 	// terminal attached, the CastHeader must have '.Width' and '.Height'
 	// fields set for the tsrecorder UI to be able to play the recording.
-	// For 'kubectl exec' sessions, terminal width and height are sent as a
+	// For 'kubectl exec/attach' sessions, terminal width and height are sent as a
 	// resize message on resize stream from the client when the session
 	// starts as well as at any time the client detects a terminal change.
 	// We can intercept the resize message on Read calls. As there is no
@@ -72,15 +95,10 @@ type conn struct {
 	// writeCastHeaderOnce is used to ensure CastHeader gets sent to tsrecorder once.
 	writeCastHeaderOnce sync.Once
 	hasTerm             bool // whether the session has TTY attached
-	// initialTermSizeSet channel gets sent a value once, when the Read has
-	// received a resize message and set the initial terminal size. It must
-	// be set to a buffered channel to prevent Reads being blocked on the
-	// first stdout/stderr write reading from the channel.
-	initialTermSizeSet chan struct{}
-	// sendInitialTermSizeSetOnce is used to ensure that a value is sent to
-	// initialTermSizeSet channel only once, when the initial resize message
-	// is received.
-	sendInitialTermSizeSetOnce sync.Once
+	// initialCastHeaderSent is a boolean that is set to ensure that the cast
+	// header is the first thing that is streamed to the session recorder.
+	// Otherwise the stream will fail.
+	initialCastHeaderSent chan struct{}
 
 	log *zap.SugaredLogger
 
@@ -171,9 +189,10 @@ func (c *conn) Read(b []byte) (int, error) {
 	c.readBuf.Next(len(readMsg.raw))
 
 	if readMsg.isFinalized && !c.readMsgIsIncomplete() {
+		// we want to send stream resize messages for terminal sessions
 		// Stream IDs for websocket streams are static.
 		// https://github.com/kubernetes/client-go/blob/v0.30.0-rc.1/tools/remotecommand/websocket.go#L218
-		if readMsg.streamID.Load() == remotecommand.StreamResize {
+		if readMsg.streamID.Load() == remotecommand.StreamResize && c.hasTerm {
 			var msg tsrecorder.ResizeMsg
 			if err = json.Unmarshal(readMsg.payload, &msg); err != nil {
 				return 0, fmt.Errorf("error umarshalling resize message: %w", err)
@@ -182,22 +201,29 @@ func (c *conn) Read(b []byte) (int, error) {
 			c.ch.Width = msg.Width
 			c.ch.Height = msg.Height
 
-			// If this is initial resize message, the width and
-			// height will be sent in the CastHeader. If this is a
-			// subsequent resize message, we need to send asciinema
-			// resize message.
 			var isInitialResize bool
-			c.sendInitialTermSizeSetOnce.Do(func() {
+			c.writeCastHeaderOnce.Do(func() {
 				isInitialResize = true
-				close(c.initialTermSizeSet) // unblock sending of CastHeader
+				// If this is a session with a terminal attached,
+				// we must wait for the terminal width and
+				// height to be parsed from a resize message
+				// before sending CastHeader, else tsrecorder
+				// will not be able to play this recording.
+				err = c.rec.WriteCastHeader(c.ch)
+				close(c.initialCastHeaderSent)
 			})
+			if err != nil {
+				return 0, fmt.Errorf("error writing CastHeader: %w", err)
+			}
+
 			if !isInitialResize {
-				if err := c.rec.WriteResize(c.ch.Height, c.ch.Width); err != nil {
+				if err := c.rec.WriteResize(msg.Height, msg.Width); err != nil {
 					return 0, fmt.Errorf("error writing resize message: %w", err)
 				}
 			}
 		}
 	}
+
 	c.currentReadMsg = readMsg
 	return n, nil
 }
@@ -244,39 +270,33 @@ func (c *conn) Write(b []byte) (int, error) {
 		c.log.Errorf("write: parsing a message errored: %v", err)
 		return 0, fmt.Errorf("write: error parsing message: %v", err)
 	}
+
 	c.currentWriteMsg = writeMsg
 	if !ok { // incomplete fragment
 		return len(b), nil
 	}
+
 	c.writeBuf.Next(len(writeMsg.raw)) // advance frame
 
 	if len(writeMsg.payload) != 0 && writeMsg.isFinalized {
 		if writeMsg.streamID.Load() == remotecommand.StreamStdOut || writeMsg.streamID.Load() == remotecommand.StreamStdErr {
-			var err error
-			c.writeCastHeaderOnce.Do(func() {
-				// If this is a session with a terminal attached,
-				// we must wait for the terminal width and
-				// height to be parsed from a resize message
-				// before sending CastHeader, else tsrecorder
-				// will not be able to play this recording.
-				if c.hasTerm {
-					c.log.Debug("waiting for terminal size to be set before starting to send recorded data")
-					<-c.initialTermSizeSet
+			// we must wait for confirmation that the initial cast header was sent before proceeding with any more writes
+			select {
+			case <-c.ctx.Done():
+				return 0, c.ctx.Err()
+			case <-c.initialCastHeaderSent:
+				if err := c.rec.WriteOutput(writeMsg.payload); err != nil {
+					return 0, fmt.Errorf("error writing message to recorder: %w", err)
 				}
-				err = c.rec.WriteCastHeader(c.ch)
-			})
-			if err != nil {
-				return 0, fmt.Errorf("error writing CastHeader: %w", err)
-			}
-			if err := c.rec.WriteOutput(writeMsg.payload); err != nil {
-				return 0, fmt.Errorf("error writing message to recorder: %v", err)
 			}
 		}
 	}
+
 	_, err = c.Conn.Write(c.currentWriteMsg.raw)
 	if err != nil {
 		c.log.Errorf("write: error writing to conn: %v", err)
 	}
+
 	return len(b), nil
 }
 
@@ -321,6 +341,7 @@ func (c *conn) writeMsgIsIncomplete() bool {
 func (c *conn) readMsgIsIncomplete() bool {
 	return c.currentReadMsg != nil && !c.currentReadMsg.isFinalized
 }
+
 func (c *conn) curReadMsgType() (messageType, error) {
 	if c.currentReadMsg != nil {
 		return c.currentReadMsg.typ, nil

k8s-operator/sessionrecording/ws/conn_test.go

@@ -6,9 +6,11 @@
 package ws
 
 import (
+	"context"
 	"fmt"
 	"reflect"
 	"testing"
+	"time"
 
 	"go.uber.org/zap"
 	"k8s.io/apimachinery/pkg/util/remotecommand"
@@ -26,46 +28,69 @@ func Test_conn_Read(t *testing.T) {
 	// Resize stream ID + {"width": 10, "height": 20}
 	testResizeMsg := []byte{byte(remotecommand.StreamResize), 0x7b, 0x22, 0x77, 0x69, 0x64, 0x74, 0x68, 0x22, 0x3a, 0x31, 0x30, 0x2c, 0x22, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74, 0x22, 0x3a, 0x32, 0x30, 0x7d}
 	lenResizeMsgPayload := byte(len(testResizeMsg))
-
+	cl := tstest.NewClock(tstest.ClockOpts{})
 	tests := []struct {
-		name       string
-		inputs     [][]byte
-		wantWidth  int
-		wantHeight int
+		name                 string
+		inputs               [][]byte
+		wantCastHeaderWidth  int
+		wantCastHeaderHeight int
+		wantRecorded         []byte
 	}{
 		{
 			name:   "single_read_control_message",
 			inputs: [][]byte{{0x88, 0x0}},
 		},
 		{
-			name:       "single_read_resize_message",
-			inputs:     [][]byte{append([]byte{0x82, lenResizeMsgPayload}, testResizeMsg...)},
-			wantWidth:  10,
-			wantHeight: 20,
+			name:                 "single_read_resize_message",
+			inputs:               [][]byte{append([]byte{0x82, lenResizeMsgPayload}, testResizeMsg...)},
+			wantCastHeaderWidth:  10,
+			wantCastHeaderHeight: 20,
+			wantRecorded:         fakes.AsciinemaCastHeaderMsg(t, 10, 20),
 		},
 		{
-			name:       "two_reads_resize_message",
-			inputs:     [][]byte{{0x2, 0x9, 0x4, 0x7b, 0x22, 0x77, 0x69, 0x64, 0x74, 0x68, 0x22}, {0x80, 0x11, 0x4, 0x3a, 0x31, 0x30, 0x2c, 0x22, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74, 0x22, 0x3a, 0x32, 0x30, 0x7d}},
-			wantWidth:  10,
-			wantHeight: 20,
+			name: "resize_data_frame_many",
+			inputs: [][]byte{
+				append([]byte{0x82, lenResizeMsgPayload}, testResizeMsg...),
+				append([]byte{0x82, lenResizeMsgPayload}, testResizeMsg...),
+			},
+			wantRecorded:         append(fakes.AsciinemaCastHeaderMsg(t, 10, 20), fakes.AsciinemaCastResizeMsg(t, 10, 20)...),
+			wantCastHeaderWidth:  10,
+			wantCastHeaderHeight: 20,
 		},
 		{
-			name:       "three_reads_resize_message_with_split_fragment",
-			inputs:     [][]byte{{0x2, 0x9, 0x4, 0x7b, 0x22, 0x77, 0x69, 0x64, 0x74, 0x68, 0x22}, {0x80, 0x11, 0x4, 0x3a, 0x31, 0x30, 0x2c, 0x22, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74}, {0x22, 0x3a, 0x32, 0x30, 0x7d}},
-			wantWidth:  10,
-			wantHeight: 20,
+			name:                 "two_reads_resize_message",
+			inputs:               [][]byte{{0x2, 0x9, 0x4, 0x7b, 0x22, 0x77, 0x69, 0x64, 0x74, 0x68, 0x22}, {0x80, 0x11, 0x4, 0x3a, 0x31, 0x30, 0x2c, 0x22, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74, 0x22, 0x3a, 0x32, 0x30, 0x7d}},
+			wantCastHeaderWidth:  10,
+			wantCastHeaderHeight: 20,
+			wantRecorded:         fakes.AsciinemaCastHeaderMsg(t, 10, 20),
+		},
+		{
+			name:                 "three_reads_resize_message_with_split_fragment",
+			inputs:               [][]byte{{0x2, 0x9, 0x4, 0x7b, 0x22, 0x77, 0x69, 0x64, 0x74, 0x68, 0x22}, {0x80, 0x11, 0x4, 0x3a, 0x31, 0x30, 0x2c, 0x22, 0x68, 0x65, 0x69, 0x67, 0x68, 0x74}, {0x22, 0x3a, 0x32, 0x30, 0x7d}},
+			wantCastHeaderWidth:  10,
+			wantCastHeaderHeight: 20,
+			wantRecorded:         fakes.AsciinemaCastHeaderMsg(t, 10, 20),
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			l := zl.Sugar()
 			tc := &fakes.TestConn{}
+			sr := &fakes.TestSessionRecorder{}
+			rec := tsrecorder.New(sr, cl, cl.Now(), true, zl.Sugar())
 			tc.ResetReadBuf()
+
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
 			c := &conn{
-				Conn: tc,
-				log:  zl.Sugar(),
+				ctx:                   ctx,
+				Conn:                  tc,
+				log:                   l,
+				hasTerm:               true,
+				initialCastHeaderSent: make(chan struct{}),
+				rec:                   rec,
 			}
 			for i, input := range tt.inputs {
-				c.initialTermSizeSet = make(chan struct{})
 				if err := tc.WriteReadBufBytes(input); err != nil {
 					t.Fatalf("writing bytes to test conn: %v", err)
 				}
@@ -75,14 +100,20 @@ func Test_conn_Read(t *testing.T) {
 					return
 				}
 			}
-			if tt.wantHeight != 0 || tt.wantWidth != 0 {
-				if tt.wantWidth != c.ch.Width {
-					t.Errorf("wants width: %v, got %v", tt.wantWidth, c.ch.Width)
+
+			if tt.wantCastHeaderHeight != 0 || tt.wantCastHeaderWidth != 0 {
+				if tt.wantCastHeaderWidth != c.ch.Width {
+					t.Errorf("wants width: %v, got %v", tt.wantCastHeaderWidth, c.ch.Width)
 				}
-				if tt.wantHeight != c.ch.Height {
-					t.Errorf("want height: %v, got %v", tt.wantHeight, c.ch.Height)
+				if tt.wantCastHeaderHeight != c.ch.Height {
+					t.Errorf("want height: %v, got %v", tt.wantCastHeaderHeight, c.ch.Height)
 				}
 			}
+
+			gotRecorded := sr.Bytes()
+			if !reflect.DeepEqual(gotRecorded, tt.wantRecorded) {
+				t.Errorf("expected bytes not recorded, wants\n%v\ngot\n%v", string(tt.wantRecorded), string(gotRecorded))
+			}
 		})
 	}
 }
@@ -94,15 +125,11 @@ func Test_conn_Write(t *testing.T) {
 	}
 	cl := tstest.NewClock(tstest.ClockOpts{})
 	tests := []struct {
-		name              string
-		inputs            [][]byte
-		wantForwarded     []byte
-		wantRecorded      []byte
-		firstWrite        bool
-		width             int
-		height            int
-		hasTerm           bool
-		sendInitialResize bool
+		name          string
+		inputs        [][]byte
+		wantForwarded []byte
+		wantRecorded  []byte
+		hasTerm       bool
 	}{
 		{
 			name:          "single_write_control_frame",
@@ -130,10 +157,7 @@ func Test_conn_Write(t *testing.T) {
 			name:          "single_write_stdout_data_message_with_cast_header",
 			inputs:        [][]byte{{0x82, 0x3, 0x1, 0x7, 0x8}},
 			wantForwarded: []byte{0x82, 0x3, 0x1, 0x7, 0x8},
-			wantRecorded:  append(fakes.AsciinemaResizeMsg(t, 10, 20), fakes.CastLine(t, []byte{0x7, 0x8}, cl)...),
-			width:         10,
-			height:        20,
-			firstWrite:    true,
+			wantRecorded:  fakes.CastLine(t, []byte{0x7, 0x8}, cl),
 		},
 		{
 			name:          "two_writes_stdout_data_message",
@@ -148,15 +172,11 @@ func Test_conn_Write(t *testing.T) {
 			wantRecorded:  fakes.CastLine(t, []byte{0x7, 0x8, 0x1, 0x2, 0x3, 0x4, 0x5}, cl),
 		},
 		{
-			name:              "three_writes_stdout_data_message_with_split_fragment_cast_header_with_terminal",
-			inputs:            [][]byte{{0x2, 0x3, 0x1, 0x7, 0x8}, {0x80, 0x6, 0x1, 0x1, 0x2, 0x3}, {0x4, 0x5}},
-			wantForwarded:     []byte{0x2, 0x3, 0x1, 0x7, 0x8, 0x80, 0x6, 0x1, 0x1, 0x2, 0x3, 0x4, 0x5},
-			wantRecorded:      append(fakes.AsciinemaResizeMsg(t, 10, 20), fakes.CastLine(t, []byte{0x7, 0x8, 0x1, 0x2, 0x3, 0x4, 0x5}, cl)...),
-			height:            20,
-			width:             10,
-			hasTerm:           true,
-			firstWrite:        true,
-			sendInitialResize: true,
+			name:          "three_writes_stdout_data_message_with_split_fragment_cast_header_with_terminal",
+			inputs:        [][]byte{{0x2, 0x3, 0x1, 0x7, 0x8}, {0x80, 0x6, 0x1, 0x1, 0x2, 0x3}, {0x4, 0x5}},
+			wantForwarded: []byte{0x2, 0x3, 0x1, 0x7, 0x8, 0x80, 0x6, 0x1, 0x1, 0x2, 0x3, 0x4, 0x5},
+			wantRecorded:  fakes.CastLine(t, []byte{0x7, 0x8, 0x1, 0x2, 0x3, 0x4, 0x5}, cl),
+			hasTerm:       true,
 		},
 	}
 	for _, tt := range tests {
@@ -164,24 +184,22 @@ func Test_conn_Write(t *testing.T) {
 			tc := &fakes.TestConn{}
 			sr := &fakes.TestSessionRecorder{}
 			rec := tsrecorder.New(sr, cl, cl.Now(), true, zl.Sugar())
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
 			c := &conn{
-				Conn: tc,
-				log:  zl.Sugar(),
-				ch: sessionrecording.CastHeader{
-					Width:  tt.width,
-					Height: tt.height,
-				},
-				rec:                rec,
-				initialTermSizeSet: make(chan struct{}),
-				hasTerm:            tt.hasTerm,
-			}
-			if !tt.firstWrite {
-				// This test case does not intend to test that cast header gets written once.
-				c.writeCastHeaderOnce.Do(func() {})
-			}
-			if tt.sendInitialResize {
-				close(c.initialTermSizeSet)
+				Conn:                  tc,
+				ctx:                   ctx,
+				log:                   zl.Sugar(),
+				ch:                    sessionrecording.CastHeader{},
+				rec:                   rec,
+				initialCastHeaderSent: make(chan struct{}),
+				hasTerm:               tt.hasTerm,
 			}
+
+			c.writeCastHeaderOnce.Do(func() {
+				close(c.initialCastHeaderSent)
+			})
+
 			for i, input := range tt.inputs {
 				_, err := c.Write(input)
 				if err != nil {