feat: upgrade bubbletea to v2

go.mod

@@ -3,10 +3,10 @@ module tunnel_pls
 go 1.26.0
 
 require (
+	charm.land/bubbletea/v2 v2.0.0
 	git.fossy.my.id/bagas/tunnel-please-grpc v1.5.0
 	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
 	github.com/joho/godotenv v1.5.1
 	github.com/libdns/cloudflare v0.2.2
@@ -21,10 +21,14 @@ 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.4 // indirect
+	github.com/charmbracelet/bubbletea v1.3.10 // indirect
 	github.com/charmbracelet/colorprofile v0.4.1 // indirect
+	github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8 // indirect
 	github.com/charmbracelet/x/ansi v0.11.6 // indirect
 	github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
 	github.com/charmbracelet/x/term v0.2.2 // indirect
+	github.com/charmbracelet/x/termios v0.1.1 // indirect
+	github.com/charmbracelet/x/windows v0.2.2 // indirect
 	github.com/clipperhouse/displaywidth v0.9.0 // indirect
 	github.com/clipperhouse/stringish v0.1.1 // indirect
 	github.com/clipperhouse/uax29/v2 v2.5.0 // indirect

go.sum

@@ -1,3 +1,5 @@
+charm.land/bubbletea/v2 v2.0.0 h1:p0d6CtWyJXJ9GfzMpUUqbP/XUUhhlk06+vCKWmox1wQ=
+charm.land/bubbletea/v2 v2.0.0/go.mod h1:3LRff2U4WIYXy7MTxfbAQ+AdfM3D8Xuvz2wbsOD9OHQ=
 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=
@@ -18,14 +20,20 @@ github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco
 github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk=
 github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY=
 github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30=
+github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8 h1:eyFRbAmexyt43hVfeyBofiGSEmJ7krjLOYt/9CF5NKA=
+github.com/charmbracelet/ultraviolet v0.0.0-20260205113103-524a6607adb8/go.mod h1:SQpCTRNBtzJkwku5ye4S3HEuthAlGy2n9VXZnWkEW98=
 github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
 github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
 github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
 github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
-github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91 h1:payRxjMjKgx2PaCWLZ4p3ro9y97+TVLZNaRZgJwSVDQ=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241212170349-ad4b7ae0f25f h1:UytXHv0UxnsDFmL/7Z9Q5SBYPwSuRLXHbwx+6LycZ2w=
-github.com/charmbracelet/x/exp/golden v0.0.0-20241011142426-46044092ad91/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
+github.com/charmbracelet/x/exp/golden v0.0.0-20241212170349-ad4b7ae0f25f/go.mod h1:wDlXFlCrmJ8J+swcL/MnGUuYnqgQdW9rhSD61oNMb6U=
 github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk=
 github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI=
+github.com/charmbracelet/x/termios v0.1.1 h1:o3Q2bT8eqzGnGPOYheoYS8eEleT5ZVNYNy8JawjaNZY=
+github.com/charmbracelet/x/termios v0.1.1/go.mod h1:rB7fnv1TgOPOyyKRJ9o+AsTU/vK5WHJ2ivHeut/Pcwo=
+github.com/charmbracelet/x/windows v0.2.2 h1:IofanmuvaxnKHuV04sC0eBy/smG6kIKrWG2/jYn2GuM=
+github.com/charmbracelet/x/windows v0.2.2/go.mod h1:/8XtdKZzedat74NQFn0NGlGL4soHB0YQZrETF96h75k=
 github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA=
 github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA=
 github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=

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/internal/session/slug"
+	"tunnel_pls/session/slug"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"

internal/config/config.go

@@ -1,8 +1,6 @@
 package config
 
-import (
+import "tunnel_pls/types"
-	"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/internal/types"
+	"tunnel_pls/types"
 
 	"github.com/stretchr/testify/assert"
 )

internal/config/loader.go

@@ -6,7 +6,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/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/internal/types"
+	"tunnel_pls/types"
 
 	proto "git.fossy.my.id/bagas/tunnel-please-grpc/gen"
 	"google.golang.org/grpc"

internal/grpc/client/client_test.go

@@ -7,13 +7,14 @@ 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"
@@ -884,16 +885,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) error { return m.Called(channel).Error(0) }
+func (m *mockLifecycle) SetChannel(channel ssh.Channel)       { m.Called(channel) }
 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() lifecycle.PortRegistry {
+func (m *mockLifecycle) PortRegistry() port.Port {
 	args := m.Called()
 	if args.Get(0) == nil {
 		return nil
 	}
-	return args.Get(0).(lifecycle.PortRegistry)
+	return args.Get(0).(port.Port)
 }
 
 type mockEventServiceClient struct {

internal/registry/registry.go

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

internal/registry/registry_test.go

@@ -4,11 +4,12 @@ import (
 	"sync"
 	"testing"
 	"time"
-	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/internal/port"
-	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/session/forwarder"
-	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/session/interaction"
-	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/session/slug"
+	"tunnel_pls/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -77,15 +78,15 @@ func (ml *mockLifecycle) Connection() ssh.Conn {
 	return args.Get(0).(ssh.Conn)
 }
 
-func (ml *mockLifecycle) PortRegistry() lifecycle.PortRegistry {
+func (ml *mockLifecycle) PortRegistry() port.Port {
 	args := ml.Called()
 	if args.Get(0) == nil {
 		return nil
 	}
-	return args.Get(0).(lifecycle.PortRegistry)
+	return args.Get(0).(port.Port)
 }
 
-func (ml *mockLifecycle) SetChannel(channel ssh.Channel) error { return ml.Called(channel).Error(0) }
+func (ml *mockLifecycle) SetChannel(channel ssh.Channel)       { ml.Called(channel) }
 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/transport/http_test.go

@@ -55,7 +55,8 @@ func TestHTTPServer_Serve(t *testing.T) {
 
 	go func() {
 		time.Sleep(100 * time.Millisecond)
-		_ = listener.Close()
+		err = listener.Close()
+		assert.NoError(t, err)
 	}()
 
 	err = srv.Serve(listener)

internal/transport/httphandler.go

@@ -1,7 +1,6 @@
 package transport
 
 import (
-	"bufio"
 	"bytes"
 	"context"
 	"errors"
@@ -17,7 +16,7 @@ import (
 	"tunnel_pls/internal/http/stream"
 	"tunnel_pls/internal/middleware"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -53,37 +52,25 @@ func (hh *httpHandler) badRequest(conn net.Conn) error {
 	return nil
 }
 
-func readHTTPHeader(br *bufio.Reader, limit int) ([]byte, error) {
-	var headerBuf []byte
-	for {
-		line, err := br.ReadBytes('\n')
-		headerBuf = append(headerBuf, line...)
-		if err != nil {
-			return nil, err
-		}
-		if bytes.HasSuffix(headerBuf, []byte("\r\n\r\n")) {
-			return headerBuf, nil
-		}
-		if len(headerBuf) > limit {
-			return nil, fmt.Errorf("headers too large")
-		}
-	}
-}
-
 func (hh *httpHandler) Handler(conn net.Conn, isTLS bool) {
 	defer hh.closeConnection(conn)
 
 	_ = conn.SetReadDeadline(time.Now().Add(10 * time.Second))
-	br := bufio.NewReaderSize(conn, hh.config.HeaderSize())
+	buf := make([]byte, hh.config.HeaderSize())
-	headerBuf, err := readHTTPHeader(br, hh.config.HeaderSize())
+	n, err := conn.Read(buf)
 	if err != nil {
 		_ = hh.badRequest(conn)
 		return
 	}
 
+	if idx := bytes.Index(buf[:n], []byte("\r\n\r\n")); idx == -1 {
+		_ = hh.badRequest(conn)
+		return
+	}
+
 	_ = conn.SetReadDeadline(time.Time{})
 
-	reqhf, err := header.NewRequest(headerBuf)
+	reqhf, err := header.NewRequest(buf[:n])
 	if err != nil {
 		log.Printf("Error creating request header: %v", err)
 		_ = hh.badRequest(conn)
@@ -114,7 +101,7 @@ func (hh *httpHandler) Handler(conn net.Conn, isTLS bool) {
 		return
 	}
 
-	hw := stream.New(conn, br, conn.RemoteAddr())
+	hw := stream.New(conn, conn, conn.RemoteAddr())
 	defer func(hw stream.HTTP) {
 		err = hw.Close()
 		if err != nil {

internal/transport/httphandler_test.go

@@ -11,11 +11,11 @@ import (
 	"testing"
 	"time"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/internal/session/forwarder"
+	"tunnel_pls/session/forwarder"
-	"tunnel_pls/internal/session/interaction"
+	"tunnel_pls/session/interaction"
-	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/session/slug"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 
@@ -321,14 +321,8 @@ func TestHandler(t *testing.T) {
 			isTLS:       false,
 			redirectTLS: false,
 			request:     []byte(""),
-			expected:    []byte(""),
+			expected:    []byte("HTTP/1.1 400 Bad Request\r\n\r\n"),
-			setupConn: func() (net.Conn, net.Conn) {
+			setupMocks: func(msr *MockSessionRegistry) {
-				mc := new(MockConn)
-				mc.ReadBuffer = bytes.NewBuffer(nil)
-				mc.On("SetReadDeadline", mock.Anything).Return(nil)
-				mc.On("Write", []byte("HTTP/1.1 400 Bad Request\r\n\r\n")).Return(0, nil)
-				mc.On("Close").Return(nil)
-				return mc, nil
 			},
 		},
 		{
@@ -721,113 +715,3 @@ func TestHandler(t *testing.T) {
 		})
 	}
 }
-
-func TestHandlerForwardsPostBody(t *testing.T) {
-	mockSessionRegistry := new(MockSessionRegistry)
-	mockConfig := &MockConfig{}
-	mockConfig.On("Domain").Return("example.com")
-	mockConfig.On("HTTPPort").Return("0")
-	mockConfig.On("HeaderSize").Return(4096)
-	mockConfig.On("TLSRedirect").Return(true)
-	hh := &httpHandler{
-		sessionRegistry: mockSessionRegistry,
-		config:          mockConfig,
-	}
-
-	mockSession := new(MockSession)
-	mockForwarder := new(MockForwarder)
-	mockSSHChannel := new(MockSSHChannel)
-
-	mockSessionRegistry.On("Get", types.SessionKey{
-		Id:   "test",
-		Type: types.TunnelTypeHTTP,
-	}).Return(mockSession, nil)
-	mockSession.On("Forwarder").Return(mockForwarder)
-
-	reqCh := make(chan *ssh.Request)
-	mockForwarder.On("OpenForwardedChannel", mock.Anything, mock.Anything).Return(mockSSHChannel, (<-chan *ssh.Request)(reqCh), nil)
-
-	var mu sync.Mutex
-	var capturedHeaders []byte
-	mockSSHChannel.On("Write", mock.Anything).Run(func(args mock.Arguments) {
-		mu.Lock()
-		capturedHeaders = append(capturedHeaders, args.Get(0).([]byte)...)
-		mu.Unlock()
-	}).Return(0, nil)
-	mockSSHChannel.On("Close").Return(nil)
-
-	bodyChan := make(chan string, 1)
-	mockForwarder.On("HandleConnection", mock.Anything, mockSSHChannel).Run(func(args mock.Arguments) {
-		w := args.Get(0).(io.ReadWriter)
-		buf := make([]byte, len("hello=world"))
-		if _, err := io.ReadFull(w, buf); err != nil {
-			bodyChan <- ""
-		} else {
-			bodyChan <- string(buf)
-		}
-		_, _ = w.Write([]byte("HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok"))
-	})
-
-	go func() {
-		for range reqCh {
-		}
-	}()
-
-	serverConn, clientConn := net.Pipe()
-	defer func() {
-		_ = clientConn.Close()
-	}()
-
-	remoteAddr, _ := net.ResolveTCPAddr("tcp", "127.0.0.1:12345")
-	wrappedServerConn := &wrappedConn{Conn: serverConn, remoteAddr: remoteAddr}
-
-	go hh.Handler(wrappedServerConn, true)
-
-	request := []byte("POST / HTTP/1.1\r\nHost: test.domain\r\nContent-Type: application/x-www-form-urlencoded\r\nContent-Length: 11\r\n\r\nhello=world")
-	go func() {
-		_, _ = clientConn.Write(request)
-	}()
-
-	var response []byte
-	respDone := make(chan struct{})
-	go func() {
-		defer close(respDone)
-		buf := make([]byte, 4096)
-		for {
-			n, err := clientConn.Read(buf)
-			if err != nil {
-				break
-			}
-			response = append(response, buf[:n]...)
-			if bytes.Contains(response, []byte("\r\n\r\nok")) {
-				break
-			}
-		}
-	}()
-
-	select {
-	case body := <-bodyChan:
-		assert.Equal(t, "hello=world", body)
-	case <-time.After(5 * time.Second):
-		t.Fatal("timed out waiting for forwarded body")
-	}
-
-	select {
-	case <-respDone:
-		resStr := string(response)
-		assert.True(t, strings.HasPrefix(resStr, "HTTP/1.1 200 OK\r\n"))
-		assert.Contains(t, resStr, "Server: Tunnel Please\r\n")
-		assert.True(t, strings.HasSuffix(resStr, "\r\n\r\nok"))
-	case <-time.After(5 * time.Second):
-		t.Fatal("timed out waiting for response")
-	}
-
-	mu.Lock()
-	hdrStr := string(capturedHeaders)
-	mu.Unlock()
-	assert.Contains(t, hdrStr, "POST / HTTP/1.1\r\n")
-	assert.Contains(t, hdrStr, "Content-Length: 11\r\n")
-	assert.Contains(t, hdrStr, "X-Forwarded-For: 127.0.0.1\r\n")
-
-	mockSessionRegistry.AssertExpectations(t)
-}

internal/transport/https_test.go

@@ -63,7 +63,8 @@ func TestHTTPSServer_Serve(t *testing.T) {
 
 	go func() {
 		time.Sleep(100 * time.Millisecond)
-		_ = listener.Close()
+		err = listener.Close()
+		assert.NoError(t, err)
 	}()
 
 	err = srv.Serve(listener)

internal/transport/tcp_test.go

@@ -45,7 +45,8 @@ func TestTCPServer_Serve(t *testing.T) {
 
 	go func() {
 		time.Sleep(100 * time.Millisecond)
-		_ = listener.Close()
+		err = listener.Close()
+		assert.NoError(t, err)
 	}()
 
 	err = srv.Serve(listener)

internal/transport/tls_test.go

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

server/server.go

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

server/server_test.go

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

session/forwarder/forwarder.go

@@ -10,8 +10,8 @@ import (
 	"strconv"
 	"sync"
 	"tunnel_pls/internal/config"
-	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/session/slug"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -28,7 +28,6 @@ type Forwarder interface {
 	Close() error
 }
 type forwarder struct {
-	mu            sync.RWMutex
 	listener      net.Listener
 	tunnelType    types.TunnelType
 	forwardedPort uint16
@@ -61,7 +60,7 @@ func (f *forwarder) copyWithBuffer(dst io.Writer, src io.Reader) (written int64,
 }
 
 func (f *forwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error) {
-	payload := createForwardedTCPIPPayload(origin, f.ForwardedPort())
+	payload := createForwardedTCPIPPayload(origin, f.forwardedPort)
 	type channelResult struct {
 		channel ssh.Channel
 		reqs    <-chan *ssh.Request
@@ -142,44 +141,32 @@ func (f *forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel) {
 }
 
 func (f *forwarder) SetType(tunnelType types.TunnelType) {
-	f.mu.Lock()
-	defer f.mu.Unlock()
 	f.tunnelType = tunnelType
 }
 
 func (f *forwarder) TunnelType() types.TunnelType {
-	f.mu.RLock()
-	defer f.mu.RUnlock()
 	return f.tunnelType
 }
 
 func (f *forwarder) ForwardedPort() uint16 {
-	f.mu.RLock()
-	defer f.mu.RUnlock()
 	return f.forwardedPort
 }
 
 func (f *forwarder) SetForwardedPort(port uint16) {
-	f.mu.Lock()
-	defer f.mu.Unlock()
 	f.forwardedPort = port
 }
 
 func (f *forwarder) SetListener(listener net.Listener) {
-	f.mu.Lock()
-	defer f.mu.Unlock()
 	f.listener = listener
 }
 
 func (f *forwarder) Listener() net.Listener {
-	f.mu.RLock()
-	defer f.mu.RUnlock()
 	return f.listener
 }
 
 func (f *forwarder) Close() error {
-	if listener := f.Listener(); listener != nil {
+	if f.Listener() != nil {
-		return listener.Close()
+		return f.listener.Close()
 	}
 	return nil
 }

session/forwarder/forwarder_test.go

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

session/interaction/coming_soon.go

@@ -3,14 +3,14 @@ package interaction
 import (
 	"strings"
 
-	"github.com/charmbracelet/bubbles/textinput"
+	tea "charm.land/bubbletea/v2"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 )
 
-func (m *model) comingSoonUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+func (m *model) comingSoonUpdate(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
+	_ = msg
 	m.showingComingSoon = false
-	return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+	return m, nil
 }
 
 func (m *model) comingSoonView() string {

session/interaction/commands.go

@@ -4,9 +4,8 @@ import (
 	"strings"
 	"time"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/key"
-	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 )
 
@@ -17,24 +16,22 @@ func (m *model) handleCommandSelection(item commandItem) (tea.Model, tea.Cmd) {
 		m.editingSlug = true
 		m.slugInput.SetValue(m.interaction.slug.String())
 		m.slugInput.Focus()
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	case "tunnel-type":
 		m.showingCommands = false
 		m.showingComingSoon = true
-		return m, tea.Batch(tickCmd(5*time.Second), tea.ClearScreen, textinput.Blink)
+		return m, tickCmd(5 * time.Second)
 	default:
 		m.showingCommands = false
 		return m, nil
 	}
 }
 
-func (m *model) commandsUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+func (m *model) commandsUpdate(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
-	var cmd tea.Cmd
-
 	switch {
 	case key.Matches(msg, m.keymap.quit), msg.String() == "esc":
 		m.showingCommands = false
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	case msg.String() == "enter":
 		selectedItem := m.commandList.SelectedItem()
 		if selectedItem != nil {
@@ -42,8 +39,8 @@ func (m *model) commandsUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 			return m.handleCommandSelection(item)
 		}
 	}
-	m.commandList, cmd = m.commandList.Update(msg)
+	m.commandList, _ = m.commandList.Update(msg)
-	return m, cmd
+	return m, nil
 }
 
 func (m *model) commandsView() string {

session/interaction/dashboard.go

@@ -4,20 +4,19 @@ import (
 	"fmt"
 	"strings"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/key"
-	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 )
 
-func (m *model) dashboardUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+func (m *model) dashboardUpdate(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
 	switch {
 	case key.Matches(msg, m.keymap.quit):
 		m.quitting = true
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink, tea.Quit)
+		return m, tea.Quit
 	case key.Matches(msg, m.keymap.command):
 		m.showingCommands = true
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	}
 	return m, nil
 }

session/interaction/interaction.go

@@ -6,14 +6,14 @@ import (
 	"sync"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/random"
-	"tunnel_pls/internal/session/slug"
+	"tunnel_pls/session/slug"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/help"
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/list"
 	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 	"github.com/muesli/termenv"
 	"golang.org/x/crypto/ssh"
@@ -120,7 +120,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case tickMsg:
 		m.showingComingSoon = false
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 
 	case tea.WindowSizeMsg:
 		m.width = msg.Width
@@ -137,9 +137,9 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 
 	case tea.QuitMsg:
 		m.quitting = true
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink, tea.Quit)
+		return m, tea.Quit
 
-	case tea.KeyMsg:
+	case tea.KeyPressMsg:
 		if m.showingComingSoon {
 			return m.comingSoonUpdate(msg)
 		}
@@ -160,28 +160,28 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 
 func (i *interaction) Redraw() {
 	if i.program != nil {
-		i.program.Send(tea.ClearScreen())
+		i.program.Send(tea.WindowSizeMsg{})
 	}
 }
 
-func (m *model) View() string {
+func (m *model) View() tea.View {
+	var content string
 	if m.quitting {
-		return ""
+		content = ""
+	} else if m.showingComingSoon {
+		content = m.comingSoonView()
+	} else if m.editingSlug {
+		content = m.slugView()
+	} else if m.showingCommands {
+		content = m.commandsView()
+	} else {
+		content = m.dashboardView()
 	}
 
-	if m.showingComingSoon {
+	v := tea.NewView(content)
-		return m.comingSoonView()
+	v.AltScreen = true
-	}
+	v.MouseMode = tea.MouseModeCellMotion
-
+	return v
-	if m.editingSlug {
-		return m.slugView()
-	}
-
-	if m.showingCommands {
-		return m.commandsView()
-	}
-
-	return m.dashboardView()
 }
 
 func (i *interaction) Start() {
@@ -249,8 +249,6 @@ func (i *interaction) Start() {
 		m,
 		tea.WithInput(i.channel),
 		tea.WithOutput(i.channel),
-		tea.WithAltScreen(),
-		tea.WithMouseCellMotion(),
 		tea.WithoutSignals(),
 		tea.WithoutSignalHandler(),
 		tea.WithFPS(30),

session/interaction/interaction_test.go

@@ -7,12 +7,12 @@ import (
 	"net"
 	"testing"
 	"time"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/list"
 	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"golang.org/x/crypto/ssh"
@@ -694,7 +694,7 @@ func TestInteraction_Integration(t *testing.T) {
 func TestModel_Update_KeyMessages(t *testing.T) {
 	tests := []struct {
 		name              string
-		key               tea.KeyMsg
+		key               tea.KeyPressMsg
 		showingComingSoon bool
 		editingSlug       bool
 		showingCommands   bool
@@ -702,25 +702,25 @@ func TestModel_Update_KeyMessages(t *testing.T) {
 	}{
 		{
 			name:              "key press while showing coming soon",
-			key:               tea.KeyMsg{Type: tea.KeyEnter},
+			key:               tea.KeyPressMsg{Code: tea.KeyEnter},
 			showingComingSoon: true,
 			description:       "should call comingSoonUpdate",
 		},
 		{
 			name:        "key press while editing slug",
-			key:         tea.KeyMsg{Type: tea.KeyEnter},
+			key:         tea.KeyPressMsg{Code: tea.KeyEnter},
 			editingSlug: true,
 			description: "should call slugUpdate",
 		},
 		{
 			name:            "key press while showing commands",
-			key:             tea.KeyMsg{Type: tea.KeyEnter},
+			key:             tea.KeyPressMsg{Code: tea.KeyEnter},
 			showingCommands: true,
 			description:     "should call commandsUpdate",
 		},
 		{
 			name:        "key press in dashboard view",
-			key:         tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'c'}},
+			key:         tea.KeyPressMsg{Code: 'c', Text: "c"},
 			description: "should call dashboardUpdate",
 		},
 	}
@@ -777,7 +777,7 @@ func TestModel_SlugUpdate(t *testing.T) {
 	tests := []struct {
 		name           string
 		tunnelType     types.TunnelType
-		keyMsg         tea.KeyMsg
+		keyMsg         tea.KeyPressMsg
 		inputValue     string
 		setupMocks     func(*MockSessionRegistry, *MockSlug, *MockRandom)
 		expectedEdit   bool
@@ -787,21 +787,21 @@ func TestModel_SlugUpdate(t *testing.T) {
 		{
 			name:         "escape key cancels editing",
 			tunnelType:   types.TunnelTypeHTTP,
-			keyMsg:       tea.KeyMsg{Type: tea.KeyEsc},
+			keyMsg:       tea.KeyPressMsg{Code: tea.KeyEsc},
 			setupMocks:   func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {},
 			expectedEdit: false,
 		},
 		{
 			name:         "ctrl+c cancels editing",
 			tunnelType:   types.TunnelTypeHTTP,
-			keyMsg:       tea.KeyMsg{Type: tea.KeyCtrlC},
+			keyMsg:       tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl},
 			setupMocks:   func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {},
 			expectedEdit: false,
 		},
 		{
 			name:       "enter key saves valid slug",
 			tunnelType: types.TunnelTypeHTTP,
-			keyMsg:     tea.KeyMsg{Type: tea.KeyEnter},
+			keyMsg:     tea.KeyPressMsg{Code: tea.KeyEnter},
 			inputValue: "my-custom-slug",
 			setupMocks: func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {
 				ms.On("String").Return("old-slug")
@@ -816,7 +816,7 @@ func TestModel_SlugUpdate(t *testing.T) {
 		{
 			name:       "enter key with error shows error message",
 			tunnelType: types.TunnelTypeHTTP,
-			keyMsg:     tea.KeyMsg{Type: tea.KeyEnter},
+			keyMsg:     tea.KeyPressMsg{Code: tea.KeyEnter},
 			inputValue: "invalid",
 			setupMocks: func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {
 				ms.On("String").Return("old-slug")
@@ -831,7 +831,7 @@ func TestModel_SlugUpdate(t *testing.T) {
 		{
 			name:       "ctrl+r generates random slug",
 			tunnelType: types.TunnelTypeHTTP,
-			keyMsg:     tea.KeyMsg{Type: tea.KeyCtrlR},
+			keyMsg:     tea.KeyPressMsg{Code: 'r', Mod: tea.ModCtrl},
 			setupMocks: func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {
 				mr.On("String", 20).Return("random-generated-slug", nil)
 			},
@@ -841,7 +841,7 @@ func TestModel_SlugUpdate(t *testing.T) {
 		{
 			name:       "ctrl+r with error does nothing",
 			tunnelType: types.TunnelTypeHTTP,
-			keyMsg:     tea.KeyMsg{Type: tea.KeyCtrlR},
+			keyMsg:     tea.KeyPressMsg{Code: 'r', Mod: tea.ModCtrl},
 			setupMocks: func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {
 				mr.On("String", 20).Return("", assert.AnError)
 			},
@@ -850,14 +850,14 @@ func TestModel_SlugUpdate(t *testing.T) {
 		{
 			name:         "regular key updates input",
 			tunnelType:   types.TunnelTypeHTTP,
-			keyMsg:       tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'a'}},
+			keyMsg:       tea.KeyPressMsg{Code: 'a', Text: "a"},
 			setupMocks:   func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {},
 			expectedEdit: true,
 		},
 		{
 			name:         "tcp tunnel exits editing immediately",
 			tunnelType:   types.TunnelTypeTCP,
-			keyMsg:       tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'a'}},
+			keyMsg:       tea.KeyPressMsg{Code: 'a', Text: "a"},
 			setupMocks:   func(msr *MockSessionRegistry, ms *MockSlug, mr *MockRandom) {},
 			expectedEdit: false,
 		},
@@ -1006,19 +1006,19 @@ func TestModel_SlugView(t *testing.T) {
 func TestModel_ComingSoonUpdate(t *testing.T) {
 	tests := []struct {
 		name   string
-		keyMsg tea.KeyMsg
+		keyMsg tea.KeyPressMsg
 	}{
 		{
 			name:   "any key dismisses coming soon",
-			keyMsg: tea.KeyMsg{Type: tea.KeyEnter},
+			keyMsg: tea.KeyPressMsg{Code: tea.KeyEnter},
 		},
 		{
 			name:   "escape key dismisses",
-			keyMsg: tea.KeyMsg{Type: tea.KeyEsc},
+			keyMsg: tea.KeyPressMsg{Code: tea.KeyEsc},
 		},
 		{
 			name:   "space key dismisses",
-			keyMsg: tea.KeyMsg{Type: tea.KeySpace},
+			keyMsg: tea.KeyPressMsg{Code: ' ', Text: " "},
 		},
 	}
 
@@ -1097,7 +1097,7 @@ func TestModel_ComingSoonView(t *testing.T) {
 func TestModel_CommandsUpdate(t *testing.T) {
 	tests := []struct {
 		name             string
-		keyMsg           tea.KeyMsg
+		keyMsg           tea.KeyPressMsg
 		selectedItem     list.Item
 		expectCommands   bool
 		expectEditSlug   bool
@@ -1105,31 +1105,31 @@ func TestModel_CommandsUpdate(t *testing.T) {
 	}{
 		{
 			name:           "escape key closes commands",
-			keyMsg:         tea.KeyMsg{Type: tea.KeyEsc},
+			keyMsg:         tea.KeyPressMsg{Code: tea.KeyEsc},
 			expectCommands: false,
 		},
 		{
 			name:           "q key closes commands",
-			keyMsg:         tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'q'}},
+			keyMsg:         tea.KeyPressMsg{Code: 'q', Text: "q"},
 			expectCommands: false,
 		},
 		{
 			name:           "enter on slug command starts editing",
-			keyMsg:         tea.KeyMsg{Type: tea.KeyEnter},
+			keyMsg:         tea.KeyPressMsg{Code: tea.KeyEnter},
 			selectedItem:   commandItem{name: "slug", desc: "Set custom subdomain"},
 			expectCommands: false,
 			expectEditSlug: true,
 		},
 		{
 			name:             "enter on tunnel-type shows coming soon",
-			keyMsg:           tea.KeyMsg{Type: tea.KeyEnter},
+			keyMsg:           tea.KeyPressMsg{Code: tea.KeyEnter},
 			selectedItem:     commandItem{name: "tunnel-type", desc: "Change tunnel type"},
 			expectCommands:   false,
 			expectComingSoon: true,
 		},
 		{
 			name:           "arrow key navigates list",
-			keyMsg:         tea.KeyMsg{Type: tea.KeyDown},
+			keyMsg:         tea.KeyPressMsg{Code: tea.KeyDown},
 			expectCommands: true,
 		},
 	}
@@ -1261,28 +1261,28 @@ func TestModel_CommandsView(t *testing.T) {
 func TestModel_DashboardUpdate(t *testing.T) {
 	tests := []struct {
 		name           string
-		keyMsg         tea.KeyMsg
+		keyMsg         tea.KeyPressMsg
 		expectQuit     bool
 		expectCommands bool
 	}{
 		{
 			name:       "q key quits",
-			keyMsg:     tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'q'}},
+			keyMsg:     tea.KeyPressMsg{Code: 'q', Text: "q"},
 			expectQuit: true,
 		},
 		{
 			name:       "ctrl+c quits",
-			keyMsg:     tea.KeyMsg{Type: tea.KeyCtrlC},
+			keyMsg:     tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl},
 			expectQuit: true,
 		},
 		{
 			name:           "c key opens commands",
-			keyMsg:         tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'c'}},
+			keyMsg:         tea.KeyPressMsg{Code: 'c', Text: "c"},
 			expectCommands: true,
 		},
 		{
 			name:   "other keys do nothing",
-			keyMsg: tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'x'}},
+			keyMsg: tea.KeyPressMsg{Code: 'x', Text: "x"},
 		},
 	}
 
@@ -1922,6 +1922,10 @@ 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)

session/interaction/model.go

@@ -4,13 +4,13 @@ import (
 	"fmt"
 	"time"
 	"tunnel_pls/internal/random"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/help"
 	"github.com/charmbracelet/bubbles/key"
 	"github.com/charmbracelet/bubbles/list"
 	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 )
 
 type commandItem struct {
@@ -77,7 +77,7 @@ type keymap struct {
 type tickMsg time.Time
 
 func (m *model) Init() tea.Cmd {
-	return tea.Batch(textinput.Blink, tea.WindowSize())
+	return tea.RequestWindowSize
 }
 
 func getResponsiveWidth(screenWidth, padding, minWidth, maxWidth int) int {

session/interaction/slug.go

@@ -3,28 +3,26 @@ package interaction
 import (
 	"fmt"
 	"strings"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
+	tea "charm.land/bubbletea/v2"
 	"github.com/charmbracelet/bubbles/key"
-	"github.com/charmbracelet/bubbles/textinput"
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 )
 
-func (m *model) slugUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
+func (m *model) slugUpdate(msg tea.KeyPressMsg) (tea.Model, tea.Cmd) {
-	var cmd tea.Cmd
 
 	if m.tunnelType != types.TunnelTypeHTTP {
 		m.editingSlug = false
 		m.slugError = ""
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	}
 
 	switch msg.String() {
 	case "esc", "ctrl+c":
 		m.editingSlug = false
 		m.slugError = ""
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	case "enter":
 		inputValue := m.slugInput.Value()
 		if err := m.interaction.sessionRegistry.Update(m.interaction.user, types.SessionKey{
@@ -39,18 +37,18 @@ func (m *model) slugUpdate(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 		}
 		m.editingSlug = false
 		m.slugError = ""
-		return m, tea.Batch(tea.ClearScreen, textinput.Blink)
+		return m, nil
 	default:
 		if key.Matches(msg, m.keymap.random) {
 			newSubdomain, err := m.randomizer.String(20)
 			if err != nil {
-				return m, cmd
+				return m, nil
 			}
 			m.slugInput.SetValue(newSubdomain)
 		}
 		m.slugError = ""
-		m.slugInput, cmd = m.slugInput.Update(msg)
+		m.slugInput, _ = m.slugInput.Update(msg)
-		return m, cmd
+		return m, nil
 	}
 }
 

session/lifecycle/lifecycle.go

@@ -2,13 +2,14 @@ package lifecycle
 
 import (
 	"errors"
-	"fmt"
 	"io"
 	"net"
 	"sync"
 	"time"
-	"tunnel_pls/internal/session/slug"
+
-	"tunnel_pls/internal/types"
+	portUtil "tunnel_pls/internal/port"
+	"tunnel_pls/session/slug"
+	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -23,12 +24,6 @@ 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
@@ -39,18 +34,18 @@ type lifecycle struct {
 	slug            slug.Slug
 	startedAt       time.Time
 	sessionRegistry SessionRegistry
-	portRegistry    PortRegistry
+	portRegistry    portUtil.Port
 	user            string
 }
 
-func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port PortRegistry, sessionRegistry SessionRegistry, user string) Lifecycle {
+func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port portUtil.Port, sessionRegistry SessionRegistry, user string) Lifecycle {
 	return &lifecycle{
 		status:          types.SessionStatusINITIALIZING,
 		conn:            conn,
 		channel:         nil,
 		forwarder:       forwarder,
 		slug:            slugManager,
-		startedAt:       time.Time{},
+		startedAt:       time.Now(),
 		sessionRegistry: sessionRegistry,
 		portRegistry:    port,
 		user:            user,
@@ -60,16 +55,16 @@ func New(conn ssh.Conn, forwarder Forwarder, slugManager slug.Slug, port PortReg
 type Lifecycle interface {
 	Connection() ssh.Conn
 	Channel() ssh.Channel
-	PortRegistry() PortRegistry
+	PortRegistry() portUtil.Port
 	User() string
-	SetChannel(channel ssh.Channel) error
+	SetChannel(channel ssh.Channel)
 	SetStatus(status types.SessionStatus)
 	IsActive() bool
 	StartedAt() time.Time
 	Close() error
 }
 
-func (l *lifecycle) PortRegistry() PortRegistry {
+func (l *lifecycle) PortRegistry() portUtil.Port {
 	return l.portRegistry
 }
 
@@ -77,25 +72,11 @@ func (l *lifecycle) User() string {
 	return l.user
 }
 
-func (l *lifecycle) SetChannel(channel ssh.Channel) error {
+func (l *lifecycle) SetChannel(channel ssh.Channel) {
-	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
 }
 
@@ -106,13 +87,7 @@ 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 {
@@ -123,74 +98,50 @@ func (l *lifecycle) IsActive() bool {
 
 func (l *lifecycle) Close() error {
 	l.mu.Lock()
+	defer l.mu.Unlock()
 	if l.status == types.SessionStatusCLOSED {
-		closeErr := l.closeErr
+		return l.closeErr
-		l.mu.Unlock()
-		return closeErr
 	}
 	l.status = types.SessionStatusCLOSED
 
-	channel := l.channel
-	conn := l.conn
-	l.mu.Unlock()
-
 	var errs []error
-	if channel != nil {
+	tunnelType := l.forwarder.TunnelType()
-		if err := channel.Close(); err != nil && !isClosedError(err) {
+
-			errs = append(errs, err)
+	if l.channel != nil {
-		}
+		if err := l.channel.Close(); err != nil && !isClosedError(err) {
-	}
-	if conn != nil {
-		if err := conn.Close(); err != nil && !isClosedError(err) {
 			errs = append(errs, err)
 		}
 	}
 
-	l.cleanupRegistry()
+	if l.conn != nil {
-	if err := l.cleanupForwarder(); err != nil {
+		if err := l.conn.Close(); err != nil && !isClosedError(err) {
-		errs = append(errs, err)
+			errs = append(errs, err)
+		}
 	}
 
-	closeErr := errors.Join(errs...)
+	clientSlug := l.slug.String()
-
-	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:   slugStr,
+		Id:   clientSlug,
-		Type: l.forwarder.TunnelType(),
+		Type: tunnelType,
 	}
 	l.sessionRegistry.Remove(key)
-}
 
-func (l *lifecycle) cleanupForwarder() error {
+	if tunnelType == types.TunnelTypeTCP {
-	if l.forwarder.TunnelType() != types.TunnelTypeTCP {
+		errs = append(errs, l.PortRegistry().SetStatus(l.forwarder.ForwardedPort(), false))
-		return nil
+		errs = append(errs, l.forwarder.Close())
 	}
-	var errs []error
+
-	errs = append(errs, l.portRegistry.SetStatus(l.forwarder.ForwardedPort(), false))
+	l.closeErr = errors.Join(errs...)
-	errs = append(errs, l.forwarder.Close())
+	return l.closeErr
-	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)
+	return errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) || err.Error() == "EOF"
 }
 
 func (l *lifecycle) StartedAt() time.Time {
-	l.mu.Lock()
-	defer l.mu.Unlock()
 	return l.startedAt
 }

session/lifecycle/lifecycle_test.go

@@ -5,9 +5,8 @@ import (
 	"errors"
 	"io"
 	"net"
-	"sync"
 	"testing"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -178,14 +177,8 @@ func TestLifecycle_SetChannel(t *testing.T) {
 
 	mockSSHChannel := &MockSSHChannel{}
 
-	err := mockLifecycle.SetChannel(mockSSHChannel)
+	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())
 }
 
@@ -283,15 +276,14 @@ func TestLifecycle_Close(t *testing.T) {
 			mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
 
 			mockLifecycle.SetStatus(types.SessionStatusRUNNING)
-			err := mockLifecycle.SetChannel(mockSSHChannel)
+			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)
@@ -309,116 +301,3 @@ 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")
-}

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/internal/types"
+	"tunnel_pls/session/forwarder"
+	"tunnel_pls/session/interaction"
+	"tunnel_pls/session/lifecycle"
+	"tunnel_pls/session/slug"
+	"tunnel_pls/types"
 
 	"golang.org/x/crypto/ssh"
 )
@@ -158,15 +158,13 @@ func (s *session) setupInteractiveMode(channel ssh.NewChannel) error {
 	}
 
 	go func() {
-		err := s.HandleGlobalRequest(reqs)
+		err = s.HandleGlobalRequest(reqs)
 		if err != nil {
 			log.Printf("global request handler error: %v", err)
 		}
 	}()
 
-	if err = s.lifecycle.SetChannel(ch); err != nil {
+	s.lifecycle.SetChannel(ch)
-		return err
-	}
 	s.interaction.SetChannel(ch)
 	s.interaction.SetMode(types.InteractiveModeINTERACTIVE)
 
@@ -200,22 +198,23 @@ func (s *session) waitForSessionEnd() error {
 }
 
 func (s *session) waitForTCPIPForward() *ssh.Request {
-	for {
+	select {
-		select {
+	case req, ok := <-s.initialReq:
-		case req, ok := <-s.initialReq:
+		if !ok {
-			if !ok {
+			log.Println("Forwarding request channel closed")
-				log.Println("Forwarding request channel closed")
-				return nil
-			}
-			if req.Type == "tcpip-forward" {
-				return req
-			}
-			log.Printf("Ignoring unexpected global request: %s", req.Type)
-			_ = req.Reply(false, nil)
-		case <-time.After(500 * time.Millisecond):
-			log.Println("No tcpip-forward request received within timeout")
 			return nil
 		}
+		if req.Type == "tcpip-forward" {
+			return req
+		}
+		if err := req.Reply(false, nil); err != nil {
+			log.Printf("Failed to reply to request: %v", err)
+		}
+		log.Printf("Expected tcpip-forward request, got: %s", req.Type)
+		return nil
+	case <-time.After(500 * time.Millisecond):
+		log.Println("No forwarding request received")
+		return nil
 	}
 }
 
@@ -378,7 +377,8 @@ func (s *session) HandleTCPForward(req *ssh.Request, addr string, portToBind uin
 	}
 
 	go func() {
-		if err := tcpServer.Serve(listener); err != nil {
+		err = tcpServer.Serve(listener)
+		if err != nil {
 			log.Printf("Failed serving tcp server: %s\n", err)
 		}
 	}()

session/session_test.go

@@ -15,8 +15,8 @@ import (
 	"time"
 	"tunnel_pls/internal/config"
 	"tunnel_pls/internal/registry"
-	"tunnel_pls/internal/session/lifecycle"
+	"tunnel_pls/session/lifecycle"
-	"tunnel_pls/internal/types"
+	"tunnel_pls/types"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -807,7 +807,7 @@ func (m *mockNewChanFail) Accept() (ssh.Channel, <-chan *ssh.Request, error) {
 }
 
 func TestWaitForTCPIPForward_EdgeCases(t *testing.T) {
-	t.Run("Wrong Request Type Then Timeout", func(t *testing.T) {
+	t.Run("Wrong Request Type", func(t *testing.T) {
 		_, sReqs, _, cConn, cleanup := setupSSH(t)
 		defer cleanup()
 
@@ -817,65 +817,10 @@ func TestWaitForTCPIPForward_EdgeCases(t *testing.T) {
 			_, _, _ = cConn.SendRequest("not-tcpip-forward", true, nil)
 		}()
 
-		start := time.Now()
 		req := s.waitForTCPIPForward()
-		elapsed := time.Since(start)
-
 		if req != nil {
 			t.Error("expected nil request")
 		}
-		if elapsed < 400*time.Millisecond {
-			t.Errorf("expected timeout ~500ms, got %v", elapsed)
-		}
-	})
-
-	t.Run("Multiple Non-Forward Requests Then Success", func(t *testing.T) {
-		_, sReqs, _, cConn, cleanup := setupSSH(t)
-		defer cleanup()
-
-		s := &session{initialReq: sReqs}
-
-		go func() {
-			time.Sleep(100 * time.Millisecond)
-			_, _, _ = cConn.SendRequest("keepalive@openssh.com", false, nil)
-			time.Sleep(100 * time.Millisecond)
-			_, _, _ = cConn.SendRequest("hostkeys-00@openssh.com", false, nil)
-			time.Sleep(100 * time.Millisecond)
-			_, _, _ = cConn.SendRequest("tcpip-forward", true, nil)
-		}()
-
-		req := s.waitForTCPIPForward()
-		if req == nil {
-			t.Error("expected tcpip-forward request, got nil")
-		}
-		if req != nil && req.Type != "tcpip-forward" {
-			t.Errorf("expected tcpip-forward, got %s", req.Type)
-		}
-	})
-
-	t.Run("Timeout After Non-Forward Requests", func(t *testing.T) {
-		_, sReqs, _, cConn, cleanup := setupSSH(t)
-		defer cleanup()
-
-		s := &session{initialReq: sReqs}
-
-		go func() {
-			time.Sleep(100 * time.Millisecond)
-			_, _, _ = cConn.SendRequest("keepalive@openssh.com", false, nil)
-			time.Sleep(100 * time.Millisecond)
-			_, _, _ = cConn.SendRequest("hostkeys-00@openssh.com", false, nil)
-		}()
-
-		start := time.Now()
-		req := s.waitForTCPIPForward()
-		elapsed := time.Since(start)
-
-		if req != nil {
-			t.Error("expected nil request after timeout")
-		}
-		if elapsed < 400*time.Millisecond {
-			t.Errorf("expected timeout ~500ms after last request, got %v", elapsed)
-		}
 	})
 
 	t.Run("Channel Closed", func(t *testing.T) {

session/slug/slug.go

@@ -1,14 +1,11 @@
 package slug
 
-import "sync"
-
 type Slug interface {
 	String() string
 	Set(slug string)
 }
 
 type slug struct {
-	mu   sync.RWMutex
 	slug string
 }
 
@@ -19,13 +16,9 @@ func New() Slug {
 }
 
 func (s *slug) String() string {
-	s.mu.RLock()
-	defer s.mu.RUnlock()
 	return s.slug
 }
 
 func (s *slug) Set(slug string) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
 	s.slug = slug
 }