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 {

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) }

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)