refactor: inject SessionRegistry interface instead of individual functions

session/lifecycle/lifecycle.go

@@ -19,30 +19,36 @@ type Forwarder interface {
 	GetForwardedPort() uint16
 }
 
-type Lifecycle struct {
-	status           types.Status
-	conn             ssh.Conn
-	channel          ssh.Channel
-	forwarder        Forwarder
-	slugManager      slug.Manager
-	unregisterClient func(key types.SessionKey)
-	startedAt        time.Time
+type SessionRegistry interface {
+	Remove(key types.SessionKey)
 }
 
-func NewLifecycle(conn ssh.Conn, forwarder Forwarder, slugManager slug.Manager) *Lifecycle {
+type Lifecycle struct {
+	status          types.Status
+	conn            ssh.Conn
+	channel         ssh.Channel
+	forwarder       Forwarder
+	sessionRegistry SessionRegistry
+	slugManager     slug.Manager
+	startedAt       time.Time
+	user            string
+}
+
+func NewLifecycle(conn ssh.Conn, forwarder Forwarder, slugManager slug.Manager, user string) *Lifecycle {
 	return &Lifecycle{
-		status:           types.INITIALIZING,
-		conn:             conn,
-		channel:          nil,
-		forwarder:        forwarder,
-		slugManager:      slugManager,
-		unregisterClient: nil,
-		startedAt:        time.Now(),
+		status:          types.INITIALIZING,
+		conn:            conn,
+		channel:         nil,
+		forwarder:       forwarder,
+		slugManager:     slugManager,
+		sessionRegistry: nil,
+		startedAt:       time.Now(),
+		user:            user,
 	}
 }
 
-func (l *Lifecycle) SetUnregisterClient(unregisterClient func(key types.SessionKey)) {
-	l.unregisterClient = unregisterClient
+func (l *Lifecycle) SetSessionRegistry(registry SessionRegistry) {
+	l.sessionRegistry = registry
 }
 
 type SessionLifecycle interface {
@@ -50,12 +56,17 @@ type SessionLifecycle interface {
 	SetStatus(status types.Status)
 	GetConnection() ssh.Conn
 	GetChannel() ssh.Channel
+	GetUser() string
 	SetChannel(channel ssh.Channel)
-	SetUnregisterClient(unregisterClient func(key types.SessionKey))
+	SetSessionRegistry(registry SessionRegistry)
 	IsActive() bool
 	StartedAt() time.Time
 }
 
+func (l *Lifecycle) GetUser() string {
+	return l.user
+}
+
 func (l *Lifecycle) GetChannel() ssh.Channel {
 	return l.channel
 }
@@ -94,13 +105,13 @@ func (l *Lifecycle) Close() error {
 	}
 
 	clientSlug := l.slugManager.Get()
-	if clientSlug != "" && l.unregisterClient != nil {
+	if clientSlug != "" && l.sessionRegistry.Remove != nil {
 		key := types.SessionKey{Id: clientSlug, Type: l.forwarder.GetTunnelType()}
-		l.unregisterClient(key)
+		l.sessionRegistry.Remove(key)
 	}
 
 	if l.forwarder.GetTunnelType() == types.TCP {
-		err := portUtil.Default.SetPortStatus(l.forwarder.GetForwardedPort(), false)
+		err = portUtil.Default.SetPortStatus(l.forwarder.GetForwardedPort(), false)
 		if err != nil {
 			return err
 		}