From efbff0f7f7e66fee3239ef34fbbd25a946ff9151 Mon Sep 17 00:00:00 2001 From: Bagas Date: Sat, 11 Jul 2026 19:20:22 +0700 Subject: [PATCH 1/7] fix: use mutex on lifecycle --- internal/session/lifecycle/lifecycle.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index 286d8fb..abf15dc 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/internal/session/lifecycle/lifecycle.go @@ -73,10 +73,14 @@ func (l *lifecycle) User() string { } func (l *lifecycle) SetChannel(channel ssh.Channel) { + l.mu.Lock() + defer l.mu.Unlock() l.channel = channel } func (l *lifecycle) Channel() ssh.Channel { + l.mu.Lock() + defer l.mu.Unlock() return l.channel } -- 2.52.0 From 0a5f673c68b1829ffee8ec95e07d70cff319795d Mon Sep 17 00:00:00 2001 From: Bagas Date: Sun, 12 Jul 2026 01:20:22 +0700 Subject: [PATCH 2/7] feat: seperate close function into tiny helper function --- internal/grpc/client/client_test.go | 2 +- internal/registry/registry_test.go | 2 +- internal/session/lifecycle/lifecycle.go | 79 +++++++++++++------- internal/session/lifecycle/lifecycle_test.go | 15 +++- internal/session/session.go | 4 +- 5 files changed, 67 insertions(+), 35 deletions(-) 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) -- 2.52.0 From 43d1ec0e6d3490bde3d0f3da8cc0ef64e06f89ca Mon Sep 17 00:00:00 2001 From: Bagas Date: Sun, 12 Jul 2026 08:33:21 +0700 Subject: [PATCH 3/7] fix: change start date of lifecycle after running status --- internal/session/lifecycle/lifecycle.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index 1ccfb5c..9a9ed04 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/internal/session/lifecycle/lifecycle.go @@ -46,7 +46,7 @@ func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUti channel: nil, forwarder: forwarder, slug: slugManager, - startedAt: time.Now(), + startedAt: time.Time{}, sessionRegistry: sessionRegistry, portRegistry: port, user: user, @@ -97,6 +97,9 @@ func (l *lifecycle) SetStatus(status types.SessionStatus) { l.mu.Lock() defer l.mu.Unlock() l.status = status + if status == types.SessionStatusRUNNING && l.startedAt.IsZero() { + l.startedAt = time.Now() + } } func (l *lifecycle) IsActive() bool { -- 2.52.0 From f1df3d26c20e9216a7abfc41f2fd34f115860d58 Mon Sep 17 00:00:00 2001 From: Bagas Date: Mon, 13 Jul 2026 02:42:21 +0700 Subject: [PATCH 4/7] feat: make PortRegistry interface segregation --- internal/grpc/client/client_test.go | 5 ++--- internal/registry/registry_test.go | 5 ++--- internal/session/lifecycle/lifecycle.go | 22 +++++++++++++++------- 3 files changed, 19 insertions(+), 13 deletions(-) diff --git a/internal/grpc/client/client_test.go b/internal/grpc/client/client_test.go index 790aa46..5ebd36d 100644 --- a/internal/grpc/client/client_test.go +++ b/internal/grpc/client/client_test.go @@ -13,7 +13,6 @@ import ( "tunnel_pls/internal/session/slug" "tunnel_pls/internal/types" - "tunnel_pls/internal/port" "tunnel_pls/internal/registry" proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen" @@ -889,12 +888,12 @@ func (m *mockLifecycle) SetChannel(channel ssh.Channel) error { return m.Called( 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) } -func (m *mockLifecycle) PortRegistry() port.Port { +func (m *mockLifecycle) PortRegistry() lifecycle.PortRegistry { args := m.Called() if args.Get(0) == nil { return nil } - return args.Get(0).(port.Port) + return args.Get(0).(lifecycle.PortRegistry) } type mockEventServiceClient struct { diff --git a/internal/registry/registry_test.go b/internal/registry/registry_test.go index 2bf5631..489122c 100644 --- a/internal/registry/registry_test.go +++ b/internal/registry/registry_test.go @@ -4,7 +4,6 @@ import ( "sync" "testing" "time" - "tunnel_pls/internal/port" "tunnel_pls/internal/session/forwarder" "tunnel_pls/internal/session/interaction" "tunnel_pls/internal/session/lifecycle" @@ -78,12 +77,12 @@ func (ml *mockLifecycle) Connection() ssh.Conn { return args.Get(0).(ssh.Conn) } -func (ml *mockLifecycle) PortRegistry() port.Port { +func (ml *mockLifecycle) PortRegistry() lifecycle.PortRegistry { args := ml.Called() if args.Get(0) == nil { return nil } - return args.Get(0).(port.Port) + return args.Get(0).(lifecycle.PortRegistry) } func (ml *mockLifecycle) SetChannel(channel ssh.Channel) error { return ml.Called(channel).Error(0) } diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index 9a9ed04..b291fee 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/internal/session/lifecycle/lifecycle.go @@ -10,8 +10,6 @@ import ( "tunnel_pls/internal/session/slug" "tunnel_pls/internal/types" - portUtil "tunnel_pls/internal/port" - "golang.org/x/crypto/ssh" ) @@ -25,6 +23,12 @@ type SessionRegistry interface { Remove(key types.SessionKey) } +type PortRegistry interface { + Unassigned() (uint16, bool) + Claim(port uint16) bool + SetStatus(port uint16, assigned bool) error +} + type lifecycle struct { mu sync.Mutex status types.SessionStatus @@ -35,11 +39,11 @@ type lifecycle struct { slug slug.Slug startedAt time.Time sessionRegistry SessionRegistry - portRegistry portUtil.Port + portRegistry PortRegistry user string } -func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUtil.Port, sessionRegistry SessionRegistry, user string) Lifecycle { +func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port PortRegistry, sessionRegistry SessionRegistry, user string) Lifecycle { return &lifecycle{ status: types.SessionStatusINITIALIZING, conn: conn, @@ -56,7 +60,7 @@ func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUti type Lifecycle interface { Connection() ssh.Conn Channel() ssh.Channel - PortRegistry() portUtil.Port + PortRegistry() PortRegistry User() string SetChannel(channel ssh.Channel) error SetStatus(status types.SessionStatus) @@ -65,7 +69,7 @@ type Lifecycle interface { Close() error } -func (l *lifecycle) PortRegistry() portUtil.Port { +func (l *lifecycle) PortRegistry() PortRegistry { return l.portRegistry } @@ -148,8 +152,12 @@ func (l *lifecycle) closeConnection() error { } func (l *lifecycle) cleanupRegistry() { + slugStr := l.slug.String() + if slugStr == "" { + return + } key := types.SessionKey{ - Id: l.slug.String(), + Id: slugStr, Type: l.forwarder.TunnelType(), } l.sessionRegistry.Remove(key) -- 2.52.0 From 6ec57268dbecd5d28a58008b61ce4de1101c470f Mon Sep 17 00:00:00 2001 From: Bagas Date: Tue, 14 Jul 2026 11:42:26 +0700 Subject: [PATCH 5/7] fix: read starttedAt with mutex on lifecycle --- internal/session/lifecycle/lifecycle.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index b291fee..c04473d 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/internal/session/lifecycle/lifecycle.go @@ -181,5 +181,7 @@ func isClosedError(err error) bool { } func (l *lifecycle) StartedAt() time.Time { + l.mu.Lock() + defer l.mu.Unlock() return l.startedAt } -- 2.52.0 From 6aa8c3bd486d557dff15e2667441a579d3c967a6 Mon Sep 17 00:00:00 2001 From: Bagas Date: Wed, 15 Jul 2026 12:23:33 +0700 Subject: [PATCH 6/7] fix: race condition test --- .../session/interaction/interaction_test.go | 4 -- internal/session/lifecycle/lifecycle_test.go | 47 +++++++++++++++++++ internal/transport/http_test.go | 3 +- internal/transport/https_test.go | 3 +- internal/transport/tcp_test.go | 3 +- 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/internal/session/interaction/interaction_test.go b/internal/session/interaction/interaction_test.go index 679ec8a..8b6f15b 100644 --- a/internal/session/interaction/interaction_test.go +++ b/internal/session/interaction/interaction_test.go @@ -1922,10 +1922,6 @@ func TestInteraction_Start_ProtocolSelection(t *testing.T) { time.Sleep(50 * time.Millisecond) i := mockInteraction.(*interaction) - if i.program != nil { - assert.NotNil(t, i.program, "program should be initialized") - } - i.Stop() mockConfig.AssertExpectations(t) diff --git a/internal/session/lifecycle/lifecycle_test.go b/internal/session/lifecycle/lifecycle_test.go index 3f07d69..e16fd48 100644 --- a/internal/session/lifecycle/lifecycle_test.go +++ b/internal/session/lifecycle/lifecycle_test.go @@ -5,6 +5,7 @@ import ( "errors" "io" "net" + "sync" "testing" "tunnel_pls/internal/types" @@ -308,3 +309,49 @@ func TestLifecycle_Close(t *testing.T) { }) } } + +func TestLifecycle_ConcurrentClose(t *testing.T) { + mockSSHConn := &MockSSHConn{} + mockSSHConn.On("Close").Return(nil) + + mockForwarder := &MockForwarder{} + mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP) + + mockSlug := &MockSlug{} + mockSlug.On("String").Return("test-slug") + + mockPort := &MockPort{} + + mockSessionRegistry := &MockSessionRegistry{} + mockSessionRegistry.On("Remove", mock.Anything).Return() + + mockSSHChannel := &MockSSHChannel{} + mockSSHChannel.On("Close").Return(nil) + + mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad") + mockLifecycle.SetStatus(types.SessionStatusRUNNING) + err := mockLifecycle.SetChannel(mockSSHChannel) + assert.NoError(t, err) + + const numGoroutines = 10 + var wg sync.WaitGroup + errChan := make(chan error, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func() { + defer wg.Done() + err := mockLifecycle.Close() + errChan <- err + }() + } + + wg.Wait() + close(errChan) + + for err := range errChan { + assert.NoError(t, err) + } + + assert.False(t, mockLifecycle.IsActive()) +} diff --git a/internal/transport/http_test.go b/internal/transport/http_test.go index cd3cf68..847931c 100644 --- a/internal/transport/http_test.go +++ b/internal/transport/http_test.go @@ -55,8 +55,7 @@ func TestHTTPServer_Serve(t *testing.T) { go func() { time.Sleep(100 * time.Millisecond) - err = listener.Close() - assert.NoError(t, err) + _ = listener.Close() }() err = srv.Serve(listener) diff --git a/internal/transport/https_test.go b/internal/transport/https_test.go index 6081d97..42bbe72 100644 --- a/internal/transport/https_test.go +++ b/internal/transport/https_test.go @@ -63,8 +63,7 @@ func TestHTTPSServer_Serve(t *testing.T) { go func() { time.Sleep(100 * time.Millisecond) - err = listener.Close() - assert.NoError(t, err) + _ = listener.Close() }() err = srv.Serve(listener) diff --git a/internal/transport/tcp_test.go b/internal/transport/tcp_test.go index c4c4963..761b902 100644 --- a/internal/transport/tcp_test.go +++ b/internal/transport/tcp_test.go @@ -45,8 +45,7 @@ func TestTCPServer_Serve(t *testing.T) { go func() { time.Sleep(100 * time.Millisecond) - err = listener.Close() - assert.NoError(t, err) + _ = listener.Close() }() err = srv.Serve(listener) -- 2.52.0 From cbab4f761248ecb75d3e5ea5d8e15bd402882972 Mon Sep 17 00:00:00 2001 From: Bagas Date: Thu, 16 Jul 2026 15:20:49 +0700 Subject: [PATCH 7/7] fix: reject closing already closed lifecycle --- internal/session/lifecycle/lifecycle.go | 61 ++++++++++-------- internal/session/lifecycle/lifecycle_test.go | 67 ++++++++++++++++++++ 2 files changed, 102 insertions(+), 26 deletions(-) diff --git a/internal/session/lifecycle/lifecycle.go b/internal/session/lifecycle/lifecycle.go index c04473d..0191e40 100644 --- a/internal/session/lifecycle/lifecycle.go +++ b/internal/session/lifecycle/lifecycle.go @@ -80,6 +80,12 @@ func (l *lifecycle) User() string { func (l *lifecycle) SetChannel(channel ssh.Channel) error { l.mu.Lock() defer l.mu.Unlock() + if l.status == types.SessionStatusCLOSED { + return fmt.Errorf("lifecycle is closed") + } + if channel == nil { + return fmt.Errorf("channel cannot be nil") + } if l.channel != nil { return fmt.Errorf("channel already set") } @@ -100,6 +106,9 @@ func (l *lifecycle) Connection() ssh.Conn { func (l *lifecycle) SetStatus(status types.SessionStatus) { l.mu.Lock() defer l.mu.Unlock() + if l.status == types.SessionStatusCLOSED { + return + } l.status = status if status == types.SessionStatusRUNNING && l.startedAt.IsZero() { l.startedAt = time.Now() @@ -114,41 +123,41 @@ 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 + closeErr := l.closeErr + l.mu.Unlock() + return closeErr } l.status = types.SessionStatusCLOSED + channel := l.channel + conn := l.conn + l.mu.Unlock() + var errs []error - errs = append(errs, l.closeChannel()) - errs = append(errs, l.closeConnection()) + if channel != nil { + if err := channel.Close(); err != nil && !isClosedError(err) { + errs = append(errs, err) + } + } + if conn != nil { + if err := conn.Close(); err != nil && !isClosedError(err) { + errs = append(errs, err) + } + } + l.cleanupRegistry() - errs = append(errs, l.cleanupForwarder()) + if err := l.cleanupForwarder(); err != nil { + errs = append(errs, err) + } - l.closeErr = errors.Join(errs...) - return l.closeErr -} + closeErr := errors.Join(errs...) -func (l *lifecycle) closeChannel() error { - if l.channel == nil { - return nil - } - if err := l.channel.Close(); err != nil && !isClosedError(err) { - return err - } - return nil -} + l.mu.Lock() + l.closeErr = closeErr + l.mu.Unlock() -func (l *lifecycle) closeConnection() error { - if l.conn == nil { - return nil - } - if err := l.conn.Close(); err != nil && !isClosedError(err) { - return err - } - return nil + return closeErr } func (l *lifecycle) cleanupRegistry() { diff --git a/internal/session/lifecycle/lifecycle_test.go b/internal/session/lifecycle/lifecycle_test.go index e16fd48..b5b2ff6 100644 --- a/internal/session/lifecycle/lifecycle_test.go +++ b/internal/session/lifecycle/lifecycle_test.go @@ -355,3 +355,70 @@ func TestLifecycle_ConcurrentClose(t *testing.T) { assert.False(t, mockLifecycle.IsActive()) } + +func TestLifecycle_SetChannel_AfterClose(t *testing.T) { + mockSSHConn := new(MockSSHConn) + mockSSHConn.On("Close").Return(nil) + mockForwarder := &MockForwarder{} + mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP) + mockSlug := &MockSlug{} + mockSlug.On("String").Return("test-slug") + mockPort := &MockPort{} + mockSessionRegistry := &MockSessionRegistry{} + mockSessionRegistry.On("Remove", mock.Anything).Return() + mockSSHChannel := &MockSSHChannel{} + mockSSHChannel.On("Close").Return(nil) + + mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad") + mockLifecycle.SetStatus(types.SessionStatusRUNNING) + err := mockLifecycle.SetChannel(mockSSHChannel) + assert.NoError(t, err) + + err = mockLifecycle.Close() + assert.NoError(t, err) + + anotherChannel := &MockSSHChannel{} + err = mockLifecycle.SetChannel(anotherChannel) + assert.Error(t, err) + assert.Contains(t, err.Error(), "lifecycle is closed") +} + +func TestLifecycle_SetChannel_Nil(t *testing.T) { + mockSSHConn := new(MockSSHConn) + mockForwarder := &MockForwarder{} + mockSlug := &MockSlug{} + mockPort := &MockPort{} + mockSessionRegistry := &MockSessionRegistry{} + + mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad") + + err := mockLifecycle.SetChannel(nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "channel cannot be nil") +} + +func TestLifecycle_SetStatus_AfterClose(t *testing.T) { + mockSSHConn := new(MockSSHConn) + mockSSHConn.On("Close").Return(nil) + mockForwarder := &MockForwarder{} + mockForwarder.On("TunnelType").Return(types.TunnelTypeHTTP) + mockSlug := &MockSlug{} + mockSlug.On("String").Return("test-slug") + mockPort := &MockPort{} + mockSessionRegistry := &MockSessionRegistry{} + mockSessionRegistry.On("Remove", mock.Anything).Return() + mockSSHChannel := &MockSSHChannel{} + mockSSHChannel.On("Close").Return(nil) + + mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad") + mockLifecycle.SetStatus(types.SessionStatusRUNNING) + err := mockLifecycle.SetChannel(mockSSHChannel) + assert.NoError(t, err) + + err = mockLifecycle.Close() + assert.NoError(t, err) + assert.False(t, mockLifecycle.IsActive()) + + mockLifecycle.SetStatus(types.SessionStatusRUNNING) + assert.False(t, mockLifecycle.IsActive(), "SetStatus should be ignored after Close") +} -- 2.52.0