internal/grpc/client/client_test.go

@@ -885,7 +885,7 @@ func (m *mockLifecycle) Connection() ssh.Conn {
 	return args.Get(0).(ssh.Conn)
 }
 func (m *mockLifecycle) User() string                         { return m.Called().String(0) }
-func (m *mockLifecycle) SetChannel(channel ssh.Channel)       { m.Called(channel) }
+func (m *mockLifecycle) SetChannel(channel ssh.Channel) error { return m.Called(channel).Error(0) }
 func (m *mockLifecycle) SetStatus(status types.SessionStatus) { m.Called(status) }
 func (m *mockLifecycle) IsActive() bool                       { return m.Called().Bool(0) }
 func (m *mockLifecycle) StartedAt() time.Time                 { return m.Called().Get(0).(time.Time) }

internal/registry/registry_test.go

@@ -86,7 +86,7 @@ func (ml *mockLifecycle) PortRegistry() port.Port {
 	return args.Get(0).(port.Port)
 }
 
-func (ml *mockLifecycle) SetChannel(channel ssh.Channel)       { ml.Called(channel) }
+func (ml *mockLifecycle) SetChannel(channel ssh.Channel) error { return ml.Called(channel).Error(0) }
 func (ml *mockLifecycle) SetStatus(status types.SessionStatus) { ml.Called(status) }
 func (ml *mockLifecycle) IsActive() bool                       { return ml.Called().Bool(0) }
 func (ml *mockLifecycle) StartedAt() time.Time                 { return ml.Called().Get(0).(time.Time) }

internal/session/lifecycle/lifecycle.go

@@ -2,6 +2,7 @@ package lifecycle
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"sync"
@@ -57,7 +58,7 @@ type Lifecycle interface {
 	Channel() ssh.Channel
 	PortRegistry() portUtil.Port
 	User() string
-	SetChannel(channel ssh.Channel)
+	SetChannel(channel ssh.Channel) error
 	SetStatus(status types.SessionStatus)
 	IsActive() bool
 	StartedAt() time.Time
@@ -72,10 +73,14 @@ func (l *lifecycle) User() string {
 	return l.user
 }
 
-func (l *lifecycle) SetChannel(channel ssh.Channel) {
+func (l *lifecycle) SetChannel(channel ssh.Channel) error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
+	if l.channel != nil {
+		return fmt.Errorf("channel already set")
+	}
 	l.channel = channel
+	return nil
 }
 
 func (l *lifecycle) Channel() ssh.Channel {
@@ -103,47 +108,65 @@ func (l *lifecycle) IsActive() bool {
 func (l *lifecycle) Close() error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
+
 	if l.status == types.SessionStatusCLOSED {
 		return l.closeErr
 	}
 	l.status = types.SessionStatusCLOSED
 
 	var errs []error
-	tunnelType := l.forwarder.TunnelType()
-
-	if l.channel != nil {
-		if err := l.channel.Close(); err != nil && !isClosedError(err) {
-			errs = append(errs, err)
-		}
-	}
-
-	if l.conn != nil {
-		if err := l.conn.Close(); err != nil && !isClosedError(err) {
-			errs = append(errs, err)
-		}
-	}
-
-	clientSlug := l.slug.String()
-	key := types.SessionKey{
-		Id:   clientSlug,
-		Type: tunnelType,
-	}
-	l.sessionRegistry.Remove(key)
-
-	if tunnelType == types.TunnelTypeTCP {
-		errs = append(errs, l.PortRegistry().SetStatus(l.forwarder.ForwardedPort(), false))
-		errs = append(errs, l.forwarder.Close())
-	}
+	errs = append(errs, l.closeChannel())
+	errs = append(errs, l.closeConnection())
+	l.cleanupRegistry()
+	errs = append(errs, l.cleanupForwarder())
 
 	l.closeErr = errors.Join(errs...)
 	return l.closeErr
 }
 
+func (l *lifecycle) closeChannel() error {
+	if l.channel == nil {
+		return nil
+	}
+	if err := l.channel.Close(); err != nil && !isClosedError(err) {
+		return err
+	}
+	return nil
+}
+
+func (l *lifecycle) closeConnection() error {
+	if l.conn == nil {
+		return nil
+	}
+	if err := l.conn.Close(); err != nil && !isClosedError(err) {
+		return err
+	}
+	return nil
+}
+
+func (l *lifecycle) cleanupRegistry() {
+	key := types.SessionKey{
+		Id:   l.slug.String(),
+		Type: l.forwarder.TunnelType(),
+	}
+	l.sessionRegistry.Remove(key)
+}
+
+func (l *lifecycle) cleanupForwarder() error {
+	if l.forwarder.TunnelType() != types.TunnelTypeTCP {
+		return nil
+	}
+	var errs []error
+	errs = append(errs, l.portRegistry.SetStatus(l.forwarder.ForwardedPort(), false))
+	errs = append(errs, l.forwarder.Close())
+	return errors.Join(errs...)
+}
+
 func isClosedError(err error) bool {
 	if err == nil {
 		return false
 	}
-	return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || err.Error() == "EOF"
+	return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed)
 }
 
 func (l *lifecycle) StartedAt() time.Time {

internal/session/lifecycle/lifecycle_test.go

@@ -177,8 +177,14 @@ func TestLifecycle_SetChannel(t *testing.T) {
 
 	mockSSHChannel := &MockSSHChannel{}
 
-	mockLifecycle.SetChannel(mockSSHChannel)
+	err := mockLifecycle.SetChannel(mockSSHChannel)
+	assert.NoError(t, err)
+	assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
 
+	anotherChannel := &MockSSHChannel{}
+	err = mockLifecycle.SetChannel(anotherChannel)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "channel already set")
 	assert.Equal(t, mockSSHChannel, mockLifecycle.Channel())
 }
 
@@ -276,14 +282,15 @@ func TestLifecycle_Close(t *testing.T) {
 			mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
 
 			mockLifecycle.SetStatus(types.SessionStatusRUNNING)
-			mockLifecycle.SetChannel(mockSSHChannel)
+			err := mockLifecycle.SetChannel(mockSSHChannel)
+			assert.NoError(t, err)
 
 			if tt.alreadyClosed {
-				err := mockLifecycle.Close()
+				err = mockLifecycle.Close()
 				assert.NoError(t, err)
 			}
 
-			err := mockLifecycle.Close()
+			err = mockLifecycle.Close()
 
 			if tt.expectErr {
 				assert.Error(t, err)

internal/session/session.go

@@ -164,7 +164,9 @@ func (s *session) setupInteractiveMode(channel ssh.NewChannel) error {
 		}
 	}()
 
-	s.lifecycle.SetChannel(ch)
+	if err = s.lifecycle.SetChannel(ch); err != nil {
+		return err
+	}
 	s.interaction.SetChannel(ch)
 	s.interaction.SetMode(types.InteractiveModeINTERACTIVE)