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: #150
Co-authored-by: Bagas <bagas@fossy.my.id>
Co-committed-by: Bagas <bagas@fossy.my.id>

.gitea/workflows/build.yml

@@ -79,7 +79,7 @@ jobs:
           fi
 
       - name: Build and push Docker image (release)
-        uses: docker/build-push-action@v7
+        uses: docker/build-push-action@v6
         with:
           context: .
           push: true
@@ -97,7 +97,7 @@ jobs:
         if: steps.version.outputs.IS_PRERELEASE == 'false'
 
       - name: Build and push Docker image (pre-release)
-        uses: docker/build-push-action@v7
+        uses: docker/build-push-action@v6
         with:
           context: .
           push: true

go.mod

@@ -4,7 +4,7 @@ go 1.26.0
 
 require (
 	git.fossy.my.id/bagas/tunnel-please-grpc v1.5.0
-	github.com/caddyserver/certmagic v0.25.2
+	github.com/caddyserver/certmagic v0.25.1
 	github.com/charmbracelet/bubbles v1.0.0
 	github.com/charmbracelet/bubbletea v1.3.10
 	github.com/charmbracelet/lipgloss v1.1.0
@@ -13,14 +13,14 @@ require (
 	github.com/muesli/termenv v0.16.0
 	github.com/stretchr/testify v1.11.1
 	golang.org/x/crypto v0.48.0
-	google.golang.org/grpc v1.79.1
+	google.golang.org/grpc v1.78.0
 	google.golang.org/protobuf v1.36.11
 )
 
 require (
 	github.com/atotto/clipboard v0.1.4 // indirect
 	github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
-	github.com/caddyserver/zerossl v0.1.5 // indirect
+	github.com/caddyserver/zerossl v0.1.4 // indirect
 	github.com/charmbracelet/colorprofile v0.4.1 // indirect
 	github.com/charmbracelet/x/ansi v0.11.6 // indirect
 	github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
@@ -36,8 +36,8 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/mattn/go-localereader v0.0.1 // indirect
 	github.com/mattn/go-runewidth v0.0.19 // indirect
-	github.com/mholt/acmez/v3 v3.1.6 // indirect
-	github.com/miekg/dns v1.1.72 // indirect
+	github.com/mholt/acmez/v3 v3.1.4 // indirect
+	github.com/miekg/dns v1.1.69 // indirect
 	github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
@@ -49,12 +49,12 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.1 // indirect
 	go.uber.org/zap/exp v0.3.0 // indirect
-	golang.org/x/mod v0.33.0 // indirect
-	golang.org/x/net v0.50.0 // indirect
+	golang.org/x/mod v0.32.0 // indirect
+	golang.org/x/net v0.49.0 // indirect
 	golang.org/x/sync v0.19.0 // indirect
 	golang.org/x/sys v0.41.0 // indirect
 	golang.org/x/text v0.34.0 // indirect
-	golang.org/x/tools v0.42.0 // indirect
+	golang.org/x/tools v0.41.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -8,12 +8,8 @@ github.com/aymanbagabas/go-udiff v0.3.1 h1:LV+qyBQ2pqe0u42ZsUEtPiCaUoqgA9gYRDs3v
 github.com/aymanbagabas/go-udiff v0.3.1/go.mod h1:G0fsKmG+P6ylD0r6N/KgQD/nWzgfnl8ZBcNLgcbrw8E=
 github.com/caddyserver/certmagic v0.25.1 h1:4sIKKbOt5pg6+sL7tEwymE1x2bj6CHr80da1CRRIPbY=
 github.com/caddyserver/certmagic v0.25.1/go.mod h1:VhyvndxtVton/Fo/wKhRoC46Rbw1fmjvQ3GjHYSQTEY=
-github.com/caddyserver/certmagic v0.25.2 h1:D7xcS7ggX/WEY54x0czj7ioTkmDWKIgxtIi2OcQclUc=
-github.com/caddyserver/certmagic v0.25.2/go.mod h1:llW/CvsNmza8S6hmsuggsZeiX+uS27dkqY27wDIuBWg=
 github.com/caddyserver/zerossl v0.1.4 h1:CVJOE3MZeFisCERZjkxIcsqIH4fnFdlYWnPYeFtBHRw=
 github.com/caddyserver/zerossl v0.1.4/go.mod h1:CxA0acn7oEGO6//4rtrRjYgEoa4MFw/XofZnrYwGqG4=
-github.com/caddyserver/zerossl v0.1.5 h1:dkvOjBAEEtY6LIGAHei7sw2UgqSD6TrWweXpV7lvEvE=
-github.com/caddyserver/zerossl v0.1.5/go.mod h1:CxA0acn7oEGO6//4rtrRjYgEoa4MFw/XofZnrYwGqG4=
 github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc=
 github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E=
 github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw=
@@ -70,12 +66,8 @@ github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byF
 github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs=
 github.com/mholt/acmez/v3 v3.1.4 h1:DyzZe/RnAzT3rpZj/2Ii5xZpiEvvYk3cQEN/RmqxwFQ=
 github.com/mholt/acmez/v3 v3.1.4/go.mod h1:L1wOU06KKvq7tswuMDwKdcHeKpFFgkppZy/y0DFxagQ=
-github.com/mholt/acmez/v3 v3.1.6 h1:eGVQNObP0pBN4sxqrXeg7MYqTOWyoiYpQqITVWlrevk=
-github.com/mholt/acmez/v3 v3.1.6/go.mod h1:5nTPosTGosLxF3+LU4ygbgMRFDhbAVpqMI4+a4aHLBY=
 github.com/miekg/dns v1.1.69 h1:Kb7Y/1Jo+SG+a2GtfoFUfDkG//csdRPwRLkCsxDG9Sc=
 github.com/miekg/dns v1.1.69/go.mod h1:7OyjD9nEba5OkqQ/hB4fy3PIoxafSZJtducccIelz3g=
-github.com/miekg/dns v1.1.72 h1:vhmr+TF2A3tuoGNkLDFK9zi36F2LS+hKTRW0Uf8kbzI=
-github.com/miekg/dns v1.1.72/go.mod h1:+EuEPhdHOsfk6Wk5TT2CzssZdqkmFhf8r+aVyDEToIs=
 github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
 github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
 github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
@@ -126,12 +118,8 @@ golang.org/x/exp v0.0.0-20231006140011-7918f672742d h1:jtJma62tbqLibJ5sFQz8bKtEM
 golang.org/x/exp v0.0.0-20231006140011-7918f672742d/go.mod h1:ldy0pHrwJyGW56pPQzzkH36rKxoZW1tw7ZJpeKx+hdo=
 golang.org/x/mod v0.32.0 h1:9F4d3PHLljb6x//jOyokMv3eX+YDeepZSEo3mFJy93c=
 golang.org/x/mod v0.32.0/go.mod h1:SgipZ/3h2Ci89DlEtEXWUk/HteuRin+HHhN+WbNhguU=
-golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8=
-golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w=
 golang.org/x/net v0.49.0 h1:eeHFmOGUTtaaPSGNmjBKpbng9MulQsJURQUAfUwY++o=
 golang.org/x/net v0.49.0/go.mod h1:/ysNB2EvaqvesRkuLAyjI1ycPZlQHM3q01F02UY/MV8=
-golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60=
-golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM=
 golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
 golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
 golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
@@ -144,18 +132,12 @@ golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk=
 golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA=
 golang.org/x/tools v0.41.0 h1:a9b8iMweWG+S0OBnlU36rzLp20z1Rp10w+IY2czHTQc=
 golang.org/x/tools v0.41.0/go.mod h1:XSY6eDqxVNiYgezAVqqCeihT4j1U2CCsqvH3WhQpnlg=
-golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k=
-golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0=
 gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk=
 gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b h1:Mv8VFug0MP9e5vUxfBcE3vUkV6CImK3cMNMIDFjmzxU=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b/go.mod h1:j9x/tPzZkyxcgEFkiKEEGxfvyumM01BEtsW8xzOahRQ=
 google.golang.org/grpc v1.78.0 h1:K1XZG/yGDJnzMdd/uZHAkVqJE+xIDOcmdSFZkBUicNc=
 google.golang.org/grpc v1.78.0/go.mod h1:I47qjTo4OKbMkjA/aOOwxDIiPSBofUtQUI5EfpWvW7U=
-google.golang.org/grpc v1.79.0 h1:6/+EFlxsMyoSbHbBoEDx94n/Ycx/bi0IhJ5Qh7b7LaA=
-google.golang.org/grpc v1.79.0/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
-google.golang.org/grpc v1.79.1 h1:zGhSi45ODB9/p3VAawt9a+O/MULLl9dpizzNNpq7flY=
-google.golang.org/grpc v1.79.1/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ=
 google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
 google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=

internal/bootstrap/bootstrap.go

@@ -15,10 +15,10 @@ import (
 	"tunnel_pls/internal/port"
 	"tunnel_pls/internal/random"
 	"tunnel_pls/internal/registry"
+	"tunnel_pls/internal/server"
 	"tunnel_pls/internal/transport"
+	"tunnel_pls/internal/types"
 	"tunnel_pls/internal/version"
-	"tunnel_pls/server"
-	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )

internal/bootstrap/bootstrap_test.go

@@ -14,8 +14,8 @@ import (
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/port"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"

internal/config/config.go

@@ -1,6 +1,8 @@
 package config
 
-import "tunnel_pls/types"
+import (
+	"tunnel_pls/internal/types"
+)
 
 type Config interface {
 	Domain() string

internal/config/config_test.go

@@ -3,7 +3,7 @@ package config
 import (
 	"os"
 	"testing"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 )

internal/config/loader.go

@@ -6,7 +6,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/joho/godotenv"
 )
@@ -32,7 +32,7 @@ type config struct {
 
 	bufferSize int
 	headerSize int
-	
+
 	pprofEnabled bool
 	pprofPort    string
 

internal/grpc/client/client.go

@@ -9,7 +9,7 @@ import (
 	"time"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
 	"google.golang.org/grpc"

internal/grpc/client/client_test.go

@@ -7,14 +7,13 @@ import (
 	"io"
 	"testing"
 	"time"
+	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
-	"tunnel_pls/internal/port"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session/forwarder"
-	"tunnel_pls/session/interaction"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
 
 	proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
 	"github.com/stretchr/testify/assert"
@@ -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 {

internal/registry/registry.go

@@ -3,11 +3,11 @@ package registry
 import (
 	"fmt"
 	"sync"
-	"tunnel_pls/session/forwarder"
-	"tunnel_pls/session/interaction"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 )
 
 type Key = types.SessionKey

internal/registry/registry_test.go

@@ -4,12 +4,11 @@ import (
 	"sync"
 	"testing"
 	"time"
-	"tunnel_pls/internal/port"
-	"tunnel_pls/session/forwarder"
-	"tunnel_pls/session/interaction"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -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) }

internal/server/server.go

@@ -13,7 +13,7 @@ import (
 	"tunnel_pls/internal/port"
 	"tunnel_pls/internal/random"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session"
+	"tunnel_pls/internal/session"
 
 	"golang.org/x/crypto/ssh"
 )

internal/server/server_test.go

@@ -10,8 +10,8 @@ import (
 	"testing"
 	"time"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"

internal/session/forwarder/forwarder.go

@@ -10,8 +10,8 @@ import (
 	"strconv"
 	"sync"
 	"tunnel_pls/internal/config"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"golang.org/x/crypto/ssh"
 )

internal/session/forwarder/forwarder_test.go

@@ -10,8 +10,8 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"

internal/session/interaction/interaction.go

@@ -6,8 +6,8 @@ import (
 	"sync"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/random"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"github.com/charmbracelet/bubbles/help"
 	"github.com/charmbracelet/bubbles/key"

internal/session/interaction/interaction_test.go

@@ -7,7 +7,7 @@ import (
 	"net"
 	"testing"
 	"time"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/list"
@@ -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)

internal/session/interaction/model.go

@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"time"
 	"tunnel_pls/internal/random"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/charmbracelet/bubbles/help"
 	"github.com/charmbracelet/bubbles/key"

internal/session/interaction/slug.go

@@ -3,7 +3,7 @@ package interaction
 import (
 	"fmt"
 	"strings"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/textinput"

internal/session/lifecycle/lifecycle.go

@@ -2,14 +2,13 @@ package lifecycle
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"net"
 	"sync"
 	"time"
-
-	portUtil "tunnel_pls/internal/port"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"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
 }

internal/session/lifecycle/lifecycle_test.go

@@ -5,8 +5,9 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 	"testing"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -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")
+}

internal/session/session.go

@@ -12,12 +12,12 @@ import (
 	portUtil "tunnel_pls/internal/port"
 	"tunnel_pls/internal/random"
 	"tunnel_pls/internal/registry"
+	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/session/slug"
 	"tunnel_pls/internal/transport"
-	"tunnel_pls/session/forwarder"
-	"tunnel_pls/session/interaction"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -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)
 

internal/session/session_test.go

@@ -15,8 +15,8 @@ import (
 	"time"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"

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)

internal/transport/httphandler.go

@@ -16,7 +16,7 @@ import (
 	"tunnel_pls/internal/http/stream"
 	"tunnel_pls/internal/middleware"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"golang.org/x/crypto/ssh"
 )

internal/transport/httphandler_test.go

@@ -11,11 +11,11 @@ import (
 	"testing"
 	"time"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/session/forwarder"
-	"tunnel_pls/session/interaction"
-	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/internal/types"
 
 	"golang.org/x/crypto/ssh"
 

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)

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)

internal/transport/tls_test.go

@@ -14,7 +14,7 @@ import (
 	"testing"
 	"time"
 	"tunnel_pls/internal/config"
-	"tunnel_pls/types"
+	"tunnel_pls/internal/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"