Refactor exec sessions into a shared SSH engine

internal/controller/api_vms_exec.go

@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"net"
 	"net/http"
 	"strconv"
@@ -70,33 +69,27 @@ func (controller *Controller) execVMLegacy(
 		return responderImpl
 	}
 
-	// Establish a port-forwarding connection to a VM's SSH port
+	session, err := controller.newSSHExecSession(
-	portForwardConn, err := retry.NewWithData[net.Conn](
+		ctx,
-		retry.Context(waitContext),
+		waitContext,
-		retry.DelayType(retry.FixedDelay),
+		vm,
-		retry.Delay(time.Second),
+		execSessionKey{vmName: name},
-		retry.Attempts(0),
+		command,
-		retry.LastErrorOnly(true),
+		stdin,
-	).Do(func() (net.Conn, error) {
+		nil,
-		return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22)
+		legacyExecSessionPolicy,
-	})
+	)
 	if err != nil {
 		return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
 	}
-	defer portForwardConn.Close()
-
-	// Establish an SSH connection to a VM
-	exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin)
-	if err != nil {
-		return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("failed to establish SSH connection to a VM: %v", err))
-	}
-	defer exec.Close()
 
 	// Upgrade HTTP request to a WebSocket connection
 	wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
 		OriginPatterns: []string{"*"},
 	})
 	if err != nil {
+		session.closeIfUnused()
+
 		return responder.Error(err)
 	}
 	defer func() {
@@ -107,66 +100,7 @@ func (controller *Controller) execVMLegacy(
 		_ = wsConn.CloseNow()
 	}()
 
-	// Read WebSocket frames
+	return controller.serveExecSession(ctx, wsConn, session)
-	readFramesErrCh := make(chan error, 1)
-	go func() {
-		readFramesErrCh <- controller.readFrames(ctx, wsConn, exec.Stdin())
-	}()
-
-	// Run the command
-	sshErrCh := make(chan error, 1)
-	outgoingFrames := make(chan *execstream.Frame)
-	go func() {
-		sshErrCh <- exec.Run(ctx, command, outgoingFrames)
-	}()
-
-	for {
-		select {
-		case readFramesErr := <-readFramesErrCh:
-			controller.logger.Warnf("failed to read and process frames from WebSocket: %v", readFramesErr)
-
-			return responder.Empty()
-		case outgoingFrame := <-outgoingFrames:
-			if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil {
-				controller.logger.Warnf("failed to write WebSocket frame to the client: %v", err)
-
-				return responder.Empty()
-			}
-		case sshErr := <-sshErrCh:
-			if sshErr != nil {
-				if err := execstream.WriteFrame(ctx, wsConn, &execstream.Frame{
-					Type:  execstream.FrameTypeError,
-					Error: sshErr.Error(),
-				}); err != nil {
-					controller.logger.Warnf("exec: failed to write error frame to WebSocket: %v", err)
-				}
-			}
-
-			if err := wsConn.Close(websocket.StatusNormalClosure, "Command finished"); err != nil {
-				controller.logger.Warnf("exec: failed to close WebSocket cleanly: %v", err)
-			}
-
-			if readFramesErrCh != nil {
-				// Read() on a WebSocket should unblock shortly after calling Close()
-				<-readFramesErrCh
-			}
-
-			return responder.Empty()
-		case <-time.After(controller.pingInterval):
-			pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
-
-			if err := wsConn.Ping(pingCtx); err != nil {
-				controller.logger.Warnf("port forwarding: failed to ping the client, "+
-					"connection might time out: %v", err)
-			}
-
-			pingCtxCancel()
-		case <-ctx.Done():
-			controller.logger.Warnf("client disconnected prematurely")
-
-			return responder.Empty()
-		}
-	}
 }
 
 func (controller *Controller) execVMReconnectable(
@@ -204,6 +138,52 @@ func (controller *Controller) execVMReconnectable(
 
 		var err error
 		session, _, err = controller.execSessions.getOrCreate(waitContext, key, func() (*execSession, error) {
+			return controller.newSSHExecSession(
+				ctx,
+				waitContext,
+				vm,
+				key,
+				command,
+				stdin,
+				controller.execSessions,
+				reconnectableExecSessionPolicy,
+			)
+		})
+		if err != nil {
+			return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
+		}
+
+		if !session.commandMatches(command) {
+			return responder.JSON(http.StatusConflict,
+				NewErrorResponse("exec session %q is already running a different command", sessionID))
+		}
+	}
+
+	wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
+		OriginPatterns: []string{"*"},
+	})
+	if err != nil {
+		session.closeIfUnused()
+
+		return responder.Error(err)
+	}
+	defer func() {
+		_ = wsConn.CloseNow()
+	}()
+
+	return controller.serveExecSession(ctx, wsConn, session)
+}
+
+func (controller *Controller) newSSHExecSession(
+	ctx *gin.Context,
+	waitContext context.Context,
+	vm *v1.VM,
+	key execSessionKey,
+	command string,
+	stdin bool,
+	registry *execSessionRegistry,
+	policy execSessionPolicy,
+) (*execSession, error) {
 	portForwardConn, err := retry.NewWithData[net.Conn](
 		retry.Context(waitContext),
 		retry.DelayType(retry.FixedDelay),
@@ -221,7 +201,7 @@ func (controller *Controller) execVMReconnectable(
 	if err != nil {
 		_ = portForwardConn.Close()
 
-				return nil, err
+		return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err)
 	}
 
 	return newExecSession(
@@ -229,30 +209,17 @@ func (controller *Controller) execVMReconnectable(
 		command,
 		exec,
 		portForwardConn,
-				controller.execSessions,
+		registry,
 		controller.execSessionExitTTL,
+		policy,
 	), nil
-		})
-		if err != nil {
-			return responder.JSON(http.StatusServiceUnavailable, NewErrorResponse("%v", err))
 }
 
-		if !session.commandMatches(command) {
+func (controller *Controller) serveExecSession(
-			return responder.JSON(http.StatusConflict,
+	ctx *gin.Context,
-				NewErrorResponse("exec session %q is already running a different command", sessionID))
+	wsConn *websocket.Conn,
-		}
+	session *execSession,
-	}
+) responder.Responder {
-
-	wsConn, err := websocket.Accept(ctx.Writer, ctx.Request, &websocket.AcceptOptions{
-		OriginPatterns: []string{"*"},
-	})
-	if err != nil {
-		return responder.Error(err)
-	}
-	defer func() {
-		_ = wsConn.CloseNow()
-	}()
-
 	subscriber, err := session.attach()
 	if err != nil {
 		_ = wsConn.Close(websocket.StatusNormalClosure, err.Error())
@@ -260,10 +227,11 @@ func (controller *Controller) execVMReconnectable(
 		return responder.Empty()
 	}
 	defer session.detach(subscriber)
+	session.start()
 
 	readFramesErrCh := make(chan error, 1)
 	go func() {
-		readFramesErrCh <- controller.readReconnectableFrames(ctx, wsConn, session, subscriber)
+		readFramesErrCh <- controller.readExecSessionFrames(ctx, wsConn, session, subscriber)
 	}()
 
 	for {
@@ -272,7 +240,7 @@ func (controller *Controller) execVMReconnectable(
 			if readFramesErr != nil &&
 				!errors.Is(readFramesErr, errExecSessionDetached) &&
 				!errors.Is(readFramesErr, errExecSessionClosed) {
-				controller.logger.Warnf("failed to read and process reconnectable exec frames from WebSocket: %v",
+				controller.logger.Warnf("failed to read and process exec frames from WebSocket: %v",
 					readFramesErr)
 			}
 
@@ -287,7 +255,7 @@ func (controller *Controller) execVMReconnectable(
 			}
 
 			if err := execstream.WriteFrame(ctx, wsConn, outgoingFrame); err != nil {
-				controller.logger.Warnf("failed to write reconnectable exec frame to the client: %v", err)
+				controller.logger.Warnf("failed to write exec frame to the client: %v", err)
 
 				return responder.Empty()
 			}
@@ -295,7 +263,7 @@ func (controller *Controller) execVMReconnectable(
 			pingCtx, pingCtxCancel := context.WithTimeout(ctx, 5*time.Second)
 
 			if err := wsConn.Ping(pingCtx); err != nil {
-				controller.logger.Warnf("reconnectable exec: failed to ping the client, "+
+				controller.logger.Warnf("exec: failed to ping the client, "+
 					"connection might time out: %v", err)
 			}
 
@@ -313,7 +281,7 @@ var (
 	errExecSessionClosed   = errors.New("exec session closed")
 )
 
-func (controller *Controller) readReconnectableFrames(
+func (controller *Controller) readExecSessionFrames(
 	ctx context.Context,
 	wsConn *websocket.Conn,
 	session *execSession,
@@ -346,12 +314,28 @@ func (controller *Controller) readReconnectableFrames(
 				return fmt.Errorf("failed to handle %q frame: %w", frame.Type, err)
 			}
 		case execstream.FrameTypeHistory:
+			if !session.policy.replayEnabled {
+				return fmt.Errorf("unexpected frame type received: %q", frame.Type)
+			}
+
 			session.sendHistory(subscriber, frame.Watermark)
 		case execstream.FrameTypeAck:
+			if !session.policy.replayEnabled {
+				return fmt.Errorf("unexpected frame type received: %q", frame.Type)
+			}
+
 			session.ack(frame.Watermark)
 		case execstream.FrameTypeDetach:
+			if !session.policy.replayEnabled {
+				return fmt.Errorf("unexpected frame type received: %q", frame.Type)
+			}
+
 			return errExecSessionDetached
 		case execstream.FrameTypeClose:
+			if !session.policy.replayEnabled {
+				return fmt.Errorf("unexpected frame type received: %q", frame.Type)
+			}
+
 			session.close()
 
 			return errExecSessionClosed
@@ -360,57 +344,3 @@ func (controller *Controller) readReconnectableFrames(
 		}
 	}
 }
-
-func (controller *Controller) readFrames(
-	ctx context.Context,
-	wsConn *websocket.Conn,
-	stdinHandle io.WriteCloser,
-) error {
-	for {
-		var frame execstream.Frame
-
-		messageType, payloadBytes, err := wsConn.Read(ctx)
-		if err != nil {
-			var closeErr websocket.CloseError
-			if errors.As(err, &closeErr) && closeErr.Code == websocket.StatusNormalClosure {
-				return nil
-			}
-
-			return fmt.Errorf("failed to read next frame from WebSocket: %w", err)
-		}
-
-		if messageType != websocket.MessageText {
-			continue
-		}
-
-		if err := json.Unmarshal(payloadBytes, &frame); err != nil {
-			return err
-		}
-
-		switch frame.Type {
-		case execstream.FrameTypeStdin:
-			if stdinHandle == nil {
-				return fmt.Errorf("failed to handle %q frame: this exec session "+
-					"has no stdin is enabled or already closed", frame.Type)
-			}
-
-			if len(frame.Data) == 0 {
-				if err := stdinHandle.Close(); err != nil {
-					return fmt.Errorf("failed to handle %q frame: failed to close "+
-						"stdin: %w", frame.Type, err)
-				}
-
-				stdinHandle = nil
-
-				continue
-			}
-
-			if _, err := stdinHandle.Write(frame.Data); err != nil {
-				return fmt.Errorf("failed to handle %q frame: failed to write "+
-					"to stdin: %w", frame.Type, err)
-			}
-		default:
-			return fmt.Errorf("unexpected frame type received: %q", frame.Type)
-		}
-	}
-}

internal/controller/exec_sessions.go

@@ -13,6 +13,22 @@ import (
 
 const execSessionReplayBufferBytes = 4 * 1024 * 1024
 
+type execSessionPolicy struct {
+	closeOnDetach   bool
+	retainAfterExit bool
+	replayEnabled   bool
+}
+
+var (
+	legacyExecSessionPolicy = execSessionPolicy{
+		closeOnDetach: true,
+	}
+	reconnectableExecSessionPolicy = execSessionPolicy{
+		retainAfterExit: true,
+		replayEnabled:   true,
+	}
+)
+
 type sshExecRunner interface {
 	Stdin() io.WriteCloser
 	Run(ctx context.Context, command string, outgoingFrames chan<- *execstream.Frame) error
@@ -154,6 +170,7 @@ type execSession struct {
 	transport net.Conn
 	registry  *execSessionRegistry
 	exitTTL   time.Duration
+	policy    execSessionPolicy
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -166,10 +183,12 @@ type execSession struct {
 	bufferBytes    int
 	nextWatermark  uint64
 	ackedWatermark uint64
+	started        bool
 	finished       bool
 	closed         bool
 	expiryTimer    *time.Timer
 
+	startOnce sync.Once
 	done      chan struct{}
 	doneOnce  sync.Once
 }
@@ -181,6 +200,7 @@ func newExecSession(
 	transport net.Conn,
 	registry *execSessionRegistry,
 	exitTTL time.Duration,
+	policy execSessionPolicy,
 ) *execSession {
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -191,6 +211,7 @@ func newExecSession(
 		transport:   transport,
 		registry:    registry,
 		exitTTL:     exitTTL,
+		policy:      policy,
 		ctx:         ctx,
 		cancel:      cancel,
 		stdin:       exec.Stdin(),
@@ -198,8 +219,6 @@ func newExecSession(
 		done:        make(chan struct{}),
 	}
 
-	go session.run()
-
 	return session
 }
 
@@ -207,6 +226,31 @@ func (session *execSession) commandMatches(command string) bool {
 	return command == "" || session.command == command
 }
 
+func (session *execSession) start() {
+	session.startOnce.Do(func() {
+		session.mu.Lock()
+		if session.closed {
+			session.mu.Unlock()
+
+			return
+		}
+		session.started = true
+		session.mu.Unlock()
+
+		go session.run()
+	})
+}
+
+func (session *execSession) closeIfUnused() {
+	session.mu.Lock()
+	unused := !session.started && len(session.subscribers) == 0
+	session.mu.Unlock()
+
+	if unused {
+		session.close()
+	}
+}
+
 func (session *execSession) attach() (*execSessionSubscriber, error) {
 	session.mu.Lock()
 	defer session.mu.Unlock()
@@ -222,6 +266,12 @@ func (session *execSession) attach() (*execSessionSubscriber, error) {
 }
 
 func (session *execSession) detach(subscriber *execSessionSubscriber) {
+	if session.policy.closeOnDetach {
+		session.close()
+
+		return
+	}
+
 	session.mu.Lock()
 	defer session.mu.Unlock()
 
@@ -261,6 +311,10 @@ func (session *execSession) writeStdin(data []byte) error {
 }
 
 func (session *execSession) ack(watermark uint64) {
+	if !session.policy.replayEnabled {
+		return
+	}
+
 	session.mu.Lock()
 	defer session.mu.Unlock()
 
@@ -276,6 +330,10 @@ func (session *execSession) sendHistory(
 	subscriber *execSessionSubscriber,
 	watermark uint64,
 ) {
+	if !session.policy.replayEnabled {
+		return
+	}
+
 	session.mu.Lock()
 	defer session.mu.Unlock()
 
@@ -333,8 +391,10 @@ func (session *execSession) close() {
 	if session.transport != nil {
 		_ = session.transport.Close()
 	}
+	if session.registry != nil {
 		session.registry.remove(session.key, session)
 	}
+}
 
 func (session *execSession) run() {
 	outgoingFrames := make(chan *execstream.Frame)
@@ -368,8 +428,9 @@ func (session *execSession) recordFrame(frame *execstream.Frame) {
 		return
 	}
 
-	session.nextWatermark++
 	frame = cloneExecFrame(frame)
+	if session.policy.replayEnabled {
+		session.nextWatermark++
 		frame.Watermark = session.nextWatermark
 
 		session.frames = append(session.frames, execReplayFrame{
@@ -379,6 +440,7 @@ func (session *execSession) recordFrame(frame *execstream.Frame) {
 		session.bufferBytes += execFrameSize(frame)
 		session.trimAcknowledgedLocked()
 		session.trimToLimitLocked()
+	}
 
 	for subscriber := range session.subscribers {
 		if subscriber.enqueue(frame) {
@@ -398,7 +460,8 @@ func (session *execSession) markFinished() {
 	}
 
 	session.finished = true
-	if !session.closed {
+	shouldClose := !session.policy.retainAfterExit
+	if !session.closed && session.policy.retainAfterExit {
 		session.expiryTimer = time.AfterFunc(session.exitTTL, session.expire)
 	}
 
@@ -416,6 +479,10 @@ func (session *execSession) markFinished() {
 	session.doneOnce.Do(func() {
 		close(session.done)
 	})
+
+	if shouldClose {
+		session.close()
+	}
 }
 
 func (session *execSession) expire() {

internal/controller/exec_sessions_test.go

@@ -51,6 +51,7 @@ func newManualExecSessionForTest(
 		exec:        &fakeExec{},
 		registry:    registry,
 		exitTTL:     time.Minute,
+		policy:      reconnectableExecSessionPolicy,
 		ctx:         ctx,
 		cancel:      cancel,
 		subscribers: map[*execSessionSubscriber]struct{}{},
@@ -98,6 +99,36 @@ func TestExecSessionRegistryGetOrCreateReusesInflightCreation(t *testing.T) {
 	require.EqualValues(t, 1, createCalls.Load())
 }
 
+func TestExecSessionStartRunsCommandOnlyOnce(t *testing.T) {
+	var runCalls atomic.Int32
+	runStarted := make(chan struct{})
+
+	session := newExecSession(
+		execSessionKey{vmName: "vm", sessionID: "session"},
+		"echo test",
+		&fakeExec{
+			run: func(ctx context.Context, _ string, _ chan<- *execstream.Frame) error {
+				runCalls.Add(1)
+				close(runStarted)
+				<-ctx.Done()
+
+				return ctx.Err()
+			},
+		},
+		nil,
+		nil,
+		time.Minute,
+		reconnectableExecSessionPolicy,
+	)
+	defer session.close()
+
+	session.start()
+	session.start()
+
+	<-runStarted
+	require.EqualValues(t, 1, runCalls.Load())
+}
+
 func TestExecSessionHistoryReplayAndAck(t *testing.T) {
 	registry := newExecSessionRegistry()
 	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
@@ -140,6 +171,57 @@ func TestExecSessionDetachKeepsProcessAlive(t *testing.T) {
 	require.EqualValues(t, 0, exec.closeCalls.Load())
 }
 
+func TestLegacyExecSessionDetachStopsProcess(t *testing.T) {
+	registry := newExecSessionRegistry()
+	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
+	session.policy = legacyExecSessionPolicy
+	exec := session.exec.(*fakeExec)
+
+	subscriber, err := session.attach()
+	require.NoError(t, err)
+
+	session.detach(subscriber)
+
+	require.True(t, session.closed)
+	require.EqualValues(t, 1, exec.closeCalls.Load())
+}
+
+func TestLegacyExecSessionDoesNotRetainReplayHistory(t *testing.T) {
+	registry := newExecSessionRegistry()
+	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
+	session.policy = legacyExecSessionPolicy
+
+	session.recordFrame(&execstream.Frame{Type: execstream.FrameTypeStdout, Data: []byte("out")})
+
+	require.Empty(t, session.frames)
+	require.Zero(t, session.nextWatermark)
+}
+
+func TestExecSessionCloseIfUnusedClosesIdleSession(t *testing.T) {
+	registry := newExecSessionRegistry()
+	key := execSessionKey{vmName: "vm", sessionID: "session"}
+	session := newManualExecSessionForTest(key, registry)
+	exec := session.exec.(*fakeExec)
+	registry.sessions[key] = session
+
+	session.closeIfUnused()
+
+	require.True(t, session.closed)
+	require.EqualValues(t, 1, exec.closeCalls.Load())
+}
+
+func TestExecSessionCloseIfUnusedKeepsAttachedSession(t *testing.T) {
+	registry := newExecSessionRegistry()
+	session := newManualExecSessionForTest(execSessionKey{vmName: "vm", sessionID: "session"}, registry)
+
+	_, err := session.attach()
+	require.NoError(t, err)
+
+	session.closeIfUnused()
+
+	require.False(t, session.closed)
+}
+
 func TestExecSessionCloseStopsProcessAndRemovesRegistryEntry(t *testing.T) {
 	registry := newExecSessionRegistry()
 	key := execSessionKey{vmName: "vm", sessionID: "session"}