refactor(session): add registry to manage SSH sessions

- Implement thread-safe session registry with sync.RWMutex
- Add Registry interface for session management operations
- Support Get, Register, Update, and Remove session operations
- Enable dynamic slug updates for existing sessions

server/handler.go

@@ -1,28 +0,0 @@
-package server
-
-import (
-	"log"
-	"net"
-	"tunnel_pls/session"
-
-	"golang.org/x/crypto/ssh"
-)
-
-func (s *Server) handleConnection(conn net.Conn) {
-	sshConn, chans, forwardingReqs, err := ssh.NewServerConn(conn, s.config)
-	if err != nil {
-		log.Printf("failed to establish SSH connection: %v", err)
-		err := conn.Close()
-		if err != nil {
-			log.Printf("failed to close SSH connection: %v", err)
-			return
-		}
-		return
-	}
-
-	log.Println("SSH connection established:", sshConn.User())
-
-	session.New(sshConn, forwardingReqs, chans)
-
-	return
-}

server/http.go

@@ -17,15 +17,9 @@ import (
 	"golang.org/x/crypto/ssh"
 )
 
-type Interaction interface {
-	SendMessage(message string)
-}
-
 type HTTPWriter interface {
 	io.Reader
 	io.Writer
-	SetInteraction(interaction Interaction)
-	AddInteraction(interaction Interaction)
 	GetRemoteAddr() net.Addr
 	GetWriter() io.Writer
 	AddResponseMiddleware(mw ResponseMiddleware)
@@ -35,21 +29,16 @@ type HTTPWriter interface {
 }
 
 type customWriter struct {
-	remoteAddr  net.Addr
-	writer      io.Writer
-	reader      io.Reader
-	headerBuf   []byte
-	buf         []byte
-	respHeader  ResponseHeaderManager
-	reqHeader   RequestHeaderManager
-	interaction Interaction
-	respMW      []ResponseMiddleware
-	reqStartMW  []RequestMiddleware
-	reqEndMW    []RequestMiddleware
-}
-
-func (cw *customWriter) SetInteraction(interaction Interaction) {
-	cw.interaction = interaction
+	remoteAddr net.Addr
+	writer     io.Writer
+	reader     io.Reader
+	headerBuf  []byte
+	buf        []byte
+	respHeader ResponseHeaderManager
+	reqHeader  RequestHeaderManager
+	respMW     []ResponseMiddleware
+	reqStartMW []RequestMiddleware
+	reqEndMW   []RequestMiddleware
 }
 
 func (cw *customWriter) GetRemoteAddr() net.Addr {
@@ -135,11 +124,10 @@ func (cw *customWriter) Read(p []byte) (int, error) {
 
 func NewCustomWriter(writer io.Writer, reader io.Reader, remoteAddr net.Addr) HTTPWriter {
 	return &customWriter{
-		remoteAddr:  remoteAddr,
-		writer:      writer,
-		reader:      reader,
-		buf:         make([]byte, 0, 4096),
-		interaction: nil,
+		remoteAddr: remoteAddr,
+		writer:     writer,
+		reader:     reader,
+		buf:        make([]byte, 0, 4096),
 	}
 }
 
@@ -224,13 +212,23 @@ func (cw *customWriter) Write(p []byte) (int, error) {
 	return len(p), nil
 }
 
-func (cw *customWriter) AddInteraction(interaction Interaction) {
-	cw.interaction = interaction
-}
-
 var redirectTLS = false
 
-func NewHTTPServer() error {
+type HTTPServer interface {
+	ListenAndServe() error
+	ListenAndServeTLS() error
+	handler(conn net.Conn)
+	handlerTLS(conn net.Conn)
+}
+type httpServer struct {
+	sessionRegistry session.Registry
+}
+
+func NewHTTPServer(sessionRegistry session.Registry) HTTPServer {
+	return &httpServer{sessionRegistry: sessionRegistry}
+}
+
+func (hs *httpServer) ListenAndServe() error {
 	httpPort := config.Getenv("HTTP_PORT", "8080")
 	listener, err := net.Listen("tcp", ":"+httpPort)
 	if err != nil {
@@ -251,13 +249,13 @@ func NewHTTPServer() error {
 				continue
 			}
 
-			go Handler(conn)
+			go hs.handler(conn)
 		}
 	}()
 	return nil
 }
 
-func Handler(conn net.Conn) {
+func (hs *httpServer) handler(conn net.Conn) {
 	defer func() {
 		err := conn.Close()
 		if err != nil && !errors.Is(err, net.ErrClosed) {
@@ -316,8 +314,8 @@ func Handler(conn net.Conn) {
 		return
 	}
 
-	sshSession, ok := session.Clients[slug]
-	if !ok {
+	sshSession, exist := hs.sessionRegistry.Get(slug)
+	if !exist {
 		_, err = conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" +
 			fmt.Sprintf("Location: https://tunnl.live/tunnel-not-found?slug=%s\r\n", slug) +
 			"Content-Length: 0\r\n" +

server/https.go

@@ -9,10 +9,9 @@ import (
 	"net"
 	"strings"
 	"tunnel_pls/internal/config"
-	"tunnel_pls/session"
 )
 
-func NewHTTPSServer() error {
+func (hs *httpServer) ListenAndServeTLS() error {
 	domain := config.Getenv("DOMAIN", "localhost")
 	httpsPort := config.Getenv("HTTPS_PORT", "8443")
 
@@ -38,13 +37,13 @@ func NewHTTPSServer() error {
 				continue
 			}
 
-			go HandlerTLS(conn)
+			go hs.handlerTLS(conn)
 		}
 	}()
 	return nil
 }
 
-func HandlerTLS(conn net.Conn) {
+func (hs *httpServer) handlerTLS(conn net.Conn) {
 	defer func() {
 		err := conn.Close()
 		if err != nil {
@@ -90,8 +89,8 @@ func HandlerTLS(conn net.Conn) {
 		return
 	}
 
-	sshSession, ok := session.Clients[slug]
-	if !ok {
+	sshSession, exist := hs.sessionRegistry.Get(slug)
+	if !exist {
 		_, err = conn.Write([]byte("HTTP/1.1 301 Moved Permanently\r\n" +
 			fmt.Sprintf("Location: https://tunnl.live/tunnel-not-found?slug=%s\r\n", slug) +
 			"Content-Length: 0\r\n" +

server/server.go

@@ -4,50 +4,45 @@ import (
 	"fmt"
 	"log"
 	"net"
-	"net/http"
 	"tunnel_pls/internal/config"
+	"tunnel_pls/session"
 
 	"golang.org/x/crypto/ssh"
 )
 
 type Server struct {
-	conn       *net.Listener
-	config     *ssh.ServerConfig
-	httpServer *http.Server
+	conn            *net.Listener
+	config          *ssh.ServerConfig
+	sessionRegistry session.Registry
 }
 
-func (s *Server) GetConn() *net.Listener {
-	return s.conn
-}
-
-func (s *Server) GetConfig() *ssh.ServerConfig {
-	return s.config
-}
-
-func (s *Server) GetHttpServer() *http.Server {
-	return s.httpServer
-}
-
-func NewServer(sshConfig *ssh.ServerConfig) *Server {
+func NewServer(sshConfig *ssh.ServerConfig, sessionRegistry session.Registry) (*Server, error) {
 	listener, err := net.Listen("tcp", fmt.Sprintf(":%s", config.Getenv("PORT", "2200")))
 	if err != nil {
 		log.Fatalf("failed to listen on port 2200: %v", err)
-		return nil
+		return nil, err
 	}
-	if config.Getenv("TLS_ENABLED", "false") == "true" {
-		err = NewHTTPSServer()
-		if err != nil {
-			log.Fatalf("failed to start https server: %v", err)
-		}
-	}
-	err = NewHTTPServer()
+
+	HttpServer := NewHTTPServer(sessionRegistry)
+	err = HttpServer.ListenAndServe()
 	if err != nil {
 		log.Fatalf("failed to start http server: %v", err)
+		return nil, err
 	}
+
+	if config.Getenv("TLS_ENABLED", "false") == "true" {
+		err = HttpServer.ListenAndServeTLS()
+		if err != nil {
+			log.Fatalf("failed to start https server: %v", err)
+			return nil, err
+		}
+	}
+
 	return &Server{
-		conn:   &listener,
-		config: sshConfig,
-	}
+		conn:            &listener,
+		config:          sshConfig,
+		sessionRegistry: sessionRegistry,
+	}, nil
 }
 
 func (s *Server) Start() {
@@ -62,3 +57,26 @@ func (s *Server) Start() {
 		go s.handleConnection(conn)
 	}
 }
+
+func (s *Server) handleConnection(conn net.Conn) {
+	sshConn, chans, forwardingReqs, err := ssh.NewServerConn(conn, s.config)
+	if err != nil {
+		log.Printf("failed to establish SSH connection: %v", err)
+		err := conn.Close()
+		if err != nil {
+			log.Printf("failed to close SSH connection: %v", err)
+			return
+		}
+		return
+	}
+
+	log.Println("SSH connection established:", sshConn.User())
+
+	sshSession := session.New(sshConn, forwardingReqs, chans, s.sessionRegistry)
+	err = sshSession.Start()
+	if err != nil {
+		log.Printf("SSH session ended with error: %v", err)
+		return
+	}
+	return
+}