Merge pull request 'staging' (#60) from staging into main

Reviewed-on: #60

.gitea/workflows/renovate.yml

@@ -1,21 +0,0 @@
-name: renovate
-
-on:
-  schedule:
-    - cron: "0 0 * * *"
-  push:
-    branches:
-      - staging
-
-jobs:
-  renovate:
-    runs-on: ubuntu-latest
-    container: git.fossy.my.id/renovate-clanker/renovate:latest
-    steps:
-      - uses: actions/checkout@v6
-      - run: renovate
-        env:
-          RENOVATE_CONFIG_FILE: ${{ gitea.workspace }}/renovate-config.js
-          LOG_LEVEL: "debug"
-          RENOVATE_TOKEN: ${{ secrets.RENOVATE_TOKEN }}
-          GITHUB_COM_TOKEN: ${{ secrets.COM_TOKEN }}

go.mod

@@ -3,7 +3,7 @@ module tunnel_pls
 go 1.25.5
 
 require (
-	git.fossy.my.id/bagas/tunnel-please-grpc v1.3.0
+	git.fossy.my.id/bagas/tunnel-please-grpc v1.5.0
 	github.com/caddyserver/certmagic v0.25.0
 	github.com/charmbracelet/bubbles v0.21.0
 	github.com/charmbracelet/bubbletea v1.3.10
@@ -52,4 +52,4 @@ require (
 	golang.org/x/text v0.32.0 // indirect
 	golang.org/x/tools v0.40.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b // indirect
-)
+)

go.sum

@@ -1,5 +1,9 @@
 git.fossy.my.id/bagas/tunnel-please-grpc v1.3.0 h1:RhcBKUG41/om4jgN+iF/vlY/RojTeX1QhBa4p4428ec=
 git.fossy.my.id/bagas/tunnel-please-grpc v1.3.0/go.mod h1:fG+VkArdkceGB0bNA7IFQus9GetLAwdF5Oi4jdMlXtY=
+git.fossy.my.id/bagas/tunnel-please-grpc v1.4.0 h1:tpJSKjaSmV+vxxbVx6qnStjxFVXjj2M0rygWXxLb99o=
+git.fossy.my.id/bagas/tunnel-please-grpc v1.4.0/go.mod h1:fG+VkArdkceGB0bNA7IFQus9GetLAwdF5Oi4jdMlXtY=
+git.fossy.my.id/bagas/tunnel-please-grpc v1.5.0 h1:3xszIhck4wo9CoeRq9vnkar4PhY7kz9QrR30qj2XszA=
+git.fossy.my.id/bagas/tunnel-please-grpc v1.5.0/go.mod h1:Weh6ZujgWmT8XxD3Qba7sJ6r5eyUMB9XSWynqdyOoLo=
 github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4=
 github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=

internal/grpc/client/client.go

@@ -210,114 +210,152 @@ func (c *Client) SubscribeEvents(ctx context.Context, identity, authToken string
 }
 
 func (c *Client) processEventStream(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events]) error {
+	handlers := c.eventHandlers(subscribe)
+
 	for {
 		recv, err := subscribe.Recv()
 		if err != nil {
 			return err
 		}
-		switch recv.GetType() {
-		case proto.EventType_SLUG_CHANGE:
-			oldSlug := recv.GetSlugEvent().GetOld()
-			newSlug := recv.GetSlugEvent().GetNew()
-			var userSession *session.SSHSession
-			userSession, err = c.sessionRegistry.Get(types.SessionKey{
-				Id:   oldSlug,
-				Type: types.HTTP,
-			})
-			if err != nil {
-				errSend := subscribe.Send(&proto.Node{
-					Type: proto.EventType_SLUG_CHANGE_RESPONSE,
-					Payload: &proto.Node_SlugEventResponse{
-						SlugEventResponse: &proto.SlugChangeEventResponse{
-							Success: false,
-							Message: err.Error(),
-						},
-					},
-				})
-				if errSend != nil {
-					if c.isConnectionError(errSend) {
-						return errSend
-					}
-					log.Printf("non-connection send error for slug change failure: %v", errSend)
-				}
-				continue
-			}
-			err = c.sessionRegistry.Update(types.SessionKey{
-				Id:   oldSlug,
-				Type: types.HTTP,
-			}, types.SessionKey{
-				Id:   newSlug,
-				Type: types.HTTP,
-			})
-			if err != nil {
-				errSend := subscribe.Send(&proto.Node{
-					Type: proto.EventType_SLUG_CHANGE_RESPONSE,
-					Payload: &proto.Node_SlugEventResponse{
-						SlugEventResponse: &proto.SlugChangeEventResponse{
-							Success: false,
-							Message: err.Error(),
-						},
-					},
-				})
-				if errSend != nil {
-					if c.isConnectionError(errSend) {
-						return errSend
-					}
-					log.Printf("non-connection send error for slug change failure: %v", errSend)
-				}
-				continue
-			}
-			userSession.GetInteraction().Redraw()
-			err = subscribe.Send(&proto.Node{
-				Type: proto.EventType_SLUG_CHANGE_RESPONSE,
-				Payload: &proto.Node_SlugEventResponse{
-					SlugEventResponse: &proto.SlugChangeEventResponse{
-						Success: true,
-						Message: "",
-					},
-				},
-			})
-			if err != nil {
-				if c.isConnectionError(err) {
-					log.Printf("connection error sending slug change success: %v", err)
-					return err
-				}
-				log.Printf("non-connection send error for slug change success: %v", err)
-				continue
-			}
-		case proto.EventType_GET_SESSIONS:
-			sessions := c.sessionRegistry.GetAllSessionFromUser(recv.GetGetSessionsEvent().GetIdentity())
-			var details []*proto.Detail
-			for _, ses := range sessions {
-				detail := ses.Detail()
-				details = append(details, &proto.Detail{
-					Node:           config.Getenv("DOMAIN", "localhost"),
-					ForwardingType: detail.ForwardingType,
-					Slug:           detail.Slug,
-					UserId:         detail.UserID,
-					Active:         detail.Active,
-					StartedAt:      timestamppb.New(detail.StartedAt),
-				})
-			}
-			err = subscribe.Send(&proto.Node{
-				Type: proto.EventType_GET_SESSIONS,
-				Payload: &proto.Node_GetSessionsEvent{
-					GetSessionsEvent: &proto.GetSessionsResponse{
-						Details: details,
-					},
-				},
-			})
-			if err != nil {
-				if c.isConnectionError(err) {
-					log.Printf("connection error sending sessions success: %v", err)
-					return err
-				}
-				log.Printf("non-connection send error for sessions success: %v", err)
-				continue
-			}
-		default:
+
+		handler, ok := handlers[recv.GetType()]
+		if !ok {
 			log.Printf("Unknown event type received: %v", recv.GetType())
+			continue
 		}
+
+		if err = handler(recv); err != nil {
+			return err
+		}
+	}
+}
+
+func (c *Client) eventHandlers(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events]) map[proto.EventType]func(*proto.Events) error {
+	return map[proto.EventType]func(*proto.Events) error{
+		proto.EventType_SLUG_CHANGE:       func(evt *proto.Events) error { return c.handleSlugChange(subscribe, evt) },
+		proto.EventType_GET_SESSIONS:      func(evt *proto.Events) error { return c.handleGetSessions(subscribe, evt) },
+		proto.EventType_TERMINATE_SESSION: func(evt *proto.Events) error { return c.handleTerminateSession(subscribe, evt) },
+	}
+}
+
+func (c *Client) handleSlugChange(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events], evt *proto.Events) error {
+	slugEvent := evt.GetSlugEvent()
+	user := slugEvent.GetUser()
+	oldSlug := slugEvent.GetOld()
+	newSlug := slugEvent.GetNew()
+
+	userSession, err := c.sessionRegistry.Get(types.SessionKey{Id: oldSlug, Type: types.HTTP})
+	if err != nil {
+		return c.sendNode(subscribe, &proto.Node{
+			Type: proto.EventType_SLUG_CHANGE_RESPONSE,
+			Payload: &proto.Node_SlugEventResponse{
+				SlugEventResponse: &proto.SlugChangeEventResponse{Success: false, Message: err.Error()},
+			},
+		}, "slug change failure response")
+	}
+
+	if err = c.sessionRegistry.Update(user, types.SessionKey{Id: oldSlug, Type: types.HTTP}, types.SessionKey{Id: newSlug, Type: types.HTTP}); err != nil {
+		return c.sendNode(subscribe, &proto.Node{
+			Type: proto.EventType_SLUG_CHANGE_RESPONSE,
+			Payload: &proto.Node_SlugEventResponse{
+				SlugEventResponse: &proto.SlugChangeEventResponse{Success: false, Message: err.Error()},
+			},
+		}, "slug change failure response")
+	}
+
+	userSession.GetInteraction().Redraw()
+	return c.sendNode(subscribe, &proto.Node{
+		Type: proto.EventType_SLUG_CHANGE_RESPONSE,
+		Payload: &proto.Node_SlugEventResponse{
+			SlugEventResponse: &proto.SlugChangeEventResponse{Success: true, Message: ""},
+		},
+	}, "slug change success response")
+}
+
+func (c *Client) handleGetSessions(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events], evt *proto.Events) error {
+	sessions := c.sessionRegistry.GetAllSessionFromUser(evt.GetGetSessionsEvent().GetIdentity())
+
+	var details []*proto.Detail
+	for _, ses := range sessions {
+		detail := ses.Detail()
+		details = append(details, &proto.Detail{
+			Node:           config.Getenv("DOMAIN", "localhost"),
+			ForwardingType: detail.ForwardingType,
+			Slug:           detail.Slug,
+			UserId:         detail.UserID,
+			Active:         detail.Active,
+			StartedAt:      timestamppb.New(detail.StartedAt),
+		})
+	}
+
+	return c.sendNode(subscribe, &proto.Node{
+		Type: proto.EventType_GET_SESSIONS,
+		Payload: &proto.Node_GetSessionsEvent{
+			GetSessionsEvent: &proto.GetSessionsResponse{Details: details},
+		},
+	}, "send get sessions response")
+}
+
+func (c *Client) handleTerminateSession(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events], evt *proto.Events) error {
+	terminate := evt.GetTerminateSessionEvent()
+	user := terminate.GetUser()
+	slug := terminate.GetSlug()
+
+	tunnelType, err := c.protoToTunnelType(terminate.GetTunnelType())
+	if err != nil {
+		return c.sendNode(subscribe, &proto.Node{
+			Type: proto.EventType_TERMINATE_SESSION,
+			Payload: &proto.Node_TerminateSessionEventResponse{
+				TerminateSessionEventResponse: &proto.TerminateSessionEventResponse{Success: false, Message: err.Error()},
+			},
+		}, "terminate session invalid tunnel type")
+	}
+
+	userSession, err := c.sessionRegistry.GetWithUser(user, types.SessionKey{Id: slug, Type: tunnelType})
+	if err != nil {
+		return c.sendNode(subscribe, &proto.Node{
+			Type: proto.EventType_TERMINATE_SESSION,
+			Payload: &proto.Node_TerminateSessionEventResponse{
+				TerminateSessionEventResponse: &proto.TerminateSessionEventResponse{Success: false, Message: err.Error()},
+			},
+		}, "terminate session fetch failed")
+	}
+
+	if err = userSession.GetLifecycle().Close(); err != nil {
+		return c.sendNode(subscribe, &proto.Node{
+			Type: proto.EventType_TERMINATE_SESSION,
+			Payload: &proto.Node_TerminateSessionEventResponse{
+				TerminateSessionEventResponse: &proto.TerminateSessionEventResponse{Success: false, Message: err.Error()},
+			},
+		}, "terminate session close failed")
+	}
+
+	return c.sendNode(subscribe, &proto.Node{
+		Type: proto.EventType_TERMINATE_SESSION,
+		Payload: &proto.Node_TerminateSessionEventResponse{
+			TerminateSessionEventResponse: &proto.TerminateSessionEventResponse{Success: true, Message: ""},
+		},
+	}, "terminate session success response")
+}
+
+func (c *Client) sendNode(subscribe grpc.BidiStreamingClient[proto.Node, proto.Events], node *proto.Node, context string) error {
+	if err := subscribe.Send(node); err != nil {
+		if c.isConnectionError(err) {
+			return err
+		}
+		log.Printf("%s: %v", context, err)
+	}
+	return nil
+}
+
+func (c *Client) protoToTunnelType(t proto.TunnelType) (types.TunnelType, error) {
+	switch t {
+	case proto.TunnelType_HTTP:
+		return types.HTTP, nil
+	case proto.TunnelType_TCP:
+		return types.TCP, nil
+	default:
+		return types.UNKNOWN, fmt.Errorf("unknown tunnel type received")
 	}
 }
 

renovate-config.js

@@ -1,8 +0,0 @@
-module.exports = {
-    "endpoint": "https://git.fossy.my.id/api/v1",
-    "gitAuthor": "Renovate-Clanker <renovate-bot@fossy.my.id>",
-    "platform": "gitea",
-    "onboardingConfigFileName": "renovate.json",
-    "autodiscover": true,
-    "optimizeForDisabled": true,
-};

server/server.go

@@ -2,9 +2,11 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net"
+	"time"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/grpc/client"
 	"tunnel_pls/session"
@@ -64,30 +66,32 @@ func (s *Server) Start() {
 
 func (s *Server) handleConnection(conn net.Conn) {
 	sshConn, chans, forwardingReqs, err := ssh.NewServerConn(conn, s.config)
-	defer func(sshConn *ssh.ServerConn) {
-		err = sshConn.Close()
-		if err != nil {
-			log.Printf("failed to close SSH server: %v", err)
-		}
-	}(sshConn)
-
 	if err != nil {
 		log.Printf("failed to establish SSH connection: %v", err)
-		err := conn.Close()
+		err = conn.Close()
 		if err != nil {
 			log.Printf("failed to close SSH connection: %v", err)
 			return
 		}
 		return
 	}
-	ctx := context.Background()
-	log.Println("SSH connection established:", sshConn.User())
+
+	defer func(sshConn *ssh.ServerConn) {
+		err = sshConn.Close()
+		if err != nil && !errors.Is(err, net.ErrClosed) {
+			log.Printf("failed to close SSH server: %v", err)
+		}
+	}(sshConn)
 
 	user := "UNAUTHORIZED"
 	if s.grpcClient != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 		_, u, _ := s.grpcClient.AuthorizeConn(ctx, sshConn.User())
 		user = u
+		cancel()
 	}
+
+	log.Println("SSH connection established:", sshConn.User())
 	sshSession := session.New(sshConn, forwardingReqs, chans, s.sessionRegistry, user)
 	err = sshSession.Start()
 	if err != nil {

session/interaction/interaction.go

@@ -23,15 +23,20 @@ import (
 
 type Lifecycle interface {
 	Close() error
+	GetUser() string
+}
+
+type SessionRegistry interface {
+	Update(user string, oldKey, newKey types.SessionKey) error
 }
 
 type Controller interface {
 	SetChannel(channel ssh.Channel)
 	SetLifecycle(lifecycle Lifecycle)
-	SetSlugModificator(func(oldSlug, newSlug string) error)
 	Start()
 	SetWH(w, h int)
 	Redraw()
+	SetSessionRegistry(registry SessionRegistry)
 }
 
 type Forwarder interface {
@@ -41,14 +46,14 @@ type Forwarder interface {
 }
 
 type Interaction struct {
-	channel          ssh.Channel
-	slugManager      slug.Manager
-	forwarder        Forwarder
-	lifecycle        Lifecycle
-	updateClientSlug func(oldSlug, newSlug string) error
-	program          *tea.Program
-	ctx              context.Context
-	cancel           context.CancelFunc
+	channel         ssh.Channel
+	slugManager     slug.Manager
+	forwarder       Forwarder
+	lifecycle       Lifecycle
+	sessionRegistry SessionRegistry
+	program         *tea.Program
+	ctx             context.Context
+	cancel          context.CancelFunc
 }
 
 func (i *Interaction) SetWH(w, h int) {
@@ -102,17 +107,21 @@ type tickMsg time.Time
 func NewInteraction(slugManager slug.Manager, forwarder Forwarder) *Interaction {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Interaction{
-		channel:          nil,
-		slugManager:      slugManager,
-		forwarder:        forwarder,
-		lifecycle:        nil,
-		updateClientSlug: nil,
-		program:          nil,
-		ctx:              ctx,
-		cancel:           cancel,
+		channel:         nil,
+		slugManager:     slugManager,
+		forwarder:       forwarder,
+		lifecycle:       nil,
+		sessionRegistry: nil,
+		program:         nil,
+		ctx:             ctx,
+		cancel:          cancel,
 	}
 }
 
+func (i *Interaction) SetSessionRegistry(registry SessionRegistry) {
+	i.sessionRegistry = registry
+}
+
 func (i *Interaction) SetLifecycle(lifecycle Lifecycle) {
 	i.lifecycle = lifecycle
 }
@@ -121,10 +130,6 @@ func (i *Interaction) SetChannel(channel ssh.Channel) {
 	i.channel = channel
 }
 
-func (i *Interaction) SetSlugModificator(modificator func(oldSlug, newSlug string) error) {
-	i.updateClientSlug = modificator
-}
-
 func (i *Interaction) Stop() {
 	if i.cancel != nil {
 		i.cancel()
@@ -218,7 +223,13 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				return m, tea.Batch(tea.ClearScreen, textinput.Blink)
 			case "enter":
 				inputValue := m.slugInput.Value()
-				if err := m.interaction.updateClientSlug(m.interaction.slugManager.Get(), inputValue); err != nil {
+				if err := m.interaction.sessionRegistry.Update(m.interaction.lifecycle.GetUser(), types.SessionKey{
+					Id:   m.interaction.slugManager.Get(),
+					Type: types.HTTP,
+				}, types.SessionKey{
+					Id:   inputValue,
+					Type: types.HTTP,
+				}); err != nil {
 					m.slugError = err.Error()
 					return m, nil
 				}
@@ -661,22 +672,32 @@ func (m *model) View() string {
 		MarginBottom(boxMargin).
 		Width(boxMaxWidth)
 
-	urlDisplay := m.getTunnelURL()
-	if shouldUseCompactLayout(m.width, 80) && len(urlDisplay) > m.width-20 {
-		maxLen := m.width - 25
-		if maxLen > 10 {
-			urlDisplay = truncateString(urlDisplay, maxLen)
-		}
-	}
+	authenticatedUser := m.interaction.lifecycle.GetUser()
+
+	userInfoStyle := lipgloss.NewStyle().
+		Foreground(lipgloss.Color("#FAFAFA")).
+		Bold(true)
+
+	sectionHeaderStyle := lipgloss.NewStyle().
+		Foreground(lipgloss.Color("#888888")).
+		Bold(true)
+
+	addressStyle := lipgloss.NewStyle().
+		Foreground(lipgloss.Color("#FAFAFA"))
 
 	var infoContent string
 	if shouldUseCompactLayout(m.width, 70) {
-		infoContent = fmt.Sprintf("🌐 %s", urlBoxStyle.Render(urlDisplay))
-	} else if isCompact {
-		infoContent = fmt.Sprintf("🌐  Forwarding to:\n\n     %s", urlBoxStyle.Render(urlDisplay))
+		infoContent = fmt.Sprintf("👤 %s\n\n%s\n%s",
+			userInfoStyle.Render(authenticatedUser),
+			sectionHeaderStyle.Render("🌐 FORWARDING ADDRESS:"),
+			addressStyle.Render(fmt.Sprintf("   %s", urlBoxStyle.Render(m.getTunnelURL()))))
 	} else {
-		infoContent = fmt.Sprintf("🌐  F O R W A R D I N G   T O:\n\n     %s", urlBoxStyle.Render(urlDisplay))
+		infoContent = fmt.Sprintf("👤  Authenticated as: %s\n\n%s\n     %s",
+			userInfoStyle.Render(authenticatedUser),
+			sectionHeaderStyle.Render("🌐  FORWARDING ADDRESS:"),
+			addressStyle.Render(urlBoxStyle.Render(m.getTunnelURL())))
 	}
+
 	b.WriteString(responsiveInfoBox.Render(infoContent))
 	b.WriteString("\n")
 

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
 		}

session/registry.go

@@ -10,7 +10,8 @@ type Key = types.SessionKey
 
 type Registry interface {
 	Get(key Key) (session *SSHSession, err error)
-	Update(oldKey, newKey Key) error
+	GetWithUser(user string, key Key) (session *SSHSession, err error)
+	Update(user string, oldKey, newKey Key) error
 	Register(key Key, session *SSHSession) (success bool)
 	Remove(key Key)
 	GetAllSessionFromUser(user string) []*SSHSession
@@ -44,7 +45,18 @@ func (r *registry) Get(key Key) (session *SSHSession, err error) {
 	return client, nil
 }
 
-func (r *registry) Update(oldKey, newKey Key) error {
+func (r *registry) GetWithUser(user string, key Key) (session *SSHSession, err error) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	client, ok := r.byUser[user][key]
+	if !ok {
+		return nil, fmt.Errorf("session not found")
+	}
+	return client, nil
+}
+
+func (r *registry) Update(user string, oldKey, newKey Key) error {
 	if oldKey.Type != newKey.Type {
 		return fmt.Errorf("tunnel type cannot change")
 	}
@@ -64,30 +76,24 @@ func (r *registry) Update(oldKey, newKey Key) error {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 
-	userID, ok := r.slugIndex[oldKey]
-	if !ok {
-		return fmt.Errorf("session not found")
-	}
-
 	if _, exists := r.slugIndex[newKey]; exists && newKey != oldKey {
 		return fmt.Errorf("someone already uses this subdomain")
 	}
-
-	client, ok := r.byUser[userID][oldKey]
+	client, ok := r.byUser[user][oldKey]
 	if !ok {
 		return fmt.Errorf("session not found")
 	}
 
-	delete(r.byUser[userID], oldKey)
+	delete(r.byUser[user], oldKey)
 	delete(r.slugIndex, oldKey)
 
 	client.slugManager.Set(newKey.Id)
-	r.slugIndex[newKey] = userID
+	r.slugIndex[newKey] = user
 
-	if r.byUser[userID] == nil {
-		r.byUser[userID] = make(map[Key]*SSHSession)
+	if r.byUser[user] == nil {
+		r.byUser[user] = make(map[Key]*SSHSession)
 	}
-	r.byUser[userID][newKey] = client
+	r.byUser[user][newKey] = client
 	return nil
 }
 
@@ -99,7 +105,7 @@ func (r *registry) Register(key Key, session *SSHSession) (success bool) {
 		return false
 	}
 
-	userID := session.userID
+	userID := session.lifecycle.GetUser()
 	if r.byUser[userID] == nil {
 		r.byUser[userID] = make(map[Key]*SSHSession)
 	}

session/session.go

@@ -9,7 +9,6 @@ import (
 	"tunnel_pls/session/interaction"
 	"tunnel_pls/session/lifecycle"
 	"tunnel_pls/session/slug"
-	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -29,7 +28,6 @@ type SSHSession struct {
 	forwarder     forwarder.ForwardingController
 	slugManager   slug.Manager
 	registry      Registry
-	userID        string
 }
 
 func (s *SSHSession) GetLifecycle() lifecycle.SessionLifecycle {
@@ -48,22 +46,16 @@ func (s *SSHSession) GetSlugManager() slug.Manager {
 	return s.slugManager
 }
 
-func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel, sessionRegistry Registry, userID string) *SSHSession {
+func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan ssh.NewChannel, sessionRegistry Registry, user string) *SSHSession {
 	slugManager := slug.NewManager()
 	forwarderManager := forwarder.NewForwarder(slugManager)
 	interactionManager := interaction.NewInteraction(slugManager, forwarderManager)
-	lifecycleManager := lifecycle.NewLifecycle(conn, forwarderManager, slugManager)
+	lifecycleManager := lifecycle.NewLifecycle(conn, forwarderManager, slugManager, user)
 
 	interactionManager.SetLifecycle(lifecycleManager)
-	interactionManager.SetSlugModificator(func(oldSlug, newSlug string) error {
-		oldKey := types.SessionKey{Id: oldSlug, Type: forwarderManager.GetTunnelType()}
-		newKey := types.SessionKey{Id: newSlug, Type: forwarderManager.GetTunnelType()}
-		return sessionRegistry.Update(oldKey, newKey)
-	})
 	forwarderManager.SetLifecycle(lifecycleManager)
-	lifecycleManager.SetUnregisterClient(func(key types.SessionKey) {
-		sessionRegistry.Remove(key)
-	})
+	interactionManager.SetSessionRegistry(sessionRegistry)
+	lifecycleManager.SetSessionRegistry(sessionRegistry)
 
 	return &SSHSession{
 		initialReq:    forwardingReq,
@@ -73,7 +65,6 @@ func New(conn *ssh.ServerConn, forwardingReq <-chan *ssh.Request, sshChan <-chan
 		forwarder:     forwarderManager,
 		slugManager:   slugManager,
 		registry:      sessionRegistry,
-		userID:        userID,
 	}
 }
 
@@ -89,7 +80,7 @@ func (s *SSHSession) Detail() Detail {
 	return Detail{
 		ForwardingType: string(s.forwarder.GetTunnelType()),
 		Slug:           s.slugManager.Get(),
-		UserID:         s.userID,
+		UserID:         s.lifecycle.GetUser(),
 		Active:         s.lifecycle.IsActive(),
 		StartedAt:      s.lifecycle.StartedAt(),
 	}

types/types.go

@@ -11,8 +11,9 @@ const (
 type TunnelType string
 
 const (
-	HTTP TunnelType = "HTTP"
-	TCP  TunnelType = "TCP"
+	UNKNOWN TunnelType = "UNKNOWN"
+	HTTP    TunnelType = "HTTP"
+	TCP     TunnelType = "TCP"
 )
 
 type SessionKey struct {