diff --git a/internal/grpc/client/client_test.go b/internal/grpc/client/client_test.go index 5009185..790aa46 100644 --- a/internal/grpc/client/client_test.go +++ b/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) } diff --git a/internal/registry/registry_test.go b/internal/registry/registry_test.go index 9a80d47..2bf5631 100644 --- a/internal/registry/registry_test.go +++ b/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) } diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index abf15dc..1ccfb5c 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/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 { diff --git a/internal/session/lifecycle/lifecycle_test.go b/internal/session/lifecycle/lifecycle_test.go index b8fc7bd..3f07d69 100644 --- a/internal/session/lifecycle/lifecycle_test.go +++ b/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) diff --git a/internal/session/session.go b/internal/session/session.go index 5978827..b31e829 100644 --- a/internal/session/session.go +++ b/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)