diff --git a/internal/controller/api_vms_exec.go b/internal/controller/api_vms_exec.go index 5e9a868..dafb25c 100644 --- a/internal/controller/api_vms_exec.go +++ b/internal/controller/api_vms_exec.go @@ -175,7 +175,7 @@ func (controller *Controller) execVMReconnectable( } func (controller *Controller) newSSHExecSession( - ctx *gin.Context, + _ *gin.Context, waitContext context.Context, vm *v1.VM, key execSessionKey, @@ -184,6 +184,8 @@ func (controller *Controller) newSSHExecSession( registry *execSessionRegistry, policy execSessionPolicy, ) (*execSession, error) { + sessionContext, sessionContextCancel := context.WithCancel(context.Background()) + portForwardConn, err := retry.NewWithData[net.Conn]( retry.Context(waitContext), retry.DelayType(retry.FixedDelay), @@ -191,20 +193,25 @@ func (controller *Controller) newSSHExecSession( retry.Attempts(0), retry.LastErrorOnly(true), ).Do(func() (net.Conn, error) { - return controller.portForwardConnection(ctx, waitContext, vm.Worker, vm.UID, 22) + return controller.portForwardConnection(sessionContext, waitContext, vm.Worker, vm.UID, 22) }) if err != nil { + sessionContextCancel() + return nil, err } exec, err := sshexec.New(portForwardConn, vm.SSHUsername(), vm.SSHPassword(), stdin) if err != nil { + sessionContextCancel() _ = portForwardConn.Close() return nil, fmt.Errorf("failed to establish SSH connection to a VM: %w", err) } - return newExecSession( + return newExecSessionWithContext( + sessionContext, + sessionContextCancel, key, command, exec, diff --git a/internal/controller/exec_sessions.go b/internal/controller/exec_sessions.go index 6ca6855..cb186ac 100644 --- a/internal/controller/exec_sessions.go +++ b/internal/controller/exec_sessions.go @@ -334,6 +334,34 @@ func newExecSession( ) *execSession { ctx, cancel := context.WithCancel(context.Background()) + return newExecSessionWithContext( + ctx, + cancel, + key, + command, + exec, + transport, + registry, + exitTTL, + policy, + ) +} + +func newExecSessionWithContext( + ctx context.Context, + cancel context.CancelFunc, + key execSessionKey, + command string, + exec sshExecRunner, + transport net.Conn, + registry *execSessionRegistry, + exitTTL time.Duration, + policy execSessionPolicy, +) *execSession { + if ctx == nil || cancel == nil { + ctx, cancel = context.WithCancel(context.Background()) + } + session := &execSession{ key: key, command: command,