From e8be839dacbb5abc31619781c2c349321cc909ee Mon Sep 17 00:00:00 2001 From: Bagas Date: Thu, 16 Jul 2026 17:31:14 +0700 Subject: [PATCH] Fix race conditions and improve lifecycle safety (#150) ## Summary Comprehensive fixes for race conditions, concurrency bugs, and code quality improvements in the lifecycle module and related test files. ## Changes ### Race Condition Fixes - Fixed data race in `SetChannel()` and `Channel()` by adding mutex protection - Fixed data race in `StartedAt()` by adding mutex protection - Fixed test race conditions in `transport` and `interaction` packages where variables were shared between goroutines ### Concurrency Safety Improvements - Refactored `Close()` to release mutex before calling external code, preventing potential deadlocks - Made `SetChannel()` reject calls after `Close()` to prevent setting channel on torn-down lifecycle - Made `SetStatus()` ignore calls after `Close()` to ensure CLOSED is a terminal state - Added nil check to `SetChannel()` to prevent masking upstream bugs ### Code Quality Improvements - Refactored `Close()` to extract helper methods: `closeChannel()`, `closeConnection()`, `cleanupRegistry()`, `cleanupForwarder()` - Fixed `isClosedError()` to remove redundant string comparison - Added `PortRegistry` interface for better interface segregation (removed dependency on full `port.Port`) - Made `cleanupRegistry()` conditional on non-empty slug to avoid unnecessary registry calls - Fixed `startedAt` to be set when session starts running instead of at object creation Reviewed-on: https://git.fossy.my.id/bagas/tunnel-please/pulls/150 Co-authored-by: Bagas Co-committed-by: Bagas --- internal/grpc/client/client_test.go | 7 +- internal/registry/registry_test.go | 7 +- .../session/interaction/interaction_test.go | 4 - internal/session/lifecycle/lifecycle.go | 107 +++++++++++---- internal/session/lifecycle/lifecycle_test.go | 129 +++++++++++++++++- internal/session/session.go | 4 +- internal/transport/http_test.go | 3 +- internal/transport/https_test.go | 3 +- internal/transport/tcp_test.go | 3 +- 9 files changed, 215 insertions(+), 52 deletions(-) diff --git a/internal/grpc/client/client_test.go b/internal/grpc/client/client_test.go index 5009185..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" @@ -885,16 +884,16 @@ 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) } -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 9a80d47..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,15 +77,15 @@ 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) { 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/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.go b/internal/session/lifecycle/lifecycle.go index 286d8fb..0191e40 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" @@ -9,8 +10,6 @@ import ( "tunnel_pls/internal/session/slug" "tunnel_pls/internal/types" - portUtil "tunnel_pls/internal/port" - "golang.org/x/crypto/ssh" ) @@ -24,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 @@ -34,18 +39,18 @@ 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, channel: nil, forwarder: forwarder, slug: slugManager, - startedAt: time.Now(), + startedAt: time.Time{}, sessionRegistry: sessionRegistry, portRegistry: port, user: user, @@ -55,16 +60,16 @@ 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) + SetChannel(channel ssh.Channel) error SetStatus(status types.SessionStatus) IsActive() bool StartedAt() time.Time Close() error } -func (l *lifecycle) PortRegistry() portUtil.Port { +func (l *lifecycle) PortRegistry() PortRegistry { return l.portRegistry } @@ -72,11 +77,25 @@ 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.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") + } l.channel = channel + return nil } func (l *lifecycle) Channel() ssh.Channel { + l.mu.Lock() + defer l.mu.Unlock() return l.channel } @@ -87,7 +106,13 @@ 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() + } } func (l *lifecycle) IsActive() bool { @@ -98,50 +123,74 @@ 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 - tunnelType := l.forwarder.TunnelType() - - if l.channel != nil { - if err := l.channel.Close(); err != nil && !isClosedError(err) { + 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) } } - if l.conn != nil { - if err := l.conn.Close(); err != nil && !isClosedError(err) { - errs = append(errs, err) - } + l.cleanupRegistry() + if err := l.cleanupForwarder(); err != nil { + errs = append(errs, err) } - clientSlug := l.slug.String() + closeErr := errors.Join(errs...) + + l.mu.Lock() + l.closeErr = closeErr + l.mu.Unlock() + + return closeErr +} + +func (l *lifecycle) cleanupRegistry() { + slugStr := l.slug.String() + if slugStr == "" { + return + } key := types.SessionKey{ - Id: clientSlug, - Type: tunnelType, + Id: slugStr, + Type: l.forwarder.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()) +func (l *lifecycle) cleanupForwarder() error { + if l.forwarder.TunnelType() != types.TunnelTypeTCP { + return nil } - - l.closeErr = errors.Join(errs...) - return l.closeErr + 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 { + l.mu.Lock() + defer l.mu.Unlock() return l.startedAt } diff --git a/internal/session/lifecycle/lifecycle_test.go b/internal/session/lifecycle/lifecycle_test.go index b8fc7bd..b5b2ff6 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" @@ -177,8 +178,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 +283,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) @@ -301,3 +309,116 @@ 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()) +} + +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") +} 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) 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)